from pydantic import BaseModel
from typing import Optional
from .common import UserContext, SourceInfo


class ChatRequest(BaseModel):
    query: str
    user_id: int
    conversation_id: str = ""
    context: Optional[UserContext] = None


class ChatResponse(BaseModel):
    answer: str
    conversation_id: str
    sources: list[SourceInfo] = []
    tasks: list[dict] = []
    trace_id: str = ""