Multi-Agent System: SQL vs. RAG Comparison (Single-Input Version)

# 1. INSTALL DEPENDENCIES
!pip install -q gradio openai pandas chromadb openpyxl langchain-community langchain-openai git+https://github.com/KarAnalytics/llm_cascade.git sentence-transformers

import os
import sqlite3
import pandas as pd
import gradio as gr
import chromadb

from llm_cascade import get_cascade

llm = get_cascade()

# --- GLOBAL STATE & DB CONFIG ---
DB_NAME = 'agent_database.db'
CHROMA_CLIENT = chromadb.Client()

# --- AGENT LOGIC ---

def agent_1_ingest(file_path):
    '''Data Ingestion & Logical Schema Identification'''
    xl = pd.ExcelFile(file_path)
    sheet_data = ''
    data_map = {}

    for sheet in xl.sheet_names:
        df = pd.read_excel(file_path, sheet_name=sheet)
        data_map[sheet] = df
        cols = ', '.join([f'{col} ({dtype})' for col, dtype in zip(df.columns, df.dtypes)])
        sheet_data += f'Sheet Name: {sheet}' + chr(10) + f'Columns detected: {cols}' + chr(10) + chr(10)

    # Prompting for Bulleted Logical Schema ONLY
    prompt = f'Analyze these sheets and provide a logical database schema. Use ONLY bullet points. Identify suggested Primary Keys and Foreign Keys:' + chr(10) + sheet_data
    system_prompt = 'You are a Data Architect. Output a logical schema in bulleted text. Do NOT write SQL code.'
    response = llm.generate(prompt, system_prompt=system_prompt)
    return response.text, data_map

def agent_2_sql_developer(logical_schema, data_map):
    '''SQL Developer Agent: Create & Populate'''
    # 1. Generate DDL
    prompt = f'Based on this logical schema, write SQLite CREATE TABLE statements:' + chr(10) + logical_schema
    system_prompt = 'Return ONLY raw SQL CREATE TABLE statements. No explanation.'
    response = llm.generate(prompt, system_prompt=system_prompt)
    ddl_resp = response.text.replace('```sql', '').replace('```', '')

    # 2. Execute DDL and Insert Data immediately
    if os.path.exists(DB_NAME): os.remove(DB_NAME)
    conn = sqlite3.connect(DB_NAME)
    cursor = conn.cursor()
    try:
        cursor.executescript(ddl_resp)
        # Match dataframes to newly created tables
        for table, df in data_map.items():
            # Basic cleanup: remove spaces from sheet names for SQL safety
            clean_table_name = table.replace(' ', '_')
            df.to_sql(clean_table_name, conn, if_exists='replace', index=False)
        conn.commit()
        status = 'Database Ready: Tables created and data inserted.'
    except Exception as e:
        status = f'Database Error: {str(e)}'
    finally:
        conn.close()

    return ddl_resp, status

def process_queries(queries_text, logical_schema, data_map):
    '''Handles multiple queries and runs Agent 2, 3, and 4 for each.'''
    queries = [q.strip() for q in queries_text.split(chr(10)) if q.strip()]
    results = []

    conn = sqlite3.connect(DB_NAME)

    for q in queries:
        # Agent 2: SQL Approach
        sql_gen_prompt = 'Schema:' + chr(10) + logical_schema + chr(10) + 'Query: ' + q + chr(10) + 'Generate ONLY the SQL SELECT statement. Table names might have underscores.'
        response = llm.generate(sql_gen_prompt)
        sql_query = response.text.replace('```sql', '').replace('```', '')

        try:
            sql_data = pd.read_sql_query(sql_query, conn).to_string()
        except Exception as e:
            sql_data = f'SQL Execution Error: {str(e)}'

        # Agent 3: RAG Approach
        context = chr(10).join(['Table ' + k + ':' + chr(10) + v.head(20).to_string() for k, v in data_map.items()])
        rag_prompt = 'Data Context:' + chr(10) + context + chr(10) + 'User Question: ' + q + chr(10) + 'Answer the question based on the data provided.'
        response = llm.generate(rag_prompt)
        rag_res = response.text

        # Agent 4: Similarity Scorer Only
        eval_prompt = 'Compare these two answers and provide ONLY a Similarity Score (0-100%).' + chr(10) + 'Answer A: ' + sql_data + chr(10) + 'Answer B: ' + rag_res
        system_prompt = 'You are a calculator. Output only a percentage score.'
        response = llm.generate(eval_prompt, system_prompt=system_prompt)
        score = response.text

        results.append('QUERY: ' + q + chr(10) + 'SQL Result: ' + sql_data + chr(10) + 'RAG Result: ' + rag_res + chr(10) + 'SIMILARITY SCORE: ' + score + chr(10) + '-' * 40)

    conn.close()
    return chr(10).join(results)

# --- WORKFLOW CONTROLLER ---
def run_app(file, queries):
    logical_schema, data_map = agent_1_ingest(file.name)
    ddl, db_status = agent_2_sql_developer(logical_schema, data_map)
    comparison_results = process_queries(queries, logical_schema, data_map)
    return logical_schema, ddl, db_status, comparison_results

# --- GRADIO UI ---
with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
    gr.Markdown("# 🤖 Agentic SQL vs. RAG Comparison")

    with gr.Row():
        file_input = gr.File(label="Upload Excel File")
        query_input = gr.Textbox(label="Business Queries (One per line)", lines=5,
                                 placeholder="How many sales in North?\nWho is the top employee?")

    run_btn = gr.Button("🚀 Process All Agents", variant="primary")

    with gr.Tabs():
        with gr.TabItem("Data Architecture (Agent 1 & 2)"):
            out_schema = gr.Markdown(label="Logical Schema (Agent 1)")
            out_ddl = gr.Code(label="Generated DDL (Agent 2)", language="sql")
            out_status = gr.Textbox(label="DB Population Status")

        with gr.TabItem("Comparative Results (Agent 2, 3 & 4)"):
            out_comparison = gr.Textbox(label="The Showdown Results", lines=20)

    run_btn.click(
        run_app,
        inputs=[file_input, query_input],
        outputs=[out_schema, out_ddl, out_status, out_comparison]
    )

demo.launch()