Upload 4 files

functions/query_functions.py

@@ -23,15 +23,14 @@ class SQLiteQuery:
       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):
+    def run(self, queries: AnyStr, 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)
-          column_names = list(result.columns)
-          results.append(f"{result}")
+        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}
     
@@ -48,16 +47,14 @@ class PostgreSQLQuery:
         )
 
     @component.output_types(results=List[str], queries=List[str])
-    def run(self, queries: List[str], session_hash):
+    def run(self, queries: AnyStr, 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)
-          column_names = list(result.columns)
-          results.append(f"{result}")
+        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}
 
@@ -145,7 +142,7 @@ class GraphQLQuery:
           results.append(f"{response_frame}") 
           return {"results": results, "queries": graphql_query, "csv_columns": column_names}
     
-def query_func(queries:List[str], session_hash, session_folder, args, **kwargs):
+def query_func(queries:AnyStr, session_hash, session_folder, args, **kwargs):
     try:
       print("QUERY")
       print(queries)

tools/tools.py

@@ -18,11 +18,8 @@ def tools_call(session_hash, data_source, titles):
                             "type": "object",
                             "properties": {
                                 "queries": {
-                                    "type": "array",
-                                    "description": "The query to use in the search. Infer this from the user's message. It should be a question or a statement",
-                                    "items": {
-                                        "type": "string",
-                                    }
+                                    "type": "string",
+                                    "description": "The query to use in the search. Infer this from the user's message. It should be a question or a statement.",
                                 }
                             },
                             "required": ["queries"],
@@ -42,11 +39,8 @@ def tools_call(session_hash, data_source, titles):
                         "type": "object",
                         "properties": {
                             "queries": {
-                                "type": "array",
-                                "description": "The PostgreSQL query to use in the search. Infer this from the user's message. It should be a question or a statement",
-                                "items": {
-                                    "type": "string",
-                                }
+                                "type": "string",
+                                "description": "The PostgreSQL query to use in the search. Infer this from the user's message. It should be a question or a statement.",
                             }
                         },
                         "required": ["queries"],