File size: 3,713 Bytes
c101c53 |
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 |
from typing import List
from haystack import component
import pandas as pd
pd.set_option('display.max_rows', None)
pd.set_option('display.max_columns', None)
pd.set_option('display.width', None)
pd.set_option('display.max_colwidth', None)
import sqlite3
import psycopg2
from utils import TEMP_DIR
@component
class SQLiteQuery:
def __init__(self, sql_database: str):
self.connection = sqlite3.connect(sql_database, check_same_thread=False)
@component.output_types(results=List[str], queries=List[str])
def run(self, queries: List[str], session_hash):
print("ATTEMPTING TO RUN SQLITE QUERY")
dir_path = TEMP_DIR / str(session_hash)
results = []
for query in queries:
result = pd.read_sql(query, self.connection)
result.to_csv(f'{dir_path}/file_upload/query.csv', index=False)
results.append(f"{result}")
self.connection.close()
return {"results": results, "queries": queries}
def sqlite_query_func(queries: List[str], session_hash, **kwargs):
dir_path = TEMP_DIR / str(session_hash)
sql_query = SQLiteQuery(f'{dir_path}/file_upload/data_source.db')
try:
result = sql_query.run(queries, session_hash)
if len(result["results"][0]) > 1000:
print("QUERY TOO LARGE")
return {"reply": "query result too large to be processed by llm, the query results are in our query.csv file. If you need to display the results directly, perhaps use the table_generation_func function."}
else:
return {"reply": result["results"][0]}
except Exception as e:
reply = f"""There was an error running the SQL Query = {queries}
The error is {e},
You should probably try again.
"""
return {"reply": reply}
@component
class PostgreSQLQuery:
def __init__(self, url: str, sql_port: int, sql_user: str, sql_pass: str, sql_db_name: str):
self.connection = psycopg2.connect(
database=sql_db_name,
user=sql_user,
password=sql_pass,
host=url, # e.g., "localhost" or an IP address
port=sql_port # default is 5432
)
@component.output_types(results=List[str], queries=List[str])
def run(self, queries: List[str], session_hash):
print("ATTEMPTING TO RUN POSTGRESQL QUERY")
dir_path = TEMP_DIR / str(session_hash)
results = []
for query in queries:
print(query)
result = pd.read_sql_query(query, self.connection)
result.to_csv(f'{dir_path}/sql/query.csv', index=False)
results.append(f"{result}")
self.connection.close()
return {"results": results, "queries": queries}
def sql_query_func(queries: List[str], session_hash, db_url, db_port, db_user, db_pass, db_name, **kwargs):
sql_query = PostgreSQLQuery(db_url, db_port, db_user, db_pass, db_name)
try:
result = sql_query.run(queries, session_hash)
print("RESULT")
print(result)
if len(result["results"][0]) > 1000:
print("QUERY TOO LARGE")
return {"reply": "query result too large to be processed by llm, the query results are in our query.csv file. If you need to display the results directly, perhaps use the table_generation_func function."}
else:
return {"reply": result["results"][0]}
except Exception as e:
reply = f"""There was an error running the SQL Query = {queries}
The error is {e},
You should probably try again.
"""
print(reply)
return {"reply": reply}
|