src/backend/chat.py

"""Chat utilities for managing conversation history with chat templates."""

from typing import List, Optional, Literal
from pydantic import BaseModel, Field, field_validator
from transformers import PreTrainedTokenizer


class Message(BaseModel):
    """
    Mensagem de chat no formato compatível OpenAI.
    
    Exemplo:
        msg = Message(role="user", content="Olá!")
        msg_dict = msg.model_dump()  # {"role": "user", "content": "Olá!"}
    """
    
    role: Literal["user", "assistant", "system"] = Field(
        ...,
        description="Role da mensagem: user, assistant ou system"
    )
    content: str = Field(
        ...,
        min_length=1,
        description="Conteúdo da mensagem"
    )
    
    @field_validator("content")
    @classmethod
    def validate_content(cls, v: str) -> str:
        """Valida que o conteúdo não está vazio."""
        if not v.strip():
            raise ValueError("Content não pode estar vazio")
        return v
    
    def model_dump_dict(self) -> dict:
        """Retorna como dicionário (compatível com transformers)."""
        return {"role": self.role, "content": self.content}
    
    class Config:
        """Configuração do Pydantic."""
        json_schema_extra = {
            "example": {
                "role": "user",
                "content": "Olá! Como você está?"
            }
        }


def _format_chat_prompt(
    tokenizer: PreTrainedTokenizer,
    messages: List[Message],
    add_generation_prompt: bool = True,
) -> str:
    """
    Formata histórico de chat usando o template do modelo (função auxiliar interna).
    
    Args:
        tokenizer: Tokenizer do modelo (deve ter chat_template configurado)
        messages: Lista de mensagens (Message ou dict)
        add_generation_prompt: Se True, adiciona prompt de geração ao final
        
    Returns:
        String formatada pronta para ser enviada ao modelo
    """
    # Converte Message para dict se necessário
    messages_dict = [
        msg.model_dump_dict() if isinstance(msg, Message) else msg
        for msg in messages
    ]
    
    if not hasattr(tokenizer, "apply_chat_template") or tokenizer.chat_template is None:
        # Fallback: concatena mensagens simplesmente
        formatted = ""
        for msg in messages_dict:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            formatted += f"{role}: {content}\n"
        return formatted.strip()
    
    return tokenizer.apply_chat_template(
        messages_dict,
        tokenize=False,
        add_generation_prompt=add_generation_prompt,
    )


def _get_conversation_summary(messages: List[Message], max_length: int = 100) -> str:
    """
    Retorna resumo da conversa (função auxiliar interna).
    
    Args:
        messages: Lista de mensagens
        max_length: Comprimento máximo do resumo
        
    Returns:
        String resumida da conversa
    """
    summary_parts = []
    for msg in messages[-5:]:  # Últimas 5 mensagens
        if isinstance(msg, Message):
            role = msg.role
            content = msg.content[:50]
        else:
            role = msg.get("role", "unknown")
            content = msg.get("content", "")[:50]
        summary_parts.append(f"{role}: {content}...")
    
    summary = " | ".join(summary_parts)
    if len(summary) > max_length:
        return summary[:max_length] + "..."
    return summary


class Conversation(BaseModel):
    """
    Gerencia histórico de conversa de forma orientada a objetos com Pydantic.
    
    Exemplo:
        conv = Conversation()
        conv.add_user_message("Olá")
        conv.add_assistant_message("Oi! Como posso ajudar?")
        messages = conv.messages
    """
    
    messages: List[Message] = Field(default_factory=list)
    system_prompt: Optional[str] = Field(default=None)
    
    def __init__(self, system_prompt: Optional[str] = None, **data):
        """
        Inicializa uma nova conversa.
        
        Args:
            system_prompt: Prompt do sistema (opcional)
        """
        super().__init__(**data)
        if system_prompt and not self.messages:
            self.set_system_prompt(system_prompt)
    
    def add_message(self, role: Literal["user", "assistant", "system"], content: str) -> None:
        """
        Adiciona uma mensagem ao histórico.
        
        Args:
            role: Role da mensagem ("user", "assistant", "system")
            content: Conteúdo da mensagem
        """
        message = Message(role=role, content=content)
        if role == "system":
            # Mensagens do sistema sempre vão no início
            self.messages.insert(0, message)
        else:
            self.messages.append(message)
    
    def add_user_message(self, content: str) -> None:
        """Adiciona mensagem do usuário."""
        self.add_message("user", content)
    
    def add_assistant_message(self, content: str) -> None:
        """Adiciona mensagem do assistente."""
        self.add_message("assistant", content)
    
    def set_system_prompt(self, content: str) -> None:
        """
        Define ou atualiza o prompt do sistema.
        
        Args:
            content: Conteúdo do prompt do sistema
        """
        # Remove mensagens do sistema existentes
        self.messages = [msg for msg in self.messages if msg.role != "system"]
        # Adiciona nova mensagem do sistema no início
        self.messages.insert(0, Message(role="system", content=content))
    
    def clear(self, keep_system: bool = True) -> None:
        """
        Limpa o histórico de conversa.
        
        Args:
            keep_system: Se True, mantém mensagens do sistema
        """
        if keep_system:
            self.messages = [msg for msg in self.messages if msg.role == "system"]
        else:
            self.messages = []
    
    def get_summary(self, max_length: int = 100) -> str:
        """
        Retorna resumo da conversa.
        
        Args:
            max_length: Comprimento máximo do resumo
            
        Returns:
            String resumida da conversa
        """
        return _get_conversation_summary(self.messages, max_length)
    
    def model_dump_messages(self) -> List[dict]:
        """Retorna mensagens como lista de dicionários (compatível com transformers)."""
        return [msg.model_dump_dict() for msg in self.messages]
    
    def __len__(self) -> int:
        """Retorna número de mensagens."""
        return len(self.messages)
    
    def __repr__(self) -> str:
        """Representação string da conversa."""
        return f"Conversation({len(self.messages)} messages)"