File size: 2,222 Bytes
0276b0d | 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 | import re
import sqlite3 # Added import for sqlite3
from datetime import datetime
from db.database import Database
from nlp.sql_model import SQLModel
class QueryProcessor:
def __init__(self, db):
self.db = db
self.sql_model = SQLModel()
def convert_date(self, date_str):
formats = [
'%d/%m/%Y',
'%d %B %Y',
'%B %Y',
'%b %Y',
'%Y-%m-%d'
]
for fmt in formats:
try:
return datetime.strptime(date_str, fmt).date()
except ValueError:
continue
raise ValueError('Invalid date format')
def extract_and_convert_date(self, user_query):
date_pattern = r'\b(\b[0-9]{1,2}[/-]?[0-9]{1,2}[/-]?[0-9]{2,4}\b|\b(?:\b(?:January|February|March|April|May|June|July|August|September|October|November|December|Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\b[- ]?[0-9]{1,2}[- ]?[0-9]{2,4}\b|\b[0-9]{2,4}[- ]?(?:January|February|March|April|May|June|July|August|September|October|November|December|Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\b)\b)'
match = re.search(date_pattern, user_query, re.I)
if match:
date_str = match.group(0)
try:
return self.convert_date(date_str)
except ValueError:
return None
return None
def process_query(self, user_query):
conn = self.db.connect()
if not conn:
return "Failed to connect to the database."
response = ''
try:
# Generate SQL query using the LLM model
sql_query = self.sql_model.generate_sql(user_query)
# Execute generated SQL query against the database
cursor = self.db.execute_query(sql_query)
if cursor:
results = cursor.fetchall()
if results:
response = f'Results: {results}'
else:
response = 'No results found.'
else:
response = 'Error executing query.'
finally:
self.db.close()
return response |