Spaces:
Sleeping
Sleeping
| from typing import List | |
| from typing import AnyStr | |
| from haystack import component | |
| import pandas as pd | |
| from pandasql import sqldf | |
| 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 pymongo import MongoClient | |
| import pymongoarrow.monkey | |
| import json | |
| import pluck | |
| from utils import TEMP_DIR | |
| import ast | |
| class SQLiteQuery: | |
| def __init__(self, sql_database: str): | |
| self.connection = sqlite3.connect(sql_database, check_same_thread=False) | |
| def run(self, queries: AnyStr, session_hash): | |
| print("ATTEMPTING TO RUN SQLITE QUERY") | |
| dir_path = TEMP_DIR / str(session_hash) | |
| results = [] | |
| result = pd.read_sql(queries, self.connection) | |
| result.to_csv(f'{dir_path}/file_upload/query.csv', index=False) | |
| column_names = list(result.columns) | |
| results.append(f"{result}") | |
| self.connection.close() | |
| return {"results": results, "queries": queries, "csv_columns": column_names} | |
| 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 | |
| ) | |
| def run(self, queries: AnyStr, session_hash): | |
| print("ATTEMPTING TO RUN POSTGRESQL QUERY") | |
| dir_path = TEMP_DIR / str(session_hash) | |
| results = [] | |
| result = pd.read_sql_query(queries, self.connection) | |
| result.to_csv(f'{dir_path}/sql/query.csv', index=False) | |
| column_names = list(result.columns) | |
| results.append(f"{result}") | |
| self.connection.close() | |
| return {"results": results, "queries": queries, "csv_columns": column_names} | |
| class DocDBQuery: | |
| def __init__(self, connection_string: str, doc_db_name: str): | |
| client = MongoClient(connection_string) | |
| self.client = client | |
| self.connection = client[doc_db_name] | |
| def run(self, aggregation_pipeline: List[str], db_collection, session_hash): | |
| pymongoarrow.monkey.patch_all() | |
| print("ATTEMPTING TO RUN MONGODB QUERY") | |
| dir_path = TEMP_DIR / str(session_hash) | |
| results = [] | |
| print(aggregation_pipeline) | |
| aggregation_pipeline = aggregation_pipeline.replace(" ", "") | |
| false_replace = [':false', ': false'] | |
| false_value = ':False' | |
| true_replace = [':true', ': true'] | |
| true_value = ':True' | |
| for replace in false_replace: | |
| aggregation_pipeline = aggregation_pipeline.replace(replace, false_value) | |
| for replace in true_replace: | |
| aggregation_pipeline = aggregation_pipeline.replace(replace, true_value) | |
| query_list = ast.literal_eval(aggregation_pipeline) | |
| print("QUERY List") | |
| print(query_list) | |
| print(db_collection) | |
| db = self.connection | |
| collection = db[db_collection] | |
| print(collection) | |
| docs = collection.aggregate_pandas_all(query_list) | |
| print("DATA FRAME COMPLETE") | |
| docs.to_csv(f'{dir_path}/doc_db/query.csv', index=False) | |
| column_names = list(docs.columns) | |
| print("CSV COMPLETE") | |
| results.append(f"{docs}") | |
| self.client.close() | |
| return {"results": results, "queries": aggregation_pipeline, "csv_columns": column_names} | |
| class GraphQLQuery: | |
| def __init__(self): | |
| self.connection = pluck | |
| def run(self, graphql_query, graphql_api_string, graphql_api_token, graphql_token_header, session_hash): | |
| print("ATTEMPTING TO RUN GRAPHQL QUERY") | |
| dir_path = TEMP_DIR / str(session_hash) | |
| results = [] | |
| headers = {"Content-Type": "application/json"} | |
| if graphql_token_header and graphql_api_token: | |
| headers[graphql_token_header] = graphql_api_token | |
| print(graphql_query) | |
| response = self.connection.execute(url=graphql_api_string, headers=headers, query=graphql_query, column_names="short") | |
| if response.errors: | |
| raise ValueError(response.errors) | |
| elif response.data: | |
| print("DATA FRAME COMPLETE") | |
| print(response) | |
| response_frame = response.frames['default'] | |
| print("RESPONSE FRAME") | |
| #print(response_frame) | |
| response_frame.to_csv(f'{dir_path}/graphql/query.csv', index=False) | |
| column_names = list(response_frame.columns) | |
| print("CSV COMPLETE") | |
| results.append(f"{response_frame}") | |
| return {"results": results, "queries": graphql_query, "csv_columns": column_names} | |
| def query_func(queries:AnyStr, session_hash, session_folder, args, **kwargs): | |
| try: | |
| print("QUERY") | |
| print(queries) | |
| if session_folder == "file_upload": | |
| dir_path = TEMP_DIR / str(session_hash) | |
| sql_query = SQLiteQuery(f'{dir_path}/file_upload/data_source.db') | |
| result = sql_query.run(queries, session_hash) | |
| elif session_folder == "sql": | |
| sql_query = PostgreSQLQuery(args[0], args[1], args[2], args[3], args[4]) | |
| result = sql_query.run(queries, session_hash) | |
| elif session_folder == 'doc_db': | |
| doc_db_query = DocDBQuery(args[0], args[1]) | |
| result = doc_db_query.run(queries, kwargs['db_collection'], session_hash) | |
| elif session_folder == 'graphql': | |
| graphql_object = GraphQLQuery() | |
| result = graphql_object.run(queries, args[0], args[1], args[2], session_hash) | |
| print("RESULT") | |
| print(result["csv_columns"]) | |
| if len(result["results"][0]) > 1000: | |
| print("QUERY TOO LARGE") | |
| return {"reply": f"""query result too large to be processed by llm, the query results are in our query.csv file. | |
| The column names of this query.csv file are: {result["csv_columns"]}. | |
| 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 {session_folder} Query = {queries} | |
| The error is {e}, | |
| You should probably try again. | |
| """ | |
| print(reply) | |
| return {"reply": reply} | |
| def graphql_schema_query(graphql_type: AnyStr, session_hash, **kwargs): | |
| dir_path = TEMP_DIR / str(session_hash) | |
| try: | |
| with open(f'{dir_path}/graphql/schema.json', 'r') as file: | |
| data = json.load(file) | |
| types_list = data["types"] | |
| result = list(filter(lambda item: item["name"] == graphql_type, types_list)) | |
| print("SCHEMA RESULT") | |
| print(graphql_type) | |
| print(str(result)) | |
| return {"reply": str(result)} | |
| except Exception as e: | |
| reply = f"""There was an error querying our schema.json file with the type:{graphql_type} | |
| The error is {e}, | |
| You should probably try again. | |
| """ | |
| print(reply) | |
| return {"reply": reply} | |
| def graphql_csv_query(csv_query: AnyStr, session_hash, **kwargs): | |
| dir_path = TEMP_DIR / str(session_hash) | |
| try: | |
| query = pd.read_csv(f'{dir_path}/graphql/query.csv') | |
| query.Name = 'query' | |
| print("GRAPHQL CSV QUERY") | |
| print(csv_query) | |
| queried_df = sqldf(csv_query, locals()) | |
| print(queried_df) | |
| column_names = list(queried_df.columns) | |
| queried_df.to_csv(f'{dir_path}/graphql/query.csv', index=False) | |
| if len(queried_df) > 1000: | |
| print("CSV QUERY TOO LARGE") | |
| return {"reply": f"""The new query results are in our query.csv file. | |
| The column names of this query.csv file are: {column_names}. | |
| If you need to display the results directly, perhaps use the table_generation_func function."""} | |
| else: | |
| return {"reply": str(queried_df)} | |
| except Exception as e: | |
| reply = f"""There was an error querying our query.csv file with the query:{csv_query} | |
| The error is {e}, | |
| You should probably try again. | |
| """ | |
| print(reply) | |
| return {"reply": reply} |