nodes/llm_agent.py

"""
Nó do agente LLM para o AgentPDF.

Este nó é responsável por gerar respostas inteligentes usando GPT-4o-mini
baseadas no contexto recuperado do PDF e na pergunta do usuário.
"""

from typing import Dict, Any
from langchain_openai import ChatOpenAI
from langchain_core.messages import AIMessage, SystemMessage, HumanMessage
from langchain_core.runnables import RunnableConfig
from langchain_core.prompts import ChatPromptTemplate

from agents.state import PDFState, ProcessingStatus
from utils.config import Config, get_openai_api_key
from utils.logger import log_node_execution, main_logger


def llm_agent_node(state: PDFState, config: RunnableConfig) -> Dict[str, Any]:
    """
    Nó responsável por gerar respostas usando o LLM.
    
    Este nó:
    1. Recebe a pergunta e o contexto recuperado
    2. Constrói um prompt otimizado
    3. Chama o GPT-4o-mini para gerar a resposta
    4. Processa e valida a resposta
    5. Atualiza o estado com a resposta final
    
    Args:
        state: Estado atual do grafo
        config: Configuração do LangGraph
    
    Returns:
        Dict[str, Any]: Atualizações para o estado
    """
    log_node_execution("LLM_AGENT", "START", "Iniciando geração de resposta")
    
    try:
        # Verifica se há pergunta e contexto
        user_question = state.get("user_question")
        retrieved_context = state.get("retrieved_context", [])
        
        if not user_question:
            error_msg = "Pergunta do usuário não encontrada"
            log_node_execution("LLM_AGENT", "ERROR", error_msg)
            return {
                "processing_status": ProcessingStatus.ERROR,
                "error_message": error_msg
            }
        
        # Verifica API key
        api_key = get_openai_api_key()
        if not api_key:
            error_msg = "Chave da API OpenAI não configurada"
            log_node_execution("LLM_AGENT", "ERROR", error_msg)
            return {
                "processing_status": ProcessingStatus.ERROR,
                "error_message": error_msg
            }
        
        log_node_execution(
            "LLM_AGENT", 
            "PROCESSING", 
            f"Gerando resposta para: '{user_question[:100]}...'"
        )
        
        # Cria o modelo LLM
        llm = create_llm_model()
        
        # Constrói o prompt
        prompt = build_prompt(user_question, retrieved_context)
        
        # Gera a resposta
        response = generate_response(llm, prompt)
        
        # Processa a resposta
        final_answer = process_response(response, user_question)
        
        # Cria mensagem de resposta
        ai_message = AIMessage(content=final_answer)
        
        log_node_execution(
            "LLM_AGENT", 
            "SUCCESS", 
            f"Resposta gerada: {len(final_answer)} caracteres"
        )
        
        return {
            "final_answer": final_answer,
            "messages": [ai_message],
            "processing_status": ProcessingStatus.COMPLETED,
            "error_message": None
        }
        
    except Exception as e:
        error_msg = f"Erro na geração de resposta: {str(e)}"
        log_node_execution("LLM_AGENT", "ERROR", error_msg)
        main_logger.exception("Erro detalhado na geração de resposta:")
        
        return {
            "processing_status": ProcessingStatus.ERROR,
            "error_message": error_msg
        }


def create_llm_model() -> ChatOpenAI:
    """
    Cria e configura o modelo LLM GPT-4o-mini.
    
    Returns:
        ChatOpenAI: Modelo LLM configurado
    """
    model_config = Config.get_model_config()
    
    llm = ChatOpenAI(
        openai_api_key=get_openai_api_key(),
        model_name=model_config["model"],
        temperature=model_config["temperature"],
        max_tokens=model_config["max_tokens"],
        timeout=60,
        max_retries=3
    )
    
    main_logger.debug(f"Modelo LLM criado: {model_config['model']}")
    return llm


def build_prompt(question: str, context_chunks: list) -> ChatPromptTemplate:
    """
    Constrói um prompt otimizado para o LLM.
    
    Args:
        question: Pergunta do usuário
        context_chunks: Lista de chunks de contexto
    
    Returns:
        ChatPromptTemplate: Prompt construído
    """
    # Combina o contexto
    context_text = "\n\n".join(context_chunks) if context_chunks else ""
    
    # Sistema de prompt em português
    system_prompt = """Você é um assistente especializado em análise de documentos PDF. Sua função é responder perguntas baseadas exclusivamente no conteúdo fornecido.

INSTRUÇÕES IMPORTANTES:
1. Use APENAS as informações do contexto fornecido para responder
2. Se a informação não estiver no contexto, diga claramente que não encontrou a informação no documento
3. Seja preciso, claro e objetivo em suas respostas
4. Cite trechos relevantes do documento quando apropriado
5. Mantenha um tom profissional e educativo
6. Se a pergunta for ambígua, peça esclarecimentos
7. Organize sua resposta de forma estruturada quando necessário

FORMATO DA RESPOSTA:
- Responda diretamente à pergunta
- Use parágrafos para organizar ideias complexas
- Inclua citações do documento quando relevante
- Termine com um resumo se a resposta for longa"""

    # Template do prompt
    prompt_template = ChatPromptTemplate.from_messages([
        ("system", system_prompt),
        ("human", """CONTEXTO DO DOCUMENTO:
{context}

PERGUNTA DO USUÁRIO:
{question}

Por favor, responda à pergunta baseando-se exclusivamente no contexto fornecido.""")
    ])
    
    return prompt_template.partial(context=context_text, question=question)


def generate_response(llm: ChatOpenAI, prompt: ChatPromptTemplate) -> str:
    """
    Gera a resposta usando o LLM.
    
    Args:
        llm: Modelo LLM
        prompt: Prompt construído
    
    Returns:
        str: Resposta gerada
    """
    try:
        # Cria a chain
        chain = prompt | llm
        
        # Gera a resposta
        response = chain.invoke({})
        
        # Extrai o conteúdo da resposta
        if hasattr(response, 'content'):
            return response.content
        else:
            return str(response)
            
    except Exception as e:
        main_logger.error(f"Erro na geração da resposta: {e}")
        raise


def process_response(response: str, original_question: str) -> str:
    """
    Processa e valida a resposta gerada.
    
    Args:
        response: Resposta bruta do LLM
        original_question: Pergunta original do usuário
    
    Returns:
        str: Resposta processada e validada
    """
    if not response or not response.strip():
        return "Desculpe, não consegui gerar uma resposta adequada para sua pergunta."
    
    # Limpa a resposta
    cleaned_response = response.strip()
    
    # Valida se a resposta é adequada
    if len(cleaned_response) < 20:
        return f"Resposta muito curta gerada. Pergunta original: {original_question}\n\nResposta: {cleaned_response}"
    
    # Adiciona informações contextuais se necessário
    if "não encontrei" in cleaned_response.lower() or "não há informação" in cleaned_response.lower():
        cleaned_response += "\n\n💡 **Dica**: Tente reformular sua pergunta ou verificar se o PDF contém a informação desejada."
    
    return cleaned_response


def create_fallback_response(question: str, error_msg: str = None) -> str:
    """
    Cria uma resposta de fallback quando há erro.
    
    Args:
        question: Pergunta original
        error_msg: Mensagem de erro opcional
    
    Returns:
        str: Resposta de fallback
    """
    base_response = f"""Desculpe, encontrei dificuldades para processar sua pergunta: "{question}"

Isso pode ter acontecido por alguns motivos:
1. O documento PDF pode não conter informações relacionadas à sua pergunta
2. Pode haver um problema temporário com o processamento
3. A pergunta pode precisar ser mais específica

**Sugestões:**
- Tente reformular sua pergunta de forma mais específica
- Verifique se o PDF foi carregado corretamente
- Certifique-se de que o documento contém a informação desejada"""

    if error_msg:
        base_response += f"\n\n**Detalhes técnicos:** {error_msg}"
    
    return base_response


def validate_response_quality(response: str, question: str) -> tuple[bool, str]:
    """
    Valida a qualidade da resposta gerada.
    
    Args:
        response: Resposta gerada
        question: Pergunta original
    
    Returns:
        tuple[bool, str]: (é_válida, motivo_se_inválida)
    """
    if not response or len(response.strip()) < 10:
        return False, "Resposta muito curta ou vazia"
    
    # Verifica se a resposta é apenas uma repetição da pergunta
    if question.lower() in response.lower() and len(response) < len(question) * 2:
        return False, "Resposta parece ser apenas repetição da pergunta"
    
    # Verifica se há conteúdo substantivo
    words = response.split()
    if len(words) < 5:
        return False, "Resposta com muito poucas palavras"
    
    # Verifica padrões de resposta inadequada
    inadequate_patterns = [
        "não posso responder",
        "não tenho informação",
        "desculpe, mas não",
        "não é possível"
    ]
    
    response_lower = response.lower()
    inadequate_count = sum(1 for pattern in inadequate_patterns if pattern in response_lower)
    
    if inadequate_count > 1:
        return False, "Resposta contém muitos padrões de inadequação"
    
    return True, "Resposta válida"


def enhance_response_with_metadata(response: str, context_used: bool, num_sources: int) -> str:
    """
    Melhora a resposta adicionando metadados úteis.
    
    Args:
        response: Resposta original
        context_used: Se contexto foi usado
        num_sources: Número de fontes consultadas
    
    Returns:
        str: Resposta melhorada
    """
    enhanced_response = response
    
    # Adiciona informação sobre as fontes
    if context_used and num_sources > 0:
        enhanced_response += f"\n\n---\n📚 *Resposta baseada em {num_sources} seção(ões) do documento.*"
    elif not context_used:
        enhanced_response += "\n\n---\n⚠️ *Resposta gerada sem contexto específico do documento.*"
    
    return enhanced_response