Update models/chat.py

models/chat.py

@@ -0,0 +1,34 @@
+from sqlmodel import SQLModel, Field, Relationship
+from typing import Optional, List
+from datetime import datetime
+
+class ChatSessionBase(SQLModel):
+    user_id: int = Field(foreign_key="user.id")
+    start_time: datetime = Field(default_factory=datetime.utcnow)
+    title: Optional[str] = None # e.g., "Consultation on 2023-10-26"
+
+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
+    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):
+    session_id: int = Field(foreign_key="chatsession.id")
+    role: str  # "user", "assistant", "system", "tool"
+    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
+
+
+class ChatMessage(ChatMessageBase, table=True):
+    id: Optional[int] = Field(default=None, primary_key=True)
+    session: Optional[ChatSession] = Relationship(back_populates="messages")
+
+class ChatMessageCreate(ChatMessageBase):
+    pass