File size: 1,100 Bytes
24371db 2336a25 24371db fb65c41 24371db fb65c41 24371db |
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 |
from typing import List
from haystack import component
import pandas as pd
import sqlite3
@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]):
print("ATTEMPTING TO RUN QUERY")
results = []
for query in queries:
result = pd.read_sql(query, self.connection)
results.append(f"{result}")
self.connection.close()
return {"results": results, "queries": queries}
def sqlite_query_func(queries: List[str], session_hash):
sql_query = SQLiteQuery(f'data_source_{session_hash}.db')
try:
result = sql_query.run(queries)
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}
|