add missing function

app.py

@@ -22,6 +22,164 @@ logger = logging.getLogger(__name__)
 # Flask app initialization
 flask_app = Flask(__name__)
 
+# Importaciones adicionales necesarias
+import os
+from dotenv import load_dotenv
+from sqlalchemy import create_engine, text
+from sqlalchemy.exc import SQLAlchemyError
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_community.agent_toolkits import create_sql_agent
+from langchain_community.utilities import SQLDatabase
+from langgraph.prebuilt import create_react_agent
+
+# Cargar variables de entorno
+load_dotenv()
+
+def initialize_llm():
+    """Inicializar el modelo LLM de Google Gemini."""
+    try:
+        api_key = os.getenv('GOOGLE_API_KEY')
+        if not api_key:
+            return None, "No se encontró GOOGLE_API_KEY en las variables de entorno"
+        
+        llm = ChatGoogleGenerativeAI(
+            model="gemini-2.0-flash",
+            google_api_key=api_key,
+            temperature=0.1,
+            convert_system_message_to_human=True
+        )
+        return llm, None
+    except Exception as e:
+        return None, str(e)
+
+def setup_database_connection():
+    """Configurar la conexión a la base de datos."""
+    try:
+        # Obtener configuración de la base de datos
+        db_user = os.getenv('DB_USER')
+        db_password = os.getenv('DB_PASSWORD')
+        db_host = os.getenv('DB_HOST', 'localhost')
+        db_name = os.getenv('DB_NAME')
+        
+        if not all([db_user, db_password, db_name]):
+            return None, "Faltan variables de entorno para la conexión a la base de datos"
+        
+        # Crear string de conexión
+        connection_string = f"mysql+pymysql://{db_user}:{db_password}@{db_host}/{db_name}"
+        
+        # Crear engine de SQLAlchemy y probar la conexión
+        engine = create_engine(connection_string)
+        with engine.connect() as conn:
+            conn.execute(text("SELECT 1"))
+        
+        return connection_string, None
+    except Exception as e:
+        return None, str(e)
+
+def create_agent(llm, connection_string):
+    """Crear el agente SQL."""
+    try:
+        if not llm or not connection_string:
+            return None, "LLM o conexión a base de datos no proporcionados"
+        
+        # Crear la base de datos SQL
+        db = SQLDatabase.from_uri(connection_string)
+        
+        # Crear el agente SQL
+        agent = create_sql_agent(
+            llm=llm,
+            db=db,
+            agent_type="zero-shot-react-description",
+            verbose=True,
+            return_intermediate_steps=True
+        )
+        
+        return agent, None
+    except Exception as e:
+        return None, str(e)
+
+def stream_agent_response(question: str, chat_history: List[List[str]] = None) -> Tuple[str, Optional[go.Figure]]:
+    """Procesar la respuesta del agente y generar visualizaciones si es necesario."""
+    try:
+        # Inicializar componentes
+        llm, llm_error = initialize_llm()
+        if llm_error:
+            return f"Error al inicializar LLM: {llm_error}", None
+            
+        connection_string, db_error = setup_database_connection()
+        if db_error:
+            return f"Error de conexión a base de datos: {db_error}", None
+            
+        agent, agent_error = create_agent(llm, connection_string)
+        if agent_error:
+            return f"Error al crear el agente: {agent_error}", None
+        
+        # Ejecutar la consulta
+        response = agent.invoke({"input": question})
+        
+        # Extraer la respuesta
+        if hasattr(response, 'output'):
+            response_text = response.output
+        elif isinstance(response, dict) and 'output' in response:
+            response_text = response['output']
+        else:
+            response_text = str(response)
+        
+        # Verificar si hay resultados de SQL que podrían visualizarse
+        chart_fig = None
+        if hasattr(response, 'intermediate_steps'):
+            for step in response.intermediate_steps:
+                if len(step) > 1 and 'sql_query' in str(step[0]).lower():
+                    # Intentar ejecutar la consulta y crear visualización
+                    try:
+                        query = str(step[0]).split('sql_query:')[1].split('\n')[0].strip()
+                        if 'SELECT' in query.upper():
+                            df = pd.read_sql_query(query, create_engine(connection_string))
+                            if len(df.columns) >= 2:
+                                fig = px.bar(df, x=df.columns[0], y=df.columns[1])
+                                chart_fig = fig
+                    except:
+                        pass
+        
+        return response_text, chart_fig
+        
+    except Exception as e:
+        return f"Error al procesar la solicitud: {str(e)}", None
+
+def create_ui():
+    """Crear la interfaz de usuario de Gradio."""
+    with gr.Blocks(title="🤖 Asistente SQL con Gemini", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("# 🤖 Asistente de Base de Datos SQL")
+        gr.Markdown("Pregunta cualquier cosa sobre tu base de datos en lenguaje natural")
+        
+        with gr.Row():
+            with gr.Column(scale=2):
+                chatbot = gr.Chatbot(
+                    label="Chat",
+                    type="messages",
+                    height=400
+                )
+                
+                with gr.Row():
+                    question_input = gr.Textbox(
+                        label="Tu pregunta",
+                        placeholder="Ej: ¿Cuántos usuarios hay registrados?",
+                        lines=2,
+                        scale=4
+                    )
+                    submit_button = gr.Button("📤 Enviar", scale=1)
+                
+            with gr.Column(scale=1):
+                chart_display = gr.Plot(
+                    label="Visualización de datos",
+                    height=400
+                )
+        
+        # Campo oculto para el estado de salida
+        streaming_output_display = gr.HTML(visible=False)
+        
+        return demo, chatbot, chart_display, question_input, submit_button, streaming_output_display
+
 # Almacenamiento en memoria de los mensajes
 message_store: Dict[str, str] = {}
 
@@ -68,11 +226,11 @@ def handle_ask():
         if llm_error:
             return jsonify({'error': f'Error al inicializar LLM: {llm_error}'}), 500
             
-        db_connection, db_error = setup_database_connection()
+        connection_string, db_error = setup_database_connection()
         if db_error:
             return jsonify({'error': f'Error de conexión a la base de datos: {db_error}'}), 500
             
-        agent, agent_error = create_agent(llm, db_connection)
+        agent, agent_error = create_agent(llm, connection_string)
         if agent_error:
             return jsonify({'error': f'Error al crear el agente: {agent_error}'}), 500
         
@@ -130,7 +288,7 @@ def create_application():
 
         return "", chat_history
     
-    async def bot_response(chat_history: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], Optional[go.Figure]]:
+    def bot_response(chat_history: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], Optional[go.Figure]]:
         """Generate bot response for messages-format chat history and return optional chart figure."""
         if not chat_history:
             return chat_history, None
@@ -162,7 +320,7 @@ def create_application():
                     i += 1
 
             # Call the agent for this new user question
-            assistant_message, chart_fig = await stream_agent_response(question, pair_history)
+            assistant_message, chart_fig = stream_agent_response(question, pair_history)
 
             # Append assistant message back into messages history
             chat_history.append({"role": "assistant", "content": assistant_message})
@@ -184,7 +342,7 @@ def create_application():
             fn=user_message,
             inputs=[question_input, chatbot],
             outputs=[question_input, chatbot],
-            queue=True
+            queue=False
         ).then(
             fn=bot_response,
             inputs=[chatbot],
@@ -197,7 +355,7 @@ def create_application():
             fn=user_message,
             inputs=[question_input, chatbot],
             outputs=[question_input, chatbot],
-            queue=True
+            queue=False
         ).then(
             fn=bot_response,
             inputs=[chatbot],