import gradio as gr
from huggingface_hub import InferenceClient
import os
import torch
import gradio as gr
from sqlalchemy import create_engine
from langchain.sql_database import SQLDatabase
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from langchain.llms import HuggingFacePipeline
from langchain.agents import create_sql_agent
from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit

# 🔒 Obtener la conexión de la base de datos desde variables de entorno (más seguro)
DB_CONNECTION_STRING = os.getenv("DB_CONNECTION_STRING")

# Conectar a la base de datos PostgreSQL
engine = create_engine("postgresql://chatbot2025:IRjZzW5GhhuhVsjk8klAxuJ2qGYh9sQg@dpg-cu8b3han91rc738jrq0g-a.frankfurt-postgres.render.com/chatbot2025")
database = SQLDatabase(engine, view_support=True, schema="public")


# 🔥 Cargar modelo SQLCoder-7B
tokenizer = AutoTokenizer.from_pretrained("defog/sqlcoder-7b")
modelo = AutoModelForCausalLM.from_pretrained("defog/sqlcoder-7b")

# Crear pipeline de Hugging Face para generación de texto
text_pipeline = pipeline("text-generation", model=modelo, tokenizer=tokenizer, max_new_tokens=500)

# Convertir el pipeline en un modelo compatible con LangChain
llm = HuggingFacePipeline(pipeline=text_pipeline)

# Conectar la base de datos con el modelo
toolkit = SQLDatabaseToolkit(db=database, llm=llm)

# Crear el agente para consultas SQL
agent = create_sql_agent(llm=llm, toolkit=toolkit, verbose=True, agent_type="zero-shot-react-description")

# 🔹 Función para generar consultas SQL con el agente
def generate_sql(query):
    response = agent.invoke(query)
    return response

# 🔹 Crear interfaz con Gradio
iface = gr.Interface(fn=generate_sql, inputs="text", outputs="text", title="Generador de Consultas SQL")

# 🔥 Lanzar la aplicación
if __name__ == "__main__":
    iface.launch()