src/streamlit_app.py

import streamlit as st
import sqlite3
import urllib.request
import os
import sqlite3
import urllib.request
from pathlib import Path
from sqlalchemy.exc import OperationalError
from sqlalchemy import create_engine, text
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain.prompts import PromptTemplate
from langchain_groq import ChatGroq
from langchain.agents import create_react_agent, AgentExecutor
from langchain import hub
import tempfile
from peft import LoraConfig, get_peft_model
import torch
from langchain_community.agent_toolkits import create_sql_agent
from langchain_groq import ChatGroq
from langchain.agents.agent_types import AgentType

from langchain_community.agent_toolkits import SQLDatabaseToolkit


st.set_page_config(page_title="Asistente para base de datos", page_icon="🧠", layout="wide")


if "db_loaded" not in st.session_state:
    st.session_state.db_loaded = False
if "messages" not in st.session_state:
    st.session_state.messages = []
if "agent" not in st.session_state:
    st.session_state.agent = None

def load_database(db_path):
    try:
        engine = create_engine(f'sqlite:///{db_path}')
        db = SQLDatabase(engine)
        

        tables = db.get_usable_table_names()
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        table_counts = {}
        for table in tables[:5]: 
            try:
                cursor.execute(f"SELECT COUNT(*) FROM {table}")
                count = cursor.fetchone()[0]
                table_counts[table] = count
            except sqlite3.Error:
                table_counts[table] = "Error"
        conn.close()
        
        return db, tables, table_counts
    except Exception as e:
        st.error(f"Error al cargar la base de datos: {str(e)}")
        return None, [], {}
    
def create_system_message(db, custom_instructions=""):
    """Crea mensaje del sistema adaptado a la base de datos del usuario"""

    tables = db.get_usable_table_names()
    table_info = []
    
    for table in tables[:4]:  # tablas para el prompt
        try:
            table_schema = db.get_table_info([table])
            table_info.append(f"Tabla {table}:\n{table_schema}")
        except:
            table_info.append(f"Tabla {table}: (esquema no disponible)")
    
    schema_info = "\n\n".join(table_info)
    
    base_prompt = hub.pull('langchain-ai/sql-agent-system-prompt')
    
    enhanced_prompt = f"""
{base_prompt.messages[0].prompt.template}

INFORMACION ESPECIFICA DE LA BASE DE DATOS:
Dialecto: SQLite
Tablas disponibles: {', '.join(tables)}

ESQUEMAS DE TABLAS:
{schema_info}

INSTRUCCIONES ADICIONALES:
- Siempre usa LIMIT para consultas exploratorias
- Verifica la existencia de columnas antes de usarlas
- Usa nombres de tablas exactos (case-sensitive)
- Para consultas complejas, dividelas en pasos
{custom_instructions}
"""
    
    return enhanced_prompt

def build_sql_agent(db, custom_instructions=""):
    groq_api_key = os.getenv("GROQ_API_KEY")
    if not groq_api_key:
        groq_api_key = st.secrets.get("GROQ_API_KEY", None)
    if not groq_api_key:
        st.error("No se encontró la clave GROQ_API_KEY. Configúrala en tu entorno o en st.secrets.")
        return None
    llm = ChatGroq(
        temperature=st.session_state.temperature,
        model_name="deepseek-r1-distill-llama-70b",
        groq_api_key=groq_api_key
    )
    system_message = create_system_message(db, custom_instructions)
    prompt_template = hub.pull('langchain-ai/react-agent-template')
    prompt = prompt_template.partial(instructions=system_message)
    return create_sql_agent(
        llm=llm,
        db=db,
        prompt=prompt,
    )


st.title("🧠 Consultas de base de datos")
st.caption("Carga una base de datos SQLite y haz preguntas en lenguaje natural")


with st.sidebar:
    st.header("📂 Cargar Base de Datos")
    option = st.radio("Seleccione fuente:", 
                      ["Subir archivo", "Ingresar URL", "Base de ejemplo (Chinook)"])
    
    db_path = None
    
    if option == "Subir archivo":
        uploaded_file = st.file_uploader("Seleccione archivo SQLite", type=["sqlite", "db", "sqlite3"])
        if uploaded_file:
            with tempfile.NamedTemporaryFile(delete=False, suffix='.sqlite') as tmpfile:
                tmpfile.write(uploaded_file.getvalue())
                db_path = tmpfile.name
    
    elif option == "Ingresar URL":
        url = st.text_input("URL de la base de datos SQLite:")
        if url:
            if st.button("Descargar base de datos"):
                with st.spinner("Descargando..."):
                    try:
                        with tempfile.NamedTemporaryFile(delete=False, suffix='.sqlite') as tmpfile:
                            urllib.request.urlretrieve(url, tmpfile.name)
                            db_path = tmpfile.name
                    except Exception as e:
                        st.error(f"Error en descarga: {str(e)}")
    
    elif option == "Base de ejemplo (Chinook)":
        if st.button("Cargar base de ejemplo"):
            url = "https://github.com/lerocha/chinook-database/raw/master/ChinookDatabase/DataSources/Chinook_Sqlite.sqlite"
            with st.spinner("Descargando base de ejemplo..."):
                try:
                    with tempfile.NamedTemporaryFile(delete=False, suffix='.sqlite') as tmpfile:
                        urllib.request.urlretrieve(url, tmpfile.name)
                        db_path = tmpfile.name
                except Exception as e:
                    st.error(f"Error en descarga: {str(e)}")


    with st.sidebar:
        
        st.header(" ")
    
        #Slider para la temperatura
        if 'temperature' not in st.session_state:
            st.session_state.temperature = 0.7  
    
    
        temperature_value = st.slider(
            label="Temperatura",
            min_value=0.0,
            max_value=1.0,  
            value=st.session_state.temperature,
            step=0.01,
            help="Controla la aleatoriedad de la respuesta. Valores mas bajos son mas deterministas, valores mas altos son más creativos."
        )
        st.session_state.temperature = temperature_value
        st.info(f"Temperatura actual: **{st.session_state.temperature:.2f}**")

     
    if db_path and os.path.exists(db_path):
        with st.spinner("Cargando base de datos..."):
            db, tables, table_counts = load_database(db_path)
            
            if db:
                st.session_state.db = db
                st.session_state.db_path = db_path
                st.session_state.db_loaded = True
                
                with st.spinner("Creando agente SQL..."):
                    st.session_state.agent = build_sql_agent(db)
                
                st.success("¡Base de datos cargada exitosamente!")
                
                st.subheader("📊 Metadatos de la base de datos")
                st.write(f"**Tablas:** {', '.join(tables)}")

                st.write("**Conteo de registros:**")
                for table, count in table_counts.items():
                    st.write(f"- {table}: {count} registros")

#Chat principal
if st.session_state.db_loaded:
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    if prompt := st.chat_input("Haz una pregunta sobre la base de datos..."):

        st.session_state.messages.append({"role": "user", "content": prompt})

        with st.chat_message("user"):
            st.markdown(prompt)

        with st.chat_message("assistant"):
            with st.spinner("Pensando..."):
                try:
                    response = st.session_state.agent.invoke({"input": prompt})
                    answer = response["output"]
                except Exception as e:
                    answer = f"⚠️ Error: {str(e)}"
                
                st.markdown(answer)

        st.session_state.messages.append({"role": "assistant", "content": answer})
else:
    st.info("Por favor carga una base de datos SQLite desde el panel lateral")

st.divider()
st.caption("SQL Assistant with Groq")