| 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: |
| |
| 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 |
|
|
| |
| 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" |
| |
| |
| 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}") |
| |
| |
| connection = pymysql.connect( |
| host=db_host, |
| user=db_user, |
| password=db_password, |
| database=db_name, |
| connect_timeout=5, |
| cursorclass=pymysql.cursors.DictCursor |
| ) |
| connection.close() |
| |
| |
| 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_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...") |
| |
| |
| memory = ConversationBufferWindowMemory( |
| memory_key="chat_history", |
| k=5, |
| return_messages=True, |
| output_key="output" |
| ) |
| |
| |
| toolkit = SQLDatabaseToolkit( |
| db=db_connection, |
| llm=llm |
| ) |
| |
| |
| 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 |
