Spaces:
Sleeping
Sleeping
File size: 4,130 Bytes
58de15f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
# app/db/schemas.py
import uuid
from datetime import datetime
from pydantic import BaseModel, Field
import enum as py_enum # Use alias
from typing import List, Dict, Any # Added List, Dict, Any
# --- Define Enums Here ---
class MessageRole(str, py_enum.Enum):
"""Enum for message roles."""
USER = "user"
BOT = "bot"
class FeedbackTypeEnum(str, py_enum.Enum):
"""Enum for feedback types."""
ACCEPT = "accept"
REJECT = "reject"
# --- End Enums ---
# --- Base Models (for shared fields) ---
class MessageBase(BaseModel):
"""Base schema for message properties."""
role: MessageRole
content: str
class SessionBase(BaseModel):
"""Base schema for session properties."""
name: str | None = None
# --- Schemas for specific data structures ---
class TopResultPreview(BaseModel):
"""Schema for the preview of the top search result."""
id: str
score: float
content_preview: str
original_file: str | None = None
class QueryResultDetail(BaseModel):
"""Schema for detailed information about a query result."""
session_id: uuid.UUID
message_id: uuid.UUID | None = None # Bot message ID
attempt_number: int = 1
retrieved_ids: list[str] = []
search_scores: list[float] = []
original_file: str | None = None # Top result's original file
# --- Schemas for Creation (request bodies) ---
class MessageCreate(MessageBase):
"""Schema for creating a new message (used internally)."""
pass
class SessionCreate(SessionBase):
"""Schema for creating a new session."""
name: str = Field(..., min_length=1, example="My New Chat Session")
class FeedbackCreate(BaseModel):
"""Schema for submitting feedback."""
message_id: uuid.UUID
feedback_type: FeedbackTypeEnum
feedback_comment: str | None = None
class QueryRequest(BaseModel):
"""Schema for the user query request."""
query: str
session_id: uuid.UUID | None = None # Session ID is optional
# --- Schemas for Reading (responses, includes fields from DB) ---
class ChatMessage(MessageBase):
"""Schema for reading chat messages from the database."""
id: uuid.UUID
session_id: uuid.UUID
created_at: datetime
original_query: str | None = None
retrieved_context_ids: list[str] | None = None
used_context_ids: list[str] | None = None # IDs used for this specific message
attempt_number: int = 1
cumulative_used_context_ids: list[str] | None = None # All IDs used up to this attempt
class Config:
from_attributes = True # Use from_attributes for Pydantic v2
# orm_mode = True # Use this for Pydantic v1
class FeedbackLog(FeedbackCreate):
"""Schema for reading feedback logs from the database."""
id: uuid.UUID
created_at: datetime
class Config:
from_attributes = True # Use from_attributes for Pydantic v2
# orm_mode = True # Use this for Pydantic v1
class FeedbackLogResponse(FeedbackLog):
"""Explicit response schema for feedback log entries."""
pass
class ChatSession(SessionBase):
"""Schema for reading chat sessions from the database."""
id: uuid.UUID
created_at: datetime
# Optionally include messages: list[ChatMessage] = []
class Config:
from_attributes = True # Use from_attributes for Pydantic v2
# orm_mode = True # Use this for Pydantic v1
# --- Schemas for Endpoint Responses ---
class QueryResponse(BaseModel):
"""Schema for the standard query response."""
llm_answer: str
context_used_preview: str | None = None
top_result_preview: TopResultPreview | None = None
details: QueryResultDetail
class RegeneratedResponse(BaseModel):
"""Schema for the response when regeneration occurs."""
new_llm_answer: str
new_context_used_preview: str | None = None
new_top_result_preview: TopResultPreview | None = None
new_details: QueryResultDetail # Uses the same detail structure
# --- Schemas for Updating ---
class ChatSessionUpdate(BaseModel):
"""Schema for updating a chat session (e.g., renaming)."""
name: str = Field(..., min_length=1, example="Renamed Session")
|