Spaces:
Sleeping
Sleeping
File size: 5,517 Bytes
e982206 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
"""
Nó para seleção do tipo de conexão (csv ou postgresql)
"""
import logging
from typing import Dict, Any
from utils.validation import validate_connection_state
async def connection_selection_node(state: Dict[str, Any]) -> Dict[str, Any]:
"""
Nó para determinar o tipo de conexão baseado na entrada do usuário
Args:
state: Estado atual do agente
Returns:
Estado atualizado com tipo de conexão definido
"""
try:
logging.info("[CONNECTION_SELECTION] Iniciando seleção de tipo de conexão")
# Verifica se o tipo de conexão já foi definido
connection_type = state.get("connection_type")
if not connection_type:
# Se não foi definido, assume csv como padrão (compatibilidade)
connection_type = "csv"
logging.info("[CONNECTION_SELECTION] Tipo de conexão não definido, usando csv como padrão")
# Valida tipo de conexão
valid_types = ["csv", "postgresql"]
if connection_type.upper() not in [t.upper() for t in valid_types]:
error_msg = f"Tipo de conexão inválido: {connection_type}. Tipos válidos: {valid_types}"
logging.error(f"[CONNECTION_SELECTION] {error_msg}")
state.update({
"connection_type": "csv", # Fallback para csv
"connection_error": error_msg,
"connection_success": False
})
return state
# Atualiza estado com tipo de conexão validado
state.update({
"connection_type": connection_type,
"connection_error": None,
"connection_success": True
})
logging.info(f"[CONNECTION_SELECTION] Tipo de conexão selecionado: {connection_type}")
return state
except Exception as e:
error_msg = f"Erro na seleção de tipo de conexão: {e}"
logging.error(f"[CONNECTION_SELECTION] {error_msg}")
# Fallback para csv em caso de erro
state.update({
"connection_type": "csv",
"connection_error": error_msg,
"connection_success": False
})
return state
def route_by_connection_type(state: Dict[str, Any]) -> str:
"""
Função de roteamento baseada no tipo de conexão
Args:
state: Estado atual do agente
Returns:
Nome do próximo nó baseado no tipo de conexão
"""
connection_type = state.get("connection_type", "csv")
file_path = state.get("file_path")
db_id = state.get("db_id")
engine_id = state.get("engine_id")
logging.info(f"[CONNECTION_ROUTING] Roteando para tipo: {connection_type}")
logging.info(f"[CONNECTION_ROUTING] DB ID existente: {db_id}")
logging.info(f"[CONNECTION_ROUTING] Engine ID existente: {engine_id}")
# Se já tem conexão estabelecida, pula para get_db_sample
# Verifica se o sistema já foi inicializado
from utils.object_manager import get_object_manager
obj_manager = get_object_manager()
# Verifica se há agentes SQL já criados (indicando sistema inicializado)
stats = obj_manager.get_stats()
has_sql_agents = stats.get("sql_agents", 0) > 0
has_databases = stats.get("databases", 0) > 0
if has_sql_agents and has_databases:
logging.info("[CONNECTION_ROUTING] Sistema já inicializado com agentes e bancos, pulando para get_db_sample")
return "get_db_sample"
# Fallback: verifica IDs específicos
if db_id and engine_id:
logging.info("[CONNECTION_ROUTING] Conexão já estabelecida via IDs, pulando para get_db_sample")
return "get_db_sample"
if connection_type.upper() == "POSTGRESQL":
return "postgresql_connection"
elif file_path:
# Há arquivo csv para processar
return "csv_processing"
else:
# Usar banco existente
return "load_database"
async def validate_connection_input_node(state: Dict[str, Any]) -> Dict[str, Any]:
"""
Nó para validar entrada de conexão antes do processamento
Args:
state: Estado atual do agente
Returns:
Estado atualizado com validação
"""
try:
logging.info("[CONNECTION_VALIDATION] Validando entrada de conexão")
connection_type = state.get("connection_type", "csv")
# Usa validação centralizada
is_valid, validation_error = validate_connection_state(state)
if not is_valid:
logging.error(f"[CONNECTION_VALIDATION] {validation_error}")
state.update({
"connection_error": validation_error,
"connection_success": False
})
return state
logging.info(f"[CONNECTION_VALIDATION] Validação de conexão {connection_type} bem-sucedida")
# Validação bem-sucedida
state.update({
"connection_error": None,
"connection_success": True
})
logging.info("[CONNECTION_VALIDATION] Validação de conexão concluída com sucesso")
return state
except Exception as e:
error_msg = f"Erro na validação de conexão: {e}"
logging.error(f"[CONNECTION_VALIDATION] {error_msg}")
state.update({
"connection_error": error_msg,
"connection_success": False
})
return state
|