Spaces:
Sleeping
Sleeping
File size: 5,516 Bytes
69e9d44 d779a9b 69e9d44 d779a9b 70cfb8f 86284cb 69e9d44 d779a9b 69e9d44 d779a9b 69e9d44 d779a9b 69e9d44 d779a9b 69e9d44 d779a9b 69e9d44 | 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 | from langchain_community.utilities import SQLDatabase
from langchain.chains import create_sql_query_chain
from langchain_openai import ChatOpenAI # or your LLaMA adapter
from config.settings import settings
from langchain.prompts import PromptTemplate, ChatPromptTemplate
from app.db.schema_reader import get_schema
from app.db.db_connector import get_db_engine
from app.db.db_connector import get_connection
import re
# llm = ChatOpenAI(model="gpt-4o-mini", temperature=0, api_key=settings.OPENAI_API_KEY)
print(f"OPENAI_API_KEY: {settings.OPENAI_API_KEY}")
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0, api_key=settings.OPENAI_API_KEY)
def detect_foreign_key_relationships_only(schema):
"""Only detect foreign key relationships without any assumptions about data types"""
relationships = []
for table_name, columns in schema.items():
for column in columns:
# Look for foreign key patterns: *_id columns
if column.endswith('_id') and column != 'id':
# Extract potential referenced table name
ref_name = column[:-3] # Remove '_id'
# Check for exact match or plural/singular variations
possible_refs = [ref_name, ref_name + 's', ref_name[:-1] if ref_name.endswith('s') else ref_name + 's']
for possible_ref in possible_refs:
if possible_ref in schema:
relationships.append({
'from_table': table_name,
'from_column': column,
'to_table': possible_ref,
'to_column': 'id'
})
break
return relationships
def generate_pure_schema_context(schema):
"""Generate schema context without any assumptions - let AI figure out everything"""
# Just provide the raw schema structure
context = "DATABASE SCHEMA:\n\n"
# List all tables and their columns
context += "TABLES AND COLUMNS:\n"
for table_name, columns in schema.items():
context += f"- {table_name}: {', '.join(columns)}\n"
# Detect only foreign key relationships (no assumptions about data types)
relationships = detect_foreign_key_relationships_only(schema)
# Show detected relationships
if relationships:
context += "\nDETECTED FOREIGN KEY RELATIONSHIPS:\n"
for rel in relationships:
context += f"- {rel['from_table']}.{rel['from_column']} → {rel['to_table']}.{rel['to_column']}\n"
return context
def generate_sql_query(user_question: str):
schema = get_schema()
# Generate pure schema context - no assumptions, let AI figure it out
schema_context = generate_pure_schema_context(schema)
prompt = ChatPromptTemplate.from_template("""
You are an expert SQLite query generator with advanced analytical capabilities.
{schema_context}
INTELLIGENT ANALYSIS INSTRUCTIONS:
1. EXAMINE the schema above - look at table names and column names to understand what data each table contains
2. ANALYZE the column names to infer what type of information is stored (names, addresses, amounts, dates, etc.)
3. USE the detected foreign key relationships to understand how tables connect
4. For queries needing data from multiple tables, follow the relationship paths and use JOINs
5. INFER from context - if user asks about location and you see city/address columns, use those tables
6. THINK logically about which tables would contain the requested information
QUERY GENERATION STRATEGY:
- Read the user's question carefully
- Identify what data they want (customers, transactions, locations, etc.)
- Look at the schema to find which tables likely contain that data
- Check if multiple tables are needed and use the foreign key relationships
- Generate appropriate SQL with JOINs if data spans multiple tables
RESPONSE FORMAT:
- Return ONLY valid SQLite syntax
- No explanations, comments, or extra text
- Use proper JOIN syntax when needed
- If you cannot determine how to query the available schema, respond: "I cannot generate this query."
Example thinking process:
- User asks "customers in New York" → Look for tables with customer info AND location info → Use relationships to JOIN them
- User asks "transactions over $1000" → Look for tables with transaction/amount data
- User asks "branch managers" → Look for tables with employee/staff info and management roles
User question: {question}
""")
messages = prompt.format_messages(schema_context=schema_context, question=user_question)
sql_query = llm.invoke(messages).content.strip()
# If the LLM says it cannot generate a query, treat as normal chat
if sql_query.lower().startswith("i cannot generate this query"):
# Use the LLM as a chatbot for normal conversation
chat_prompt = (
"You are an AI assistant whose primary role is to generate SQL queries from natural language questions. "
"If the user's question is not related to SQL or databases, respond conversationally as a helpful assistant."
f"\n\nUser question: {user_question}"
)
chat_response = llm.invoke([{"role": "user", "content": chat_prompt}]).content.strip()
return {"chat_message": chat_response}
return sql_query
|