src/generation.py

"""
Geração de respostas usando LLMs
"""
from typing import Optional, List, Dict, Any, Iterator
from .config import LLM_PROVIDER, DEFAULT_TEMPERATURE, DEFAULT_MAX_TOKENS
from .llms.factory import create_llm
from .llms.base import BaseLLM


class GenerationManager:
    """Gerenciador de geração de texto com suporte a múltiplos providers"""

    def __init__(self, provider: Optional[str] = None, model_id: Optional[str] = None):
        """
        Inicializa gerenciador de geração

        Args:
            provider: Nome do provider (huggingface, openai, anthropic, ollama)
                     Se None, usa LLM_PROVIDER do .env
            model_id: ID do modelo. Se None, usa default do provider
        """
        self.provider_name = provider or LLM_PROVIDER
        self.model_id = model_id
        self.llm: Optional[BaseLLM] = None

    def get_client(self) -> Optional[BaseLLM]:
        """Obtém cliente LLM (lazy loading com fallback)"""
        if self.llm is None:
            self.llm = create_llm(
                provider=self.provider_name,
                model_id=self.model_id,
                fallback=True
            )
        return self.llm

    def build_rag_prompt(
        self,
        question: str,
        contexts: List[Dict[str, Any]],
        system_prompt: Optional[str] = None
    ) -> str:
        """
        Constrói prompt para RAG

        Args:
            question: Pergunta do usuário
            contexts: Lista de contextos recuperados
            system_prompt: Prompt de sistema customizado

        Returns:
            Prompt formatado
        """
        if system_prompt is None:
            system_prompt = "Use os trechos fornecidos para responder à pergunta de forma precisa e concisa."

        context_text = "\n\n".join([
            f"Trecho {i+1} (fonte: {ctx['title']}):\n{ctx['content']}"
            for i, ctx in enumerate(contexts)
        ])

        prompt = f"""{system_prompt}

Contexto:
{context_text}

Pergunta: {question}

Resposta:"""

        return prompt

    def generate(
        self,
        prompt: str,
        temperature: float = DEFAULT_TEMPERATURE,
        max_tokens: int = DEFAULT_MAX_TOKENS
    ) -> str:
        """
        Gera resposta usando LLM

        Args:
            prompt: Prompt para o modelo
            temperature: Temperatura de geração
            max_tokens: Máximo de tokens a gerar

        Returns:
            Texto gerado
        """
        client = self.get_client()

        if client is None:
            return "Erro: Nenhum provider LLM disponível. Verifique as configurações no .env"

        if not client.is_available():
            error_info = client.get_model_info()
            return f"Erro: Provider {error_info.get('provider')} indisponível. {client.last_error}"

        try:
            response = client.generate(
                prompt=prompt,
                temperature=temperature,
                max_tokens=max_tokens
            )
            return response
        except Exception as e:
            return f"Erro na geração: {str(e)}"

    def generate_stream(
        self,
        prompt: str,
        temperature: float = DEFAULT_TEMPERATURE,
        max_tokens: int = DEFAULT_MAX_TOKENS
    ) -> Iterator[str]:
        """
        Gera resposta em streaming (se suportado pelo provider)

        Args:
            prompt: Prompt para o modelo
            temperature: Temperatura de geração
            max_tokens: Máximo de tokens a gerar

        Yields:
            Tokens gerados progressivamente
        """
        # Nota: Streaming ainda não implementado para todos os providers
        # Por enquanto, retorna resposta completa
        response = self.generate(prompt, temperature, max_tokens)
        yield response

    def format_sources(self, contexts: List[Dict[str, Any]]) -> str:
        """
        Formata fontes para exibição

        Args:
            contexts: Lista de contextos

        Returns:
            String formatada com fontes
        """
        if not contexts:
            return "\n\n**Fontes:** Nenhuma fonte encontrada"

        sources = []
        for i, ctx in enumerate(contexts, 1):
            preview = ctx['content'][:150] + "..." if len(ctx['content']) > 150 else ctx['content']
            score = ctx.get('score', 0)
            sources.append(
                f"{i}. **{ctx['title']}** (relevância: {score:.2%})\n   _{preview}_"
            )

        return "\n\n**Fontes:**\n" + "\n\n".join(sources)