File size: 2,229 Bytes
e0560d9 | 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 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | import os
import re
from dotenv import load_dotenv
from langchain_groq import ChatGroq
from langchain_core.messages import HumanMessage, SystemMessage
from utils.query_engine import run_query
load_dotenv()
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
llm = ChatGroq(
groq_api_key=GROQ_API_KEY,
model_name="llama-3.3-70b-versatile",
temperature=0
)
SYSTEM_PROMPT = """
You are an expert SQL assistant.
IMPORTANT:
1. Generate ONLY SQLite SQL queries.
2. Do not explain anything.
3. Use valid SQLite syntax.
4. Return only executable SQL.
DATABASE TABLES:
customers(
customer_id,
name,
email,
city,
signup_date
)
products(
product_id,
product_name,
category,
price,
stock
)
employees(
employee_id,
employee_name,
department
)
orders(
order_id,
customer_id,
employee_id,
order_date,
total_amount
)
order_items(
order_item_id,
order_id,
product_id,
quantity
)
"""
def clean_sql(query):
query = query.replace("```sql", "")
query = query.replace("```", "")
return query.strip()
def generate_sql(question):
messages = [
SystemMessage(content=SYSTEM_PROMPT),
HumanMessage(content=question)
]
response = llm.invoke(messages)
sql_query = clean_sql(response.content)
return sql_query
def generate_summary(question, dataframe):
summary_prompt = f"""
User Question:
{question}
Query Result:
{dataframe.head(10).to_string()}
Generate a short business summary.
"""
response = llm.invoke(summary_prompt)
return response.content
def ask_agent(question):
try:
# Generate SQL
sql_query = generate_sql(question)
# Execute Query
result_df = run_query(sql_query)
# Generate Summary
summary = generate_summary(question, result_df)
return {
"sql": sql_query,
"data": result_df,
"summary": summary
}
except Exception as e:
return {
"error": str(e)
} |