Spaces:
Sleeping
Sleeping
| """ | |
| Definições do estado do agente e funções de coordenação geral | |
| """ | |
| from typing import Dict, Any, Optional, TypedDict | |
| class AgentState(TypedDict): | |
| """Estado do agente LangGraph - apenas dados serializáveis""" | |
| user_input: str | |
| selected_model: str | |
| response: str | |
| advanced_mode: bool | |
| execution_time: float | |
| error: Optional[str] | |
| intermediate_steps: list | |
| # Dados serializáveis do banco | |
| db_sample_dict: dict | |
| # IDs para recuperar objetos não-serializáveis | |
| agent_id: str | |
| engine_id: str | |
| cache_id: str | |
| # Campos relacionados ao refinamento de pergunta | |
| question_refinement_enabled: bool # Se o refinamento está habilitado | |
| original_user_input: Optional[str] # Pergunta original antes do refinamento | |
| refined_question: Optional[str] # Pergunta refinada pelo GPT-4o | |
| question_refinement_applied: bool # Se o refinamento foi aplicado | |
| question_refinement_changes: list # Lista de mudanças feitas | |
| question_refinement_justification: Optional[str] # Justificativa do refinamento | |
| question_refinement_success: bool # Se o refinamento foi bem-sucedido | |
| question_refinement_error: Optional[str] # Erro no refinamento | |
| question_refinement_has_significant_change: bool # Se houve mudança significativa | |
| # Campos relacionados a gráficos | |
| query_type: str # 'sql_query', 'sql_query_graphic', 'prediction' | |
| sql_query_extracted: Optional[str] # Query SQL extraída da resposta do agente | |
| graph_type: Optional[str] # Tipo de gráfico escolhido pela LLM | |
| graph_data: Optional[dict] # Dados preparados para o gráfico (serializável) | |
| graph_image_id: Optional[str] # ID da imagem do gráfico no ObjectManager | |
| graph_generated: bool # Se o gráfico foi gerado com sucesso | |
| graph_error: Optional[str] # Erro na geração de gráfico | |
| # Campos relacionados ao cache | |
| cache_hit: bool # Se houve hit no cache | |
| # Campos relacionados ao Processing Agent | |
| processing_enabled: bool # Se o Processing Agent está habilitado | |
| processing_model: str # Modelo usado no Processing Agent | |
| processing_agent_id: Optional[str] # ID do Processing Agent no ObjectManager | |
| suggested_query: Optional[str] # Query SQL sugerida pelo Processing Agent | |
| query_observations: Optional[str] # Observações sobre a query sugerida | |
| processing_result: Optional[dict] # Resultado completo do Processing Agent | |
| processing_success: bool # Se o processamento foi bem-sucedido | |
| processing_error: Optional[str] # Erro no processamento | |
| # Campos relacionados ao refinamento | |
| refined: bool # Se a resposta foi refinada | |
| refinement_error: Optional[str] # Erro no refinamento | |
| refinement_quality: Optional[str] # Qualidade do refinamento | |
| quality_metrics: Optional[dict] # Métricas de qualidade | |
| # Campos relacionados ao contexto SQL | |
| sql_context: Optional[str] # Contexto preparado para o agente SQL | |
| sql_result: Optional[dict] # Resultado do agente SQL | |
| # Campos relacionados ao tipo de conexão | |
| connection_type: str # "csv" | "postgresql" | |
| postgresql_config: Optional[dict] # Configuração PostgreSQL | |
| selected_table: Optional[str] # Tabela selecionada (para PostgreSQL) | |
| single_table_mode: bool # Se deve usar apenas uma tabela (PostgreSQL) | |
| connection_success: bool # Se a conexão foi estabelecida com sucesso | |
| connection_error: Optional[str] # Erro na conexão | |
| connection_info: Optional[dict] # Informações da conexão estabelecida | |
| def should_refine_response(state: Dict[str, Any]) -> str: | |
| """ | |
| Determina se deve refinar a resposta | |
| Args: | |
| state: Estado atual | |
| Returns: | |
| Nome do próximo nó | |
| """ | |
| if state.get("advanced_mode", False) and not state.get("error"): | |
| return "refine_response" | |
| else: | |
| return "cache_response" | |
| def should_generate_graph(state: Dict[str, Any]) -> str: | |
| """ | |
| Determina se deve gerar gráfico | |
| Args: | |
| state: Estado atual | |
| Returns: | |
| Nome do próximo nó | |
| """ | |
| query_type = state.get("query_type", "") | |
| if query_type == "sql_query_graphic" and not state.get("error"): | |
| return "graph_selection" | |
| elif state.get("advanced_mode", False) and not state.get("error"): | |
| return "refine_response" | |
| else: | |
| return "cache_response" | |
| def should_use_processing_agent(state: Dict[str, Any]) -> str: | |
| """ | |
| Determina se deve usar o Processing Agent | |
| Args: | |
| state: Estado atual | |
| Returns: | |
| Nome do próximo nó | |
| """ | |
| if state.get("processing_enabled", False): | |
| return "validate_processing" | |
| else: | |
| return "prepare_context" | |
| def should_refine_question(state: Dict[str, Any]) -> str: | |
| """ | |
| Determina se deve refinar a pergunta | |
| Args: | |
| state: Estado atual | |
| Returns: | |
| Nome do próximo nó | |
| """ | |
| if state.get("question_refinement_enabled", False): | |
| return "question_refinement" | |
| else: | |
| # Pula refinamento e vai para validação de processing | |
| return "validate_processing" | |
| def route_after_cache_check(state: Dict[str, Any]) -> str: | |
| """ | |
| Roteamento após verificação de cache | |
| CACHE TEMPORARIAMENTE DESATIVADO - Sempre ignora cache hit | |
| Args: | |
| state: Estado atual | |
| Returns: | |
| Nome do próximo nó | |
| """ | |
| import logging | |
| cache_hit = state.get("cache_hit", False) | |
| processing_enabled = state.get("processing_enabled", False) | |
| question_refinement_enabled = state.get("question_refinement_enabled", False) | |
| # DESATIVAÇÃO TEMPORÁRIA DO CACHE | |
| # Força cache_hit = False para sempre processar queries | |
| cache_hit = False | |
| logging.info(f"[ROUTING] Cache hit: {cache_hit} (CACHE DESATIVADO TEMPORARIAMENTE)") | |
| logging.info(f"[ROUTING] Processing enabled: {processing_enabled}") | |
| logging.info(f"[ROUTING] Question refinement enabled: {question_refinement_enabled}") | |
| if cache_hit: | |
| logging.info("[ROUTING] Direcionando para update_history (cache hit)") | |
| return "update_history" | |
| # Se não tem conexão, vai para seleção de conexão | |
| if not state.get("agent_id") or not state.get("engine_id"): | |
| logging.info("[ROUTING] Direcionando para connection_selection (sem conexão)") | |
| return "connection_selection" | |
| # Se refinamento está habilitado E ainda não foi aplicado, vai para refinamento primeiro | |
| if question_refinement_enabled and not state.get("question_refinement_applied", False): | |
| logging.info("[ROUTING] Direcionando para question_refinement (refinamento habilitado)") | |
| return "question_refinement" | |
| elif processing_enabled: | |
| logging.info("[ROUTING] Direcionando para validate_processing (processing habilitado)") | |
| return "validate_processing" | |
| else: | |
| logging.info("[ROUTING] Direcionando para connection_selection (fluxo direto)") | |
| return "connection_selection" | |