Multi-Agent System: SQL vs. RAG for Business Database Queries

# 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 chromadb
from llm_cascade import get_cascade

llm = get_cascade()

# Global database config
DB_NAME = 'agent_database.db'
CHROMA_CLIENT = chromadb.Client()

def agent_1_ingest(file_path):
    '''Agent 1 -- Data Architect: reads Excel, proposes logical schema.'''
    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)

    prompt = (
        '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


print('Agent 1 (Data Architect) ready.')

def agent_2_sql_developer(logical_schema, data_map):
    '''Agent 2 -- SQL Developer: generates DDL, creates DB, populates tables.'''
    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('```', '')

    # Execute DDL and insert data
    if os.path.exists(DB_NAME):
        os.remove(DB_NAME)
    conn = sqlite3.connect(DB_NAME)
    cursor = conn.cursor()
    try:
        cursor.executescript(ddl_resp)
        for table, df in data_map.items():
            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


print('Agent 2 (SQL Developer) ready.')

def agent_3_sql_query(question, logical_schema, conn):
    '''Agent 3 -- SQL Query Writer: writes and executes a SQL SELECT for the question.'''
    sql_gen_prompt = (
        'Schema:' + chr(10) + logical_schema + chr(10)
        + 'Query: ' + question + 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('```', '').strip()

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


print('Agent 3 (SQL Query Writer) ready.')

def agent_4_rag_analyst(question, data_map):
    '''Agent 4 -- RAG Analyst: answers from raw data passed as context.'''
    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: ' + question + chr(10)
        + 'Answer the question based on the data provided.'
    )
    response = llm.generate(rag_prompt)
    return response.text


print('Agent 4 (RAG Analyst) ready.')

def agent_5_evaluator(sql_result, rag_result):
    '''Agent 5 -- Evaluator: compares SQL and RAG answers, returns a similarity score.'''
    eval_prompt = (
        'Compare these two answers and provide ONLY a Similarity Score (0-100%).' + chr(10)
        + 'Answer A: ' + sql_result + chr(10)
        + 'Answer B: ' + rag_result
    )
    system_prompt = 'You are a calculator. Output only a percentage score.'
    response = llm.generate(eval_prompt, system_prompt=system_prompt)
    return response.text


print('Agent 5 (Evaluator) ready.')

import gradio as gr

# Module-level state so the Gradio UI can persist the schema across queries
_app_state = {'logical_schema': None, 'data_map': None, 'db_ready': False}


def ui_build_database(file):
    '''Phase 1: run Agents 1 and 2 on the uploaded file.'''
    if file is None:
        return '', '', 'Please upload an Excel file first.'

    try:
        logical_schema, data_map = agent_1_ingest(file.name)
        ddl, status = agent_2_sql_developer(logical_schema, data_map)

        _app_state['logical_schema'] = logical_schema
        _app_state['data_map'] = data_map
        _app_state['db_ready'] = 'Ready' in status
        return logical_schema, ddl, status
    except Exception as e:
        return '', '', f'Error building database: {e}'


def ui_ask_question(question, history):
    '''Phase 2: run Agents 3, 4, 5 on a single question.'''
    if not _app_state['db_ready']:
        bot_msg = 'Please upload an Excel file and click Build Database first.'
        return history + [[question, bot_msg]], ''

    if not question or not question.strip():
        return history, question

    try:
        conn = sqlite3.connect(DB_NAME)
        try:
            sql_query, sql_result = agent_3_sql_query(question, _app_state['logical_schema'], conn)
        finally:
            conn.close()

        rag_result = agent_4_rag_analyst(question, _app_state['data_map'])
        similarity = agent_5_evaluator(sql_result, rag_result)

        response = (
            '**Agent 3 (SQL):**' + chr(10) + '```sql' + chr(10) + sql_query + chr(10) + '```' + chr(10) + chr(10)
            + '**SQL Result:**' + chr(10) + '```' + chr(10) + sql_result + chr(10) + '```' + chr(10) + chr(10)
            + '**Agent 4 (RAG):** ' + rag_result + chr(10) + chr(10)
            + '**Agent 5 (Similarity):** ' + similarity
        )
    except Exception as e:
        response = f'Error: {e}'

    return history + [[question, response]], ''


with gr.Blocks(title='Multi-Agent DB') as demo:
    gr.Markdown('# Multi-Agent: Excel to SQLite to SQL vs. RAG Comparison')

    gr.Markdown('### Step 1: Upload Excel file and build the database (Agents 1 & 2)')
    with gr.Row():
        file_input = gr.File(label='Excel File', file_types=['.xlsx', '.xls'])
        build_btn = gr.Button('Build Database', variant='primary')

    with gr.Row():
        with gr.Column():
            out_schema = gr.Markdown(label='Logical Schema (Agent 1)')
        with gr.Column():
            out_ddl = gr.Code(label='SQL DDL (Agent 2)', language='sql')
    out_status = gr.Textbox(label='Status', interactive=False)

    build_btn.click(
        ui_build_database,
        inputs=[file_input],
        outputs=[out_schema, out_ddl, out_status],
    )

    gr.Markdown('### Step 2: Ask business questions (Agents 3, 4 & 5 run per query)')
    chatbot = gr.Chatbot(label='Q&A History', height=400)
    query_input = gr.Textbox(
        label='Your question',
        placeholder='e.g., How many total sales? Who is the top employee?',
    )
    submit_btn = gr.Button('Submit Query', variant='primary')

    submit_btn.click(
        ui_ask_question,
        inputs=[query_input, chatbot],
        outputs=[chatbot, query_input],
    )
    query_input.submit(
        ui_ask_question,
        inputs=[query_input, chatbot],
        outputs=[chatbot, query_input],
    )

demo.launch(debug=False)