schemas/chat.py

from typing import List, Optional
from pydantic import BaseModel

from components.llm.common import ChatRequest, Message


class MessageSchema(BaseModel):
    role: str
    content: str
    searchResults: Optional[str] = ''
    searchEntities: Optional[List[str]] = []
    reasoning: Optional[str] = ''
    
    def to_bl(self) -> Message:
        return ChatRequest.model_validate(self.model_dump())
    
    
class ChatRequestSchema(BaseModel):
    history: List[MessageSchema]
    chat_id: Optional[str]
    
    def to_bl(self) -> ChatRequest:
        return ChatRequest.model_validate(self.model_dump(exclude={"chat_id"}))
    
    @classmethod
    def from_bl(cls, bl: ChatRequest, chat_id: Optional[str] = None) -> "ChatRequestSchema":
        return cls.model_validate({
            "history": [msg.model_dump() for msg in bl.history],
            "chat_id": chat_id
        })