agents/state.py

"""
Definições de estado para o AgentPDF usando LangGraph.

Este módulo define as estruturas de estado que serão utilizadas
pelos nós do grafo para compartilhar informações durante a execução.
"""

from typing import List, Optional, Dict, Any
from typing_extensions import TypedDict
from langchain_core.messages import BaseMessage
from langgraph.graph.message import add_messages
from typing import Annotated


class PDFState(TypedDict):
    """
    Estado principal do AgentPDF.
    
    Contém todas as informações necessárias para o processamento
    de PDFs e geração de respostas.
    """
    # Mensagens da conversa
    messages: Annotated[List[BaseMessage], add_messages]
    
    # Informações do PDF
    pdf_path: Optional[str]
    pdf_text: Optional[str]
    pdf_chunks: Optional[List[str]]
    
    # Vector store e embeddings
    vector_store: Optional[Any]
    embeddings_created: bool
    
    # Contexto recuperado
    retrieved_context: Optional[List[str]]
    
    # Pergunta do usuário
    user_question: Optional[str]
    
    # Resposta final
    final_answer: Optional[str]
    
    # Status do processamento
    processing_status: str
    error_message: Optional[str]


class ProcessingState(TypedDict):
    """
    Estado específico para processamento de documentos.
    """
    document_path: str
    extracted_text: Optional[str]
    text_chunks: Optional[List[str]]
    chunk_size: int
    chunk_overlap: int


class RetrievalState(TypedDict):
    """
    Estado específico para recuperação de contexto.
    """
    query: str
    retrieved_docs: Optional[List[str]]
    similarity_scores: Optional[List[float]]
    top_k: int


class LLMState(TypedDict):
    """
    Estado específico para interação com o LLM.
    """
    system_prompt: str
    user_query: str
    context: Optional[str]
    response: Optional[str]
    model_name: str


# Estados de entrada e saída para diferentes nós
class InputState(TypedDict):
    """Estado de entrada para o grafo."""
    messages: Annotated[List[BaseMessage], add_messages]
    pdf_path: Optional[str]


class OutputState(TypedDict):
    """Estado de saída do grafo."""
    messages: Annotated[List[BaseMessage], add_messages]
    final_answer: str


# Configurações padrão
DEFAULT_CHUNK_SIZE = 1000
DEFAULT_CHUNK_OVERLAP = 200
DEFAULT_TOP_K = 5
DEFAULT_MODEL = "gpt-4o-mini"

# Status de processamento
class ProcessingStatus:
    IDLE = "idle"
    LOADING_PDF = "loading_pdf"
    PROCESSING_TEXT = "processing_text"
    CREATING_EMBEDDINGS = "creating_embeddings"
    RETRIEVING_CONTEXT = "retrieving_context"
    GENERATING_RESPONSE = "generating_response"
    COMPLETED = "completed"
    ERROR = "error"