models/chat.py

# /home/user/app/models/chat.py
from sqlmodel import SQLModel, Field, Relationship
from typing import Optional, List, TYPE_CHECKING, Dict, Any # Ensure TYPE_CHECKING is imported
from datetime import datetime

if TYPE_CHECKING:
    from .user import User

class ChatSessionBase(SQLModel):
    user_id: int = Field(foreign_key="user.id")
    start_time: datetime = Field(default_factory=datetime.utcnow)
    title: Optional[str] = None
    # NEW: For storing initial context provided by the clinician
    patient_context_summary: Optional[str] = Field(default=None, description="Brief summary of patient context for this session")
    # NEW: Could store structured context as JSON if more detail is needed
    # patient_context_structured: Optional[Dict[str, Any]] = Field(default=None, sa_column=Column(JSON)) # If using SQLAlchemy Column for JSON

class ChatSession(ChatSessionBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    user: Optional["User"] = Relationship(back_populates="chat_sessions")
    messages: List["ChatMessage"] = Relationship(back_populates="session")

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
    tool_name: Optional[str] = None
    # NEW: For AI to optionally provide source/confidence
    # source_references: Optional[List[str]] = Field(default=None, sa_column=Column(JSON)) # If using SQLAlchemy Column for JSON
    # confidence_score: Optional[float] = None

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