File size: 6,363 Bytes
c6828a0 |
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 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 |
import os
import sys
import re
import logging
from typing import Optional, Tuple
try:
from sqlalchemy import text as sa_text
except Exception:
sa_text = None
try:
# Intentar importar dependencias opcionales
from langchain_community.agent_toolkits import create_sql_agent
from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.agents.agent_types import AgentType
from langchain.memory import ConversationBufferWindowMemory
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
import pymysql
from dotenv import load_dotenv
DEPENDENCIES_AVAILABLE = True
except ImportError as e:
logging.warning(f"Some dependencies are not available: {e}")
DEPENDENCIES_AVAILABLE = False
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
def check_environment():
"""Verifica si el entorno est谩 configurado correctamente."""
if not DEPENDENCIES_AVAILABLE:
return False, "Missing required Python packages. Please install them with: pip install -r requirements.txt"
# Verificar si estamos en un entorno con variables de entorno
required_vars = ["DB_USER", "DB_PASSWORD", "DB_HOST", "DB_NAME", "GOOGLE_API_KEY"]
missing_vars = [var for var in required_vars if not os.getenv(var)]
if missing_vars:
return False, f"Missing required environment variables: {', '.join(missing_vars)}"
return True, "Environment is properly configured"
def setup_database_connection():
"""Intenta establecer una conexi贸n a la base de datos."""
if not DEPENDENCIES_AVAILABLE:
return None, "Dependencies not available"
try:
load_dotenv(override=True)
db_user = os.getenv("DB_USER")
db_password = os.getenv("DB_PASSWORD")
db_host = os.getenv("DB_HOST")
db_name = os.getenv("DB_NAME")
if not all([db_user, db_password, db_host, db_name]):
missing = [var for var, val in [
("DB_USER", db_user),
("DB_PASSWORD", "*" if db_password else ""),
("DB_HOST", db_host),
("DB_NAME", db_name)
] if not val]
logger.error(f"Missing required database configuration: {', '.join(missing)}")
return None, f"Missing database configuration: {', '.join(missing)}"
logger.info(f"Connecting to database: {db_user}@{db_host}/{db_name}")
# Probar conexi贸n
connection = pymysql.connect(
host=db_host,
user=db_user,
password=db_password,
database=db_name,
connect_timeout=5,
cursorclass=pymysql.cursors.DictCursor
)
connection.close()
# Si la conexi贸n es exitosa, crear motor SQLAlchemy
db_uri = f"mysql+pymysql://{db_user}:{db_password}@{db_host}/{db_name}"
logger.info("Database connection successful")
return SQLDatabase.from_uri(db_uri), ""
except Exception as e:
error_msg = f"Error connecting to database: {str(e)}"
logger.error(error_msg)
return None, error_msg
def initialize_llm():
"""Inicializa el modelo de lenguaje."""
if not DEPENDENCIES_AVAILABLE:
error_msg = "Dependencies not available. Make sure all required packages are installed."
logger.error(error_msg)
return None, error_msg
google_api_key = os.getenv("GOOGLE_API_KEY")
logger.info(f"GOOGLE_API_KEY found: {'Yes' if google_api_key else 'No'}")
if not google_api_key:
error_msg = "GOOGLE_API_KEY not found in environment variables. Please check your Hugging Face Space secrets."
logger.error(error_msg)
return None, error_msg
try:
logger.info("Initializing Google Generative AI...")
llm = ChatGoogleGenerativeAI(
model="gemini-2.0-flash",
temperature=0,
google_api_key=google_api_key,
convert_system_message_to_human=True
)
# Test the model with a simple prompt
test_prompt = "Hello, this is a test."
logger.info(f"Testing model with prompt: {test_prompt}")
test_response = llm.invoke(test_prompt)
logger.info(f"Model test response: {str(test_response)[:100]}...")
logger.info("Google Generative AI initialized successfully")
return llm, ""
except Exception as e:
error_msg = f"Error initializing Google Generative AI: {str(e)}"
logger.error(error_msg, exc_info=True)
return None, error_msg
def create_agent(llm, db_connection):
"""Create and return a SQL database agent with conversation memory."""
if not llm:
error_msg = "Cannot create agent: LLM is not available"
logger.error(error_msg)
return None, error_msg
if not db_connection:
error_msg = "Cannot create agent: Database connection is not available"
logger.error(error_msg)
return None, error_msg
try:
logger.info("Creating SQL agent with memory...")
# Create conversation memory
memory = ConversationBufferWindowMemory(
memory_key="chat_history",
k=5,
return_messages=True,
output_key="output"
)
# Create the database toolkit
toolkit = SQLDatabaseToolkit(
db=db_connection,
llm=llm
)
# Create the agent
agent = create_sql_agent(
llm=llm,
toolkit=toolkit,
agent_type=AgentType.OPENAI_FUNCTIONS,
verbose=True,
handle_parsing_errors=True,
max_iterations=10,
early_stopping_method="generate",
memory=memory,
return_intermediate_steps=True
)
logger.info("SQL agent created successfully")
return agent, ""
except Exception as e:
error_msg = f"Error creating SQL agent: {str(e)}"
logger.error(error_msg, exc_info=True)
return None, error_msg |