| 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}
|
|
|
