models/chat.py

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