app.py

import os
from dotenv import load_dotenv
from google.adk.sessions import InMemorySessionService
from google.adk.runners import Runner
from google.adk.agents import LlmAgent
from google.adk.tools import FunctionTool
from google.adk.models.lite_llm import LiteLlm
from typing import Dict, Any
from datetime import date, datetime
from google import genai
from google.genai import types
import requests

import google.generativeai as genai
from pydantic import BaseModel
import json
import uuid
import asyncio
import gradio as gr


# Cargar variables de entorno desde .env
load_dotenv()

from datetime import datetime

def obtener_fecha():
    return datetime.now().strftime("%Y-%m-%d")

# Establecer variables de entorno necesarias
#os.environ['GOOGLE_API_KEY'] = os.getenv("GOOGLE_API_KEY")
#os.environ['GOOGLE_GENAI_USE_VERTEXAI'] = os.getenv("GOOGLE_GENAI_USE_VERTEXAI")
#genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))

# Inicializar cliente de BigQuery - ACTUALIZA CON TU PROJECT ID
#client_bq = bigquery.Client(project="uma-datascience-dev")  # Cambia por tu project ID

class SQLResult(BaseModel):
    consulta_sql: str

# Definir función para generar SQL
def generar_sql(pregunta: str) -> Dict[str, Any]:
    """
    Genera una consulta SQL a partir de una pregunta en lenguaje natural usando un modelo de lenguaje (Gemini).

    Esta función toma una pregunta sobre datos médicos/administrativos y utiliza un LLM configurado con un prompt
    que describe el contexto, el esquema de la tabla y reglas de generación para producir una consulta SQL válida
    y segura que pueda ejecutarse sobre BigQuery.

    Parámetros:
    -----------
    pregunta : str
        Pregunta del usuario en lenguaje natural relacionada con la tabla 
        `uma-datascience-dev.Bigquery_Dataset.evaluaciones_auditoria`.

    Retorna:
    --------
    Dict[str, Any]
        Un diccionario con al menos la clave `consulta_sql`, que contiene la consulta SQL generada.
        Ejemplo:
        {
            "consulta_sql": "SELECT COUNT(*) FROM ..."
        }

    Requiere:
    ---------
    - Que esté configurada correctamente la API de Gemini (`google.generativeai`) y la clave esté cargada.
    - Que exista un esquema de Pydantic `SQLResult` para validar la estructura esperada.

    Nota:
    -----
    - Usa el modelo `gemini-2.0-flash-lite`.
    - Aplica reglas estrictas de formato y nombres de campos según el schema de BigQuery.
    - La consulta es generada en base al contexto, fecha actual y convenciones de codificación establecidas.
    """
    prompt = f"""
Sos un asistente experto en análisis de datos médicos/administrativos que responde preguntas en lenguaje natural
transformándolas en consultas SQL para BigQuery.

Tabla disponible: uma-datascience-dev.Bigquery_Dataset.evaluaciones_auditoria

Columnas disponibles:
- Fecha (DATE): Fecha del registro
- Practica_solicitada (STRING): Descripción de la práctica médica solicitada
- DNI (INTEGER): Documento Nacional de Identidad del paciente
- Profesional (STRING): Nombre del profesional médico
- Estado (STRING): Estado actual del registro/solicitud

🔎 Valores posibles para la columna `Estado`:
- "aprobado"
- "rechazado"

(No existen otros estados. No inventes valores como "pendiente", "en revisión", etc.)

Algunos valores de la columna 'practica_solicitada' son:

Análisis de sangre
Colonoscopia
Ecografía abdominal
Electrocardiograma
Endoscopia digestiva alta
Mamografía
Radiografía de cráneo
Radiografía de rodilla
Resonancia de cerebro
Resonancia de columna lumbar
Tomografía de tórax


⚠️ Importante:
- Siempre usá nombres de tabla completamente calificados: uma-datascience-dev.Bigquery_Dataset.evaluaciones_auditoria
- Usá comillas invertidas para referenciar tablas y campos cuando sea necesario
- Para fechas, usá el formato DATE en las consultas (ej: DATE('2024-01-01'))
- Si la pregunta es ambigua, asumí una interpretación razonable y explícita
- No inventes campos que no existen en el schema

📘 Ejemplos de preguntas y sus consultas SQL:

- Pregunta: ¿Cuántas prácticas se solicitaron este mes?

Esperamos una respuesta tipo:

SELECT COUNT(*) AS total_practicas FROM uma-datascience-dev.Bigquery_Dataset.evaluaciones_auditoria WHERE EXTRACT(MONTH FROM Fecha) = EXTRACT(MONTH FROM CURRENT_DATE()) AND EXTRACT(YEAR FROM Fecha) = EXTRACT(YEAR FROM CURRENT_DATE())


- Pregunta: ¿Cuántas prácticas se aprobaron este año?

Esperamos una respuesta tipo:

SELECT COUNT(*) AS total_aprobadas FROM uma-datascience-dev.Bigquery_Dataset.evaluaciones_auditoria WHERE Estado = 'aprobado' AND EXTRACT(YEAR FROM Fecha) = EXTRACT(YEAR FROM CURRENT_DATE())

---

Ejemplos de consultas típicas:
- "¿Cuántas prácticas se solicitaron este mes?" 
- "¿Qué profesionales han atendido más pacientes?"
- "¿Cuáles son los estados más comunes?"
- "¿Qué prácticas se solicitan más frecuentemente?"

La fecha de hoy es: {obtener_fecha()}


---


Consulta del usuario: {pregunta}

Genera una consulta SQL precisa y eficiente.


"""

    model = genai.GenerativeModel(
        model_name="gemini-2.0-flash-lite",
        generation_config={
            "response_mime_type": "application/json",
            "response_schema": SQLResult
        }
    )   
    response = model.generate_content(prompt)
    consulta_sql = json.loads(response.text)
    print("Consulta SQL generada:", consulta_sql)
    return consulta_sql

# Definir función para ejecutar SQL

    

import requests

def ejecutar_sql(consulta_sql: str) -> dict:
    """
    Ejecuta una consulta SQL en un endpoint de Cloud Run que expone acceso a BigQuery.

    Parámetros:
    -----------
    consulta_sql : str
        Consulta SQL a ejecutar.

    Retorna:
    --------
    dict
        Diccionario con dos claves:
        - "resultados": lista de filas devueltas por la consulta (cada fila como dict).
        - "cantidad_de_filas": cantidad total de filas en la respuesta.
    
    Lanza:
    ------
    Exception si la llamada al endpoint falla o devuelve un error HTTP.
    """
    url = "https://bigquery-tool-145741847476.us-central1.run.app"
    payload = {"sql": consulta_sql}

    response = requests.post(url, json=payload)

    if response.status_code == 200:
        data = response.json()
        resultados = data.get("results", [])
        return {
            "resultados": resultados,
            "cantidad_de_filas": len(resultados)
        }
    else:
        raise Exception(
            f"Error al llamar Cloud Run: {response.status_code} - {response.text}"
        )

APP_NAME = "predoc_app"
# Envolver funciones como herramientas de ADK
tool_generar_sql = FunctionTool(func=generar_sql)
tool_ejecutar_sql = FunctionTool(func=ejecutar_sql)

# Definir el agente
root_agent = LlmAgent(
    name="analista_datos_medicos",
    model='gemini-2.5-flash',
    instruction="""
Eres un asistente experto en análisis de datos médicos/administrativos. Tu trabajo consiste en interpretar preguntas en lenguaje natural y responderlas con datos reales almacenados en BigQuery.

Tu tabla contiene información sobre:
- Fechas de solicitudes médicas
- Prácticas médicas solicitadas
- DNI de pacientes
- Profesionales médicos
- Estados de las solicitudes

Para responder preguntas:
1. Usa la herramienta `generar_sql` para transformar la pregunta del usuario en una consulta SQL válida
2. Luego usa la herramienta `ejecutar_sql` para ejecutar esa consulta y obtener los resultados
3. Interpreta y presenta los resultados de manera clara y útil

Hazlo automáticamente sin pedir confirmación al usuario. Si encuentras errores, explica qué ocurrió y sugiere alternativas.

Ejemplos de preguntas que puedes responder:
- "¿Cuántas prácticas se solicitaron esta semana?"
- "¿Qué profesional ha atendido más casos?"
- "¿Cuáles son las prácticas más solicitadas?"
- "¿Cómo se distribuyen los estados de las solicitudes?"
- "¿Cuántos pacientes únicos hay en el sistema?"
""",
    tools=[generar_sql, ejecutar_sql],
    generate_content_config=types.GenerationConfig(
        temperature=0.0
    )
)


session_service = InMemorySessionService()

# Esta función se conecta a ChatInterface
def respond(text, history):
    # Detectar si es inicio de conversación (history vacío o con 0 mensajes)
    print("HISTORY",history)
    if history is None or len(history) == 0:
        user_id = str(uuid.uuid4())
        session_id = str(uuid.uuid4())
        print(f"🔄 Nueva sesión: {session_id}")

        async def create_session():
            await session_service.create_session(
                app_name=APP_NAME,
                user_id=user_id,
                session_id=session_id
            )
        asyncio.run(create_session())
        # Guardar en algún lado user_id y session_id para usar en las próximas llamadas
        # Por simplicidad acá lo guardamos en variables globales
        global CURRENT_USER_ID, CURRENT_SESSION_ID
        CURRENT_USER_ID = user_id
        CURRENT_SESSION_ID = session_id
    else:
        # Usar IDs existentes
        user_id = CURRENT_USER_ID
        session_id = CURRENT_SESSION_ID

    runner = Runner(agent=root_agent, app_name=APP_NAME, session_service=session_service)

    def call_agent_text(query):
        content = types.Content(role='user', parts=[types.Part(text=query)])
        events = runner.run(user_id=user_id, session_id=session_id, new_message=content)
        for event in events:
            if event.is_final_response():
                return event.content.parts[0].text
        return "No se obtuvo respuesta."



    return call_agent_text(text)


# Inicializamos demo sin el argumento state
demo = gr.ChatInterface(fn=respond, title="Agente Analista", multimodal=False)

demo.launch(debug=True)