app.py

#@title Gradio UI - 1
import gradio as gr
import os
from langchain.chat_models import init_chat_model
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain.agents import create_agent
import markdown as md
import base64

# from google.colab import userdata

# --- GLOBALS ---
model = None  # ✅ prevent "NameError" before setup
db = None
toolkit = None
agent = None

# Define available models for each provider
PROVIDER_MODELS = {
    "google_genai": ["gemini-2.0-pro", "gemini-2.5-flash-lite", "gemini-2.5-flash", "gemini-2.0-flash-thinking"],
    "openai": ["gpt-4o", "gpt-4o-mini", "gpt-3.5-turbo"],
    "anthropic": ["claude-3-opus", "claude-3-sonnet", "claude-3-haiku"],
    "azure_openai": ["gpt-4-turbo", "gpt-4o-mini", "gpt-35-turbo"],
    "bedrock": ["anthropic.claude-3-sonnet-v1", "mistral.mixtral-8x7b"],
    "xai": ["grok-2", "grok-2-mini"],
    "deepseek": ["deepseek-chat", "deepseek-coder"],
    "perplexity": ["sonar-small-chat", "sonar-medium-chat", "sonar-large-chat"]
}

def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

# Encode the images
github_logo_encoded = encode_image("Images/github-logo.png")
linkedin_logo_encoded = encode_image("Images/linkedin-logo.png")
website_logo_encoded = encode_image("Images/ai-logo.png")

def create_chatbot_interface():
    with gr.Blocks(theme=gr.themes.Ocean(font=[gr.themes.GoogleFont("Noto Sans")]),
                   css='footer {visibility: hidden}') as demo:
        gr.Markdown("# DB Assistant 📊👨🏻‍🏭")
        with gr.Accordion("📔 Description:", open=False):
            gr.Markdown(md.description)
        with gr.Accordion('🚀 Key Features', open=False):
            gr.Markdown(md.key_features)
        with gr.Accordion('⚙️ Tech Stack Overview', open=False):
            gr.Markdown(md.tech_stack_overview)
        with gr.Accordion('🧩 How It Works', open=False):
            gr.Markdown(md.how_it_works)
            
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("## LLM Setup")
                llm_provider = gr.Dropdown(
                    list(PROVIDER_MODELS.keys()), label="Select Provider", value="google_genai"
                )
                llm_model = gr.Dropdown(
                    choices=PROVIDER_MODELS["google_genai"],
                    label="Select Model",
                    value="gemini-2.5-flash-lite"
                )
                api_key = gr.Textbox(label="Enter API Key", type="password")
                setup_llm_btn = gr.Button("Setup LLM")

                gr.Markdown("## Database Connection")
                db_connection_string = gr.Textbox(label="Enter Connection String", type="password")
                connect_db_btn = gr.Button("Connect to Database")
                db_status = gr.Markdown("") 

            with gr.Column(scale=2):
                gr.Markdown("## Chatbot Interface")
                chatbot = gr.Chatbot()
                msg = gr.Textbox(label="Enter your question")
                clear = gr.Button("Clear")

        with gr.Accordion("🧠 Example Interaction", open=False):
            gr.Markdown(md.example)
        with gr.Accordion('🧑‍💻 Ideal Use Cases', open=False):
            gr.Markdown(md.use_cases)
        with gr.Accordion('🪄 Future Enhancements', open=False):
            gr.Markdown(md.enhancements)
        with gr.Accordion('💡 Credits', open=False):
            gr.Markdown(md.credits)

        gr.HTML(md.footer.format(github_logo_encoded, linkedin_logo_encoded, website_logo_encoded))


        # --- FUNCTIONS ---

        def update_model_dropdown(provider):
            models = PROVIDER_MODELS.get(provider, [])
            default = models[0] if models else None
            return gr.update(choices=models, value=default)

        def setup_llm(provider, model_name, key):
            os.environ["GOOGLE_API_KEY"] = key
            global model
            model = init_chat_model(model_name, model_provider=provider)
            return f"✅ LLM model `{model_name}` from `{provider}` setup successfully."

        def connect_to_db(db_url):
            global db, toolkit, tools, agent, system_prompt, model
            if model is None:
                return "❌ Please set up the LLM before connecting to the database."

            db = SQLDatabase.from_uri(db_url)
            toolkit = SQLDatabaseToolkit(db=db, llm=model)
            tools = toolkit.get_tools()

            system_prompt = """
              You are an agent designed to interact with a SQL database.
              Given an input question, create a syntactically correct {dialect} query to run,
              then look at the results of the query and return the answer. Unless the user
              specifies a specific number of examples they wish to obtain, always limit your
              query to at most {top_k} results.

              You can order the results by a relevant column to return the most interesting
              examples in the database. Never query for all the columns from a specific table,
              only ask for the relevant columns given the question.

              You MUST double check your query before executing it. If you get an error while
              executing a query, rewrite the query and try again.

              DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the
              database.

              To start you should ALWAYS look at the tables in the database to see what you
              can query. Do NOT skip this step.

              Then you should query the schema of the most relevant tables.
              """.format(
                  dialect=db.dialect,
                  top_k=5,
              )

            agent = create_agent(model, tools, system_prompt=system_prompt)
            tables = db.get_usable_table_names()
            return f"✅ Connected to database successfully.\n\n**Available tables:** {tables}"

        def respond(message, chat_history):
            result = ""
            for step in agent.stream(
                {"messages": [{"role": "user", "content": message}]},
                stream_mode="values",
            ):
                result += step["messages"][-1].content
            chat_history.append((message, result))
            return "", chat_history

        # --- EVENTS ---
        llm_provider.change(update_model_dropdown, inputs=llm_provider, outputs=llm_model)
        setup_llm_btn.click(setup_llm, inputs=[llm_provider, llm_model, api_key], outputs=None)
        connect_db_btn.click(connect_to_db, inputs=[db_connection_string], outputs=db_status)
        msg.submit(respond, [msg, chatbot], [msg, chatbot])
        clear.click(lambda: None, None, chatbot, queue=False)

    return demo


if __name__ == "__main__":
    demo = create_chatbot_interface()
    demo.launch(debug=True)