Update models/chat.py

models/chat.py

@@ -1,34 +1,43 @@
 from sqlmodel import SQLModel, Field, Relationship
-from typing import Optional, List
+from typing import Optional, List, TYPE_CHECKING # Ensure TYPE_CHECKING is imported
 from datetime import datetime
 
+# Forward reference for type hinting to avoid circular imports
+if TYPE_CHECKING:
+    from .user import User # For the 'user' attribute in ChatSession
+
 class ChatSessionBase(SQLModel):
-    user_id: int = Field(foreign_key="user.id")
+    user_id: int = Field(foreign_key="user.id") # Foreign key to the 'user' table's 'id' column
     start_time: datetime = Field(default_factory=datetime.utcnow)
-    title: Optional[str] = None # e.g., "Consultation on 2023-10-26"
+    title: Optional[str] = None
 
 class ChatSession(ChatSessionBase, table=True):
     id: Optional[int] = Field(default=None, primary_key=True)
     
-    user: Optional["User"] = Relationship(back_populates="chat_sessions") # Define in User model
+    # ------------ RELATIONSHIP DEFINITION ------------
+    # A ChatSession belongs to one User.
+    # `back_populates="chat_sessions"` links this to the `chat_sessions` attribute in the User model.
+    user: Optional["User"] = Relationship(back_populates="chat_sessions")
+    
+    # A ChatSession can have many ChatMessage objects.
     messages: List["ChatMessage"] = Relationship(back_populates="session")
-
-# Add to User model in user.py:
-# chat_sessions: List["ChatSession"] = Relationship(back_populates="user")
-
+    # -------------------------------------------------
 
 class ChatMessageBase(SQLModel):
+    # Foreign key to the 'chatsession' table's 'id' column.
+    # SQLModel infers table names from class names (lowercase).
     session_id: int = Field(foreign_key="chatsession.id")
-    role: str  # "user", "assistant", "system", "tool"
+    role: str
     content: str
     timestamp: datetime = Field(default_factory=datetime.utcnow)
-    tool_call_id: Optional[str] = None # For tool messages
-    tool_name: Optional[str] = None # For tool messages
-
+    tool_call_id: Optional[str] = None
+    tool_name: Optional[str] = None
 
 class ChatMessage(ChatMessageBase, table=True):
     id: Optional[int] = Field(default=None, primary_key=True)
-    session: Optional[ChatSession] = Relationship(back_populates="messages")
+    
+    # A ChatMessage belongs to one ChatSession.
+    session: Optional["ChatSession"] = Relationship(back_populates="messages")
 
 class ChatMessageCreate(ChatMessageBase):
     pass