improve query functions

functions/query_functions.py

@@ -81,8 +81,8 @@ class PostgreSQLQuery:
     
 
 
-def sql_query_func(queries: List[str], session_hash, args, **kwargs):
-    sql_query = PostgreSQLQuery(args[0], args[1], args[2], args[3], args[4])
+def sql_query_func(queries: List[str], session_hash, db_url, db_port, db_user, db_pass, db_name, **kwargs):
+    sql_query = PostgreSQLQuery(db_url, db_port, db_user, db_pass, db_name)
     try:
       result = sql_query.run(queries, session_hash)
       print("RESULT")
@@ -150,8 +150,8 @@ class DocDBQuery:
     
 
 
-def doc_db_query_func(aggregation_pipeline: List[str], db_collection: AnyStr, session_hash, args, **kwargs):
-    doc_db_query = DocDBQuery(args[0], args[1])
+def doc_db_query_func(aggregation_pipeline: List[str], db_collection: AnyStr, session_hash, connection_string, doc_db_name, **kwargs):
+    doc_db_query = DocDBQuery(connection_string, doc_db_name)
     try:
       result = doc_db_query.run(aggregation_pipeline, db_collection, session_hash)
       print("RESULT")
@@ -206,10 +206,10 @@ class GraphQLQuery:
     
 
 
-def graphql_query_func(graphql_query: AnyStr, session_hash, args, **kwargs):
+def graphql_query_func(graphql_query: AnyStr, session_hash, graphql_api_string, graphql_api_token, graphql_token_header, **kwargs):
     graphql_object = GraphQLQuery()
     try:
-      result = graphql_object.run(graphql_query, args[0], args[1], args[2], session_hash)
+      result = graphql_object.run(graphql_query, graphql_api_string, graphql_api_token, graphql_token_header, session_hash)
       print("RESULT")
       if len(result["results"][0]) > 1000:
         print("QUERY TOO LARGE")

templates/sql_db.py

@@ -1,6 +1,6 @@
 import ast
 import gradio as gr
-from functions import example_question_generator, chatbot_func
+from functions import sql_example_question_generator, sql_chatbot_with_fc
 from data_sources import connect_sql_db
 from utils import message_dict
 
@@ -55,7 +55,7 @@ with gr.Blocks() as demo:
                                         ]
                 else:
                     try:
-                        generated_examples = ast.literal_eval(example_question_generator(request.session_hash, 'sql', sql_db_name, process_message[2], ""))
+                        generated_examples = ast.literal_eval(sql_example_question_generator(request.session_hash, process_message[2], sql_db_name))
                         example_questions = [
                                                 ["Describe the dataset"]
                                             ]
@@ -75,18 +75,16 @@ with gr.Blocks() as demo:
                 db_user = gr.Textbox(visible=False, value=sql_user)
                 db_pass = gr.Textbox(visible=False, value=sql_pass)
                 db_name = gr.Textbox(visible=False, value=sql_db_name)
-                titles = gr.Textbox(value=process_message[2], interactive=False, label="SQL Tables")
-                data_source = gr.Textbox(visible=False, value='sql')
-                schema = gr.Textbox(visible=False, value='')
+                db_tables = gr.Textbox(value=process_message[2], interactive=False, label="SQL Tables")
                 bot = gr.Chatbot(type='messages', label="SQL DB Chat Window", render_markdown=True, sanitize_html=False, show_label=True, render=False, visible=True, elem_classes="chatbot")
                 chat = gr.ChatInterface(
-                                    fn=chatbot_func,
+                                    fn=sql_chatbot_with_fc,
                                     type='messages',
                                     chatbot=bot,
                                     title="Chat with your Database",
                                     examples=example_questions,
                                     concurrency_limit=None,
-                                    additional_inputs=[session_hash, data_source, titles, schema, db_url, db_port, db_user, db_pass, db_name]
+                                    additional_inputs=[session_hash, db_url, db_port, db_user, db_pass, db_name, db_tables]
                                     )
             
     def process_sql_db(url, sql_user, sql_port, sql_pass, sql_db_name, session_hash):