shared.py

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