File size: 3,713 Bytes
3d3efc0
 
 
 
 
 
 
 
 
 
 
0f58395
3d3efc0
 
 
 
0f58395
3d3efc0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b0f0c39
 
 
3d3efc0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
import sqlparse
import typer
from transformers import pipeline
from groq import Groq
import json
app = typer.Typer()
from dotenv import load_dotenv
import os
load_dotenv()
import streamlit as st
import gradio as gr
client = Groq(api_key=os.getenv("GROQ_API_KEY"))
# print("API Key Loaded:", os.getenv("GROQ_API_KEY"))
# st.title("πŸš€ SQL Copilot ")


#client = Groq(api_key=GROQ_API_KEY)

from datetime import datetime

# ---------------------------
# Tools (Agents)
# ---------------------------
def format_agent(query: str = typer.Option(..., "--query", "-q", help="The SQL query to format")):
    """Formatter agent πŸ“"""
    return sqlparse.format(query, reindent=True, keyword_case="upper")

def explainer_agent(query: str = typer.Option(..., "--query", "-q", help="The SQL query to format")):
    """AI-powered SQL explanation"""
    prompt = f"""
    You are an expert SQL analyst. 
    Analyze this SQL query and return:
    1. Business logic in plain English point wise in client or business person who is not technical understandable format.
    2. Tables used
    3. Joins used
    4. Filters/conditions
     
    

    SQL Query:
    {query}
    """
    completion = client.chat.completions.create(
        model="llama-3.3-70b-versatile",  
        messages=[{"role": "user", "content": prompt}],
    )
    
    return (completion.choices[0].message.content)

AGENTS = {
    "Formatter πŸ“": format_agent,
    "Explainer πŸ€–": explainer_agent
  
}
AGENT_COLORS = {
    "Formatter πŸ“": "#ADD8E6",  # light blue
    "Explainer πŸ€–": "#90EE90" # light green

}
# ---------------------------
# Process Queries
# ---------------------------
def process_queries(query_text, uploaded_file):
    queries = []

    if uploaded_file is not None:
        with open(uploaded_file.name, "r", encoding="utf-8") as f:
            content = f.read()
        #content = uploaded_file.read().decode("utf-8")
        queries = [q.strip() for q in sqlparse.split(content) if q.strip()]
    elif query_text and query_text.strip():
        queries = [query_text.strip()]

    if not queries:
        return [], None, None

    results = []
    formatted_list = []
    rules_list = []

    for q in queries:
        formatted_q = format_agent(q)
        rules_q = explainer_agent(q)

        results.append(["Formatter πŸ“", formatted_q])
        results.append(["Explainer πŸ€–", rules_q])

        formatted_list.append(formatted_q)
        rules_list.append(f"Query:\n{q}\n\nBusiness Rules:\n{rules_q}\n{'-'*50}\n")

    # Save files for download
    formatted_file = "formatted_queries.sql"
    rules_file = "business_rules.txt"

    with open(formatted_file, "w", encoding="utf-8") as f:
        f.write("\n\n".join(formatted_list))

    with open(rules_file, "w", encoding="utf-8") as f:
        f.write("\n".join(rules_list))

    return results, formatted_file, rules_file

# ---------------------------
# Gradio UI
# ---------------------------
with gr.Blocks(title="πŸš€ SQL Copilot Agent") as demo:
    gr.Markdown("## πŸš€ SQL Copilot Agent")

    with gr.Tab("Single Query / File"):
        query_box = gr.Textbox(label="Paste your SQL query here", lines=6)
        file_box = gr.File(label="Or upload SQL file", file_types=[".sql"])
        process_btn = gr.Button("Process")

        output_display = gr.Dataframe(headers=["Agent", "Output"], interactive=False)
        download_formatted = gr.File(label="Download Formatted SQL")
        download_rules = gr.File(label="Download Business Rules")

        process_btn.click(
            fn=process_queries,
            inputs=[query_box, file_box],
            outputs=[output_display, download_formatted, download_rules],
        )

demo.launch()