Files changed (2) hide show
  1. functions/query_functions.py +11 -14
  2. tools/tools.py +4 -10
functions/query_functions.py CHANGED
@@ -23,15 +23,14 @@ class SQLiteQuery:
23
  self.connection = sqlite3.connect(sql_database, check_same_thread=False)
24
 
25
  @component.output_types(results=List[str], queries=List[str])
26
- def run(self, queries: List[str], session_hash):
27
  print("ATTEMPTING TO RUN SQLITE QUERY")
28
  dir_path = TEMP_DIR / str(session_hash)
29
  results = []
30
- for query in queries:
31
- result = pd.read_sql(query, self.connection)
32
- result.to_csv(f'{dir_path}/file_upload/query.csv', index=False)
33
- column_names = list(result.columns)
34
- results.append(f"{result}")
35
  self.connection.close()
36
  return {"results": results, "queries": queries, "csv_columns": column_names}
37
 
@@ -48,16 +47,14 @@ class PostgreSQLQuery:
48
  )
49
 
50
  @component.output_types(results=List[str], queries=List[str])
51
- def run(self, queries: List[str], session_hash):
52
  print("ATTEMPTING TO RUN POSTGRESQL QUERY")
53
  dir_path = TEMP_DIR / str(session_hash)
54
  results = []
55
- for query in queries:
56
- print(query)
57
- result = pd.read_sql_query(query, self.connection)
58
- result.to_csv(f'{dir_path}/sql/query.csv', index=False)
59
- column_names = list(result.columns)
60
- results.append(f"{result}")
61
  self.connection.close()
62
  return {"results": results, "queries": queries, "csv_columns": column_names}
63
 
@@ -145,7 +142,7 @@ class GraphQLQuery:
145
  results.append(f"{response_frame}")
146
  return {"results": results, "queries": graphql_query, "csv_columns": column_names}
147
 
148
- def query_func(queries:List[str], session_hash, session_folder, args, **kwargs):
149
  try:
150
  print("QUERY")
151
  print(queries)
 
23
  self.connection = sqlite3.connect(sql_database, check_same_thread=False)
24
 
25
  @component.output_types(results=List[str], queries=List[str])
26
+ def run(self, queries: AnyStr, session_hash):
27
  print("ATTEMPTING TO RUN SQLITE QUERY")
28
  dir_path = TEMP_DIR / str(session_hash)
29
  results = []
30
+ result = pd.read_sql(queries, self.connection)
31
+ result.to_csv(f'{dir_path}/file_upload/query.csv', index=False)
32
+ column_names = list(result.columns)
33
+ results.append(f"{result}")
 
34
  self.connection.close()
35
  return {"results": results, "queries": queries, "csv_columns": column_names}
36
 
 
47
  )
48
 
49
  @component.output_types(results=List[str], queries=List[str])
50
+ def run(self, queries: AnyStr, session_hash):
51
  print("ATTEMPTING TO RUN POSTGRESQL QUERY")
52
  dir_path = TEMP_DIR / str(session_hash)
53
  results = []
54
+ result = pd.read_sql_query(queries, self.connection)
55
+ result.to_csv(f'{dir_path}/sql/query.csv', index=False)
56
+ column_names = list(result.columns)
57
+ results.append(f"{result}")
 
 
58
  self.connection.close()
59
  return {"results": results, "queries": queries, "csv_columns": column_names}
60
 
 
142
  results.append(f"{response_frame}")
143
  return {"results": results, "queries": graphql_query, "csv_columns": column_names}
144
 
145
+ def query_func(queries:AnyStr, session_hash, session_folder, args, **kwargs):
146
  try:
147
  print("QUERY")
148
  print(queries)
tools/tools.py CHANGED
@@ -18,11 +18,8 @@ def tools_call(session_hash, data_source, titles):
18
  "type": "object",
19
  "properties": {
20
  "queries": {
21
- "type": "array",
22
- "description": "The query to use in the search. Infer this from the user's message. It should be a question or a statement",
23
- "items": {
24
- "type": "string",
25
- }
26
  }
27
  },
28
  "required": ["queries"],
@@ -42,11 +39,8 @@ def tools_call(session_hash, data_source, titles):
42
  "type": "object",
43
  "properties": {
44
  "queries": {
45
- "type": "array",
46
- "description": "The PostgreSQL query to use in the search. Infer this from the user's message. It should be a question or a statement",
47
- "items": {
48
- "type": "string",
49
- }
50
  }
51
  },
52
  "required": ["queries"],
 
18
  "type": "object",
19
  "properties": {
20
  "queries": {
21
+ "type": "string",
22
+ "description": "The query to use in the search. Infer this from the user's message. It should be a question or a statement.",
 
 
 
23
  }
24
  },
25
  "required": ["queries"],
 
39
  "type": "object",
40
  "properties": {
41
  "queries": {
42
+ "type": "string",
43
+ "description": "The PostgreSQL query to use in the search. Infer this from the user's message. It should be a question or a statement.",
 
 
 
44
  }
45
  },
46
  "required": ["queries"],