Spaces:
Sleeping
Sleeping
File size: 7,118 Bytes
5e16d9f |
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 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 |
import sqlite3
import logging
import requests
from typing import Dict, Any, List, Optional
from dataclasses import dataclass
@dataclass
class TableInfo:
"""Store table information including schema and relationships"""
name: str
columns: List[Dict[str, Any]]
relationships: List[Dict[str, str]]
class LLMService:
def __init__(self, api_key: str, db_path: str):
self.api_key = api_key
self.db_path = db_path
self.api_url = "https://api.groq.com/openai/v1/chat/completions"
self.headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json"
}
logging.basicConfig(level=logging.INFO)
self.logger = logging.getLogger(__name__)
self.table_info = self._load_database_schema()
def _load_database_schema(self) -> Dict[str, TableInfo]:
"""Load complete database schema with relationships"""
schema_info = {}
try:
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
# Get all tables
cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
tables = cursor.fetchall()
for table in tables:
table_name = table[0]
# Get columns
cursor.execute(f"PRAGMA table_info({table_name});")
columns = [{
'name': col[1],
'type': col[2],
'nullable': not col[3],
'primary_key': bool(col[5])
} for col in cursor.fetchall()]
# Get foreign keys
cursor.execute(f"PRAGMA foreign_key_list({table_name});")
relationships = [{
'from_table': table_name,
'to_table': fk[2],
'from_column': fk[3],
'to_column': fk[4]
} for fk in cursor.fetchall()]
schema_info[table_name] = TableInfo(
name=table_name,
columns=columns,
relationships=relationships
)
return schema_info
except sqlite3.Error as e:
self.logger.error(f"Database error while loading schema: {e}")
return {}
def _prepare_schema_prompt(self) -> str:
"""Prepare a comprehensive schema description for the LLM"""
prompt = "Database Schema:\n\n"
# Add table schemas
for table_name, info in self.table_info.items():
prompt += f"Table: {table_name}\n"
prompt += "Columns:\n"
for col in info.columns:
prompt += f"- {col['name']} ({col['type']})"
if col['primary_key']:
prompt += " PRIMARY KEY"
if not col['nullable']:
prompt += " NOT NULL"
prompt += "\n"
# Add relationships
if info.relationships:
prompt += "Relationships:\n"
for rel in info.relationships:
prompt += f"- {rel['from_table']}.{rel['from_column']} -> {rel['to_table']}.{rel['to_column']}\n"
prompt += "\n"
return prompt
def convert_to_sql_query(self, natural_query: str) -> Dict[str, Any]:
"""Convert natural language to SQL query using Groq API"""
schema_prompt = self._prepare_schema_prompt()
system_prompt = f"""
You are an expert SQL query generator. Your task is to convert natural language queries into valid SQL queries based on the provided database schema.
{schema_prompt}
Rules for generating SQL queries:
1. Use proper JOIN syntax when relating multiple tables
2. Consider table relationships and use appropriate JOIN conditions
3. Handle NULL values appropriately
4. Use table aliases when necessary for clarity
5. Return only the requested columns, use * only when specifically asked
6. Include WHERE clauses based on the natural language conditions
7. Use appropriate aggregation functions when needed (COUNT, SUM, AVG, etc.)
Generate a SQL query for the following natural language request:
{natural_query}
Return only the SQL query without any explanation.
"""
payload = {
"model": "llama3-8b-8192",
"messages": [
{"role": "system", "content": system_prompt},
{"role": "user", "content": natural_query}
],
"max_tokens": 500,
"temperature": 0.1
}
try:
self.logger.info(f"Sending request to Groq API for query: {natural_query}")
response = requests.post(
self.api_url,
headers=self.headers,
json=payload,
timeout=30
)
response.raise_for_status()
result = response.json()
self.logger.info(f"Received response from Groq API")
if 'choices' in result and result['choices']:
sql_query = self._extract_sql_query(result['choices'][0]['message']['content'])
if sql_query:
return {"success": True, "query": sql_query}
return {"success": False, "error": "Failed to generate SQL query"}
except requests.exceptions.RequestException as e:
self.logger.error(f"API request error: {str(e)}")
return {"success": False, "error": f"API request failed: {str(e)}"}
def _extract_sql_query(self, content: str) -> Optional[str]:
"""Extract SQL query from LLM response"""
# Remove markdown code blocks if present
content = content.replace("```sql", "").replace("```", "").strip()
# Basic validation
if content.upper().startswith("SELECT"):
return content
return None
def execute_query(self, sql_query: str) -> Dict[str, Any]:
"""Execute the SQL query and return results"""
try:
with sqlite3.connect(self.db_path) as conn:
cursor = conn.cursor()
cursor.execute(sql_query)
results = cursor.fetchall()
columns = [description[0] for description in cursor.description]
return {
"success": True,
"columns": columns,
"results": results
}
except sqlite3.Error as e:
self.logger.error(f"Database error: {e}")
return {
"success": False,
"error": str(e)
} |