refactor chat functions (#39)

- refactor chat functions (a1e003fcf2656d4ade9f58021902a8575d321220)

data_sources/upload_file.py

@@ -84,15 +84,18 @@ def process_data_upload(data_file, session_hash):
         os.makedirs(dir_path, exist_ok=True)
 
         connection = sqlite3.connect(f'{dir_path}/data_source.db')
-        print("Opened database successfully");
-        print(df.columns)
+        print("Opened database successfully")
 
         df.to_sql('data_source', connection, if_exists='replace', index = False)
+
+        cur=connection.execute('select * from data_source')
+        columns = [i[0] for i in cur.description]
+        print(columns)
         
         connection.commit()
         connection.close()
 
-        return ["success","<p style='color:green;text-align:center;font-size:18px;'>Data upload successful</p>"]
+        return ["success","<p style='color:green;text-align:center;font-size:18px;'>Data upload successful</p>", columns]
     except Exception as e:
         print("UPLOAD ERROR")
         print(e)

functions/__init__.py

@@ -1,9 +1,9 @@
 from .query_functions import SQLiteQuery, sqlite_query_func, sql_query_func, doc_db_query_func, graphql_query_func, graphql_schema_query, graphql_csv_query
 from .chart_functions import table_generation_func, scatter_chart_generation_func, \
 line_chart_generation_func, bar_chart_generation_func, pie_chart_generation_func, histogram_generation_func, scatter_chart_fig
-from .chat_functions import sql_example_question_generator, example_question_generator, doc_db_example_question_generator, chatbot_with_fc, sql_chatbot_with_fc, doc_db_chatbot_with_fc, graphql_chatbot_with_fc, graphql_example_question_generator
+from .chat_functions import example_question_generator, chatbot_func
 from .stat_functions import regression_func
 
 __all__ = ["SQLiteQuery","sqlite_query_func","sql_query_func","doc_db_query_func","graphql_query_func","graphql_schema_query","graphql_csv_query","table_generation_func","scatter_chart_generation_func",
            "line_chart_generation_func","bar_chart_generation_func","regression_func", "pie_chart_generation_func", "histogram_generation_func",
-           "scatter_chart_fig","doc_db_example_question_generator","sql_example_question_generator","example_question_generator","chatbot_with_fc","sql_chatbot_with_fc","doc_db_chatbot_with_fc","graphql_chatbot_with_fc","graphql_example_question_generator"]
+           "scatter_chart_fig","example_question_generator","chatbot_func"]

functions/chat_functions.py

@@ -1,340 +1,161 @@
-from utils import TEMP_DIR, message_dict
+from utils import message_dict
 
 from haystack.dataclasses import ChatMessage
 from haystack.components.generators.chat import OpenAIChatGenerator
 
-import sqlite3
-
 chat_generator = OpenAIChatGenerator(model="gpt-4o")
 response = None
 
-def example_question_generator(session_hash):
-    example_response = None
-    example_messages = [
-        ChatMessage.from_system(
-            "You are a helpful and knowledgeable agent who has access to an SQLite database which has a table called 'data_source'."
-        )
-    ]
-
-    session_path = 'file_upload'
-
-    dir_path = TEMP_DIR / str(session_hash) / str(session_path)
-    connection = sqlite3.connect(f'{dir_path}/data_source.db')
-    print("Querying questions");
-    cur=connection.execute('select * from data_source')
-    columns = [i[0] for i in cur.description]
-    print("QUESTION COLUMNS")
-    print(columns)
-    cur.close()
-    connection.close()
-
-    example_messages.append(ChatMessage.from_user(text=f"""We have a SQLite database with the following {columns}. 
-                                                  We also have an AI agent with access to the same database that will be performing data analysis.
-                                                  Please return an array of seven strings, each one being a question for our data analysis agent
-                                                  that we can suggest that you believe will be insightful or helpful to a data analysis looking for
-                                                  data insights. Return nothing more than the array of questions because I need that specific data structure
-                                                  to process your response. No other response type or data structure will work."""))
-
-    example_response = chat_generator.run(messages=example_messages)
-
-    return example_response["replies"][0].text
-
-def sql_example_question_generator(session_hash, db_tables, db_name):
-    example_response = None
-    example_messages = [
-        ChatMessage.from_system(
-            f"You are a helpful and knowledgeable agent who has access to an PostgreSQL database called {db_name}."
-        )
-    ]
-
-    example_messages.append(ChatMessage.from_user(text=f"""We have a PostgreSQL database with the following tables: {db_tables}. 
-                                                  We also have an AI agent with access to the same database that will be performing data analysis.
-                                                  Please return an array of seven strings, each one being a question for our data analysis agent
-                                                  that we can suggest that you believe will be insightful or helpful to a data analysis looking for
-                                                  data insights. Return nothing more than the array of questions because I need that specific data structure
-                                                  to process your response. No other response type or data structure will work."""))
-
-    example_response = chat_generator.run(messages=example_messages)
-
-    return example_response["replies"][0].text
-
-def doc_db_example_question_generator(session_hash, db_collections, db_name, db_schema):
-    example_response = None
-    example_messages = [
-        ChatMessage.from_system(
-            f"You are a helpful and knowledgeable agent who has access to an MongoDB NoSQL document database called {db_name}."
-        )
-    ]
-
-    example_messages.append(ChatMessage.from_user(text=f"""We have a MongoDB NoSQL document database with the following collections: {db_collections}.
-                                                  The schema of these collections is: {db_schema}. 
-                                                  We also have an AI agent with access to the same database that will be performing data analysis.
-                                                  Please return an array of seven strings, each one being a question for our data analysis agent
-                                                  that we can suggest that you believe will be insightful or helpful to a data analysis looking for
-                                                  data insights. Return nothing more than the array of questions because I need that specific data structure
-                                                  to process your response. No other response type or data structure will work."""))
-
-    example_response = chat_generator.run(messages=example_messages)
-
-    return example_response["replies"][0].text
-
-def graphql_example_question_generator(session_hash, graphql_endpoint, graphql_types):
+def example_question_message(data_source, name, titles, schema):
+
+    example_message_dict = {
+        'file_upload' : ["You are a helpful and knowledgeable agent who has access to an SQLite database which has a table called 'data_source'.",
+                         f"""We have a SQLite database with the following {titles}. 
+                        We also have an AI agent with access to the same database that will be performing data analysis.
+                        Please return an array of seven strings, each one being a question for our data analysis agent
+                        that we can suggest that you believe will be insightful or helpful to a data analysis looking for
+                        data insights. Return nothing more than the array of questions because I need that specific data structure
+                        to process your response. No other response type or data structure will work."""],
+
+        'sql' : [f"You are a helpful and knowledgeable agent who has access to an MongoDB NoSQL document database called {name}.",
+                 f"""We have a PostgreSQL database with the following tables: {titles}. 
+                        We also have an AI agent with access to the same database that will be performing data analysis.
+                        Please return an array of seven strings, each one being a question for our data analysis agent
+                        that we can suggest that you believe will be insightful or helpful to a data analysis looking for
+                        data insights. Return nothing more than the array of questions because I need that specific data structure
+                        to process your response. No other response type or data structure will work."""],
+
+        'doc_db' : [f"You are a helpful and knowledgeable agent who has access to an MongoDB NoSQL document database called {name}.",
+                    f"""We have a MongoDB NoSQL document database with the following collections: {titles}.
+                        The schema of these collections is: {schema}. 
+                        We also have an AI agent with access to the same database that will be performing data analysis.
+                        Please return an array of seven strings, each one being a question for our data analysis agent
+                        that we can suggest that you believe will be insightful or helpful to a data analysis looking for
+                        data insights. Return nothing more than the array of questions because I need that specific data structure
+                        to process your response. No other response type or data structure will work."""],
+
+        'graphql' : [f"You are a helpful and knowledgeable agent who has access to an GraphQL API endpoint called {name}.",
+                     f"""We have a GraphQL API endpoint with the following types: {titles}.
+                        We also have an AI agent with access to the same GraphQL API endpoint that will be performing data analysis.
+                        Please return an array of seven strings, each one being a question for our data analysis agent
+                        that we can suggest that you believe will be insightful or helpful to a data analysis looking for
+                        data insights. Return nothing more than the array of questions because I need that specific data structure
+                        to process your response. No other response type or data structure will work."""]         
+
+    }
+
+    return example_message_dict[data_source]
+
+def example_question_generator(session_hash, data_source, name, titles, schema):
     example_response = None
+    example_message_list = example_question_message(data_source, name, titles, schema)
     example_messages = [
         ChatMessage.from_system(
-            f"You are a helpful and knowledgeable agent who has access to an GraphQL API endpoint called {graphql_endpoint}."
+            example_message_list[0]
         )
     ]
 
-    example_messages.append(ChatMessage.from_user(text=f"""We have a GraphQL API endpoint with the following types: {graphql_types}.
-                                                  We also have an AI agent with access to the same GraphQL API endpoint that will be performing data analysis.
-                                                  Please return an array of seven strings, each one being a question for our data analysis agent
-                                                  that we can suggest that you believe will be insightful or helpful to a data analysis looking for
-                                                  data insights. Return nothing more than the array of questions because I need that specific data structure
-                                                  to process your response. No other response type or data structure will work."""))
+    example_messages.append(ChatMessage.from_user(text=example_message_list[1]))
 
     example_response = chat_generator.run(messages=example_messages)
 
     return example_response["replies"][0].text
 
-def chatbot_with_fc(message, history, session_hash):
+def system_message(data_source, titles, schema=""):
+
+    system_message_dict = {
+        'file_upload' : f"""You are a helpful and knowledgeable agent who has access to an SQLite database which has a table called 'data_source' that contains the following columns: {titles}. 
+                    You also have access to a function, called table_generation_func, that can take a query.csv file generated from our sql query and returns an iframe that we should display in our chat window.
+                    You also have access to a scatter plot function, called scatter_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a scatter plot and returns an iframe that we should display in our chat window.
+                    You also have access to a line chart function, called line_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a line chart and returns an iframe that we should display in our chat window.
+                    You also have access to a bar graph function, called line_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a bar graph and returns an iframe that we should display in our chat window.
+                    You also have access to a pie chart function, called pie_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a pie chart and returns an iframe that we should display in our chat window.
+                    You also have access to a histogram function, called histogram_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a histogram and returns an iframe that we should display in our chat window.
+                    You also have access to a linear regression function, called regression_func, that can take a query.csv file generated from our sql query and a list of column names for our independent and dependent variables and return a regression data string and a regression chart which is returned as an iframe.
+                    Could you please always display the generated charts, tables, and visualizations as part of your output?""",
+
+        'sql' : f"""You are a helpful and knowledgeable agent who has access to an PostgreSQL database which has a series of tables called {titles}. 
+                    You also have access to a function, called table_generation_func, that can take a query.csv file generated from our sql query and returns an iframe that we should display in our chat window.
+                    You also have access to a scatter plot function, called scatter_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a scatter plot and returns an iframe that we should display in our chat window.
+                    You also have access to a line chart function, called line_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a line chart and returns an iframe that we should display in our chat window.
+                    You also have access to a bar graph function, called line_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a bar graph and returns an iframe that we should display in our chat window.
+                    You also have access to a pie chart function, called pie_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a pie chart and returns an iframe that we should display in our chat window.
+                    You also have access to a histogram function, called histogram_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a histogram and returns an iframe that we should display in our chat window.
+                    You also have access to a linear regression function, called regression_func, that can take a query.csv file generated from our sql query and a list of column names for our independent and dependent variables and return a regression data string and a regression chart which is returned as an iframe.
+                    Could you please always display the generated charts, tables, and visualizations as part of your output?""",
+
+        'doc_db' : f"""You are a helpful and knowledgeable agent who has access to a NoSQL MongoDB Document database which has a series of collections called {titles}.
+                    The schema of these collections is: {schema}.
+                    You also have access to a function, called table_generation_func, that can take a query.csv file generated from our MongoDB query and returns an iframe that we should display in our chat window.
+                    You also have access to a scatter plot function, called scatter_chart_generation_func, that can take a query.csv file generated from our MongoDB query and uses plotly dictionaries to generate a scatter plot and returns an iframe that we should display in our chat window.
+                    You also have access to a line chart function, called line_chart_generation_func, that can take a query.csv file generated from our MongoDB query and uses plotly dictionaries to generate a line chart and returns an iframe that we should display in our chat window.
+                    You also have access to a bar graph function, called line_chart_generation_func, that can take a query.csv file generated from our MongoDB query and uses plotly dictionaries to generate a bar graph and returns an iframe that we should display in our chat window.
+                    You also have access to a pie chart function, called pie_chart_generation_func, that can take a query.csv file generated from our MongoDB query and uses plotly dictionaries to generate a pie chart and returns an iframe that we should display in our chat window.
+                    You also have access to a histogram function, called histogram_generation_func, that can take a query.csv file generated from our MongoDB query and uses plotly dictionaries to generate a histogram and returns an iframe that we should display in our chat window.
+                    You also have access to a linear regression function, called regression_func, that can take a query.csv file generated from our MongoDB query and a list of column names for our independent and dependent variables and return a regression data string and a regression chart which is returned as an iframe.
+                    Could you please always display the generated charts, tables, and visualizations as part of your output?""",
+
+        'graphql' : f"""You are a helpful and knowledgeable agent who has access to a GraphQL API which has the following types: {titles}.
+                    We have also saved a schema.json file that contains the entire introspection query that we can use to find out more about each type before making a query.
+                    You also have access to a function, called table_generation_func, that can take a query.csv file generated from our GraphQL API query and returns an iframe that we should display in our chat window.
+                    You also have access to a scatter plot function, called scatter_chart_generation_func, that can take a query.csv file generated from our GraphQL API query and uses plotly dictionaries to generate a scatter plot and returns an iframe that we should display in our chat window.
+                    You also have access to a line chart function, called line_chart_generation_func, that can take a query.csv file generated from our GraphQL API query and uses plotly dictionaries to generate a line chart and returns an iframe that we should display in our chat window.
+                    You also have access to a bar graph function, called line_chart_generation_func, that can take a query.csv file generated from our GraphQL API query and uses plotly dictionaries to generate a bar graph and returns an iframe that we should display in our chat window.
+                    You also have access to a pie chart function, called pie_chart_generation_func, that can take a query.csv file generated from our GraphQL API query and uses plotly dictionaries to generate a pie chart and returns an iframe that we should display in our chat window.
+                    You also have access to a histogram function, called histogram_generation_func, that can take a query.csv file generated from our GraphQL API query and uses plotly dictionaries to generate a histogram and returns an iframe that we should display in our chat window.
+                    You also have access to a linear regression function, called regression_func, that can take a query.csv file generated from our GraphQL API query and a list of column names for our independent and dependent variables and return a regression data string and a regression chart which is returned as an iframe.
+                    Could you please always display the generated charts, tables, and visualizations as part of your output?"""            
+
+    }
+
+    return system_message_dict[data_source]
+
+def chatbot_func(message, history, session_hash, data_source, titles, schema, *args):
     from functions import sqlite_query_func, table_generation_func, regression_func, scatter_chart_generation_func, \
+        sql_query_func, doc_db_query_func, graphql_query_func, graphql_schema_query, graphql_csv_query, \
         line_chart_generation_func,bar_chart_generation_func,pie_chart_generation_func,histogram_generation_func
     import tools.tools as tools
 
-    available_functions = {"sql_query_func": sqlite_query_func,"table_generation_func":table_generation_func,
+    available_functions = {"sqlite_query_func": sqlite_query_func,"sql_query_func": sql_query_func,"doc_db_query_func": doc_db_query_func,
+                           "graphql_query_func": graphql_query_func,"graphql_schema_query": graphql_schema_query,"graphql_csv_query": graphql_csv_query,
+                           "table_generation_func":table_generation_func,
                            "line_chart_generation_func":line_chart_generation_func,"bar_chart_generation_func":bar_chart_generation_func,
                            "scatter_chart_generation_func":scatter_chart_generation_func, "pie_chart_generation_func":pie_chart_generation_func,
                            "histogram_generation_func":histogram_generation_func,
                            "regression_func":regression_func }
     
-    session_path = 'file_upload'
-
-    dir_path = TEMP_DIR / str(session_hash) / str(session_path)
-    connection = sqlite3.connect(f'{dir_path}/data_source.db')
-    cur=connection.execute('select * from data_source')
-    columns = [i[0] for i in cur.description]
-    cur.close()
-    connection.close()
-
-    if message_dict[session_hash]['file_upload'] != None:
-        message_dict[session_hash]['file_upload'].append(ChatMessage.from_user(message))
-    else:
-        messages = [
-            ChatMessage.from_system(
-                f"""You are a helpful and knowledgeable agent who has access to an SQLite database which has a table called 'data_source' that contains the following columns: {columns}. 
-                You also have access to a function, called table_generation_func, that can take a query.csv file generated from our sql query and returns an iframe that we should display in our chat window.
-                You also have access to a scatter plot function, called scatter_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a scatter plot and returns an iframe that we should display in our chat window.
-                You also have access to a line chart function, called line_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a line chart and returns an iframe that we should display in our chat window.
-                You also have access to a bar graph function, called line_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a bar graph and returns an iframe that we should display in our chat window.
-                You also have access to a pie chart function, called pie_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a pie chart and returns an iframe that we should display in our chat window.
-                You also have access to a histogram function, called histogram_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a histogram and returns an iframe that we should display in our chat window.
-                You also have access to a linear regression function, called regression_func, that can take a query.csv file generated from our sql query and a list of column names for our independent and dependent variables and return a regression data string and a regression chart which is returned as an iframe.
-                Could you please always display the generated charts, tables, and visualizations as part of your output?"""
-            )
-        ]
-        messages.append(ChatMessage.from_user(message))
-        message_dict[session_hash]['file_upload'] = messages
-
-    response = chat_generator.run(messages=message_dict[session_hash]['file_upload'], generation_kwargs={"tools": tools.data_file_tools_call(session_hash)})
-
-    while True:
-        # if OpenAI response is a tool call
-        if response and response["replies"][0].meta["finish_reason"] == "tool_calls" or response["replies"][0].tool_calls:
-            function_calls = response["replies"][0].tool_calls
-            for function_call in function_calls:
-                message_dict[session_hash]['file_upload'].append(ChatMessage.from_assistant(tool_calls=[function_call]))
-                ## Parse function calling information
-                function_name = function_call.tool_name
-                function_args = function_call.arguments
-
-                ## Find the corresponding function and call it with the given arguments
-                function_to_call = available_functions[function_name]
-                function_response = function_to_call(**function_args, session_hash=session_hash, session_folder='file_upload')
-                print(function_name)
-                ## Append function response to the messages list using `ChatMessage.from_tool`
-                message_dict[session_hash]['file_upload'].append(ChatMessage.from_tool(tool_result=function_response['reply'], origin=function_call))
-                response = chat_generator.run(messages=message_dict[session_hash]['file_upload'], generation_kwargs={"tools": tools.data_file_tools_call(session_hash)})
-
-        # Regular Conversation
-        else:
-            message_dict[session_hash]['file_upload'].append(response["replies"][0])
-            break
-
-    return response["replies"][0].text
-
-def sql_chatbot_with_fc(message, history, session_hash, db_url, db_port, db_user, db_pass, db_name, db_tables):
-    from functions import sql_query_func, table_generation_func, regression_func, scatter_chart_generation_func, \
-        line_chart_generation_func,bar_chart_generation_func,pie_chart_generation_func,histogram_generation_func
-    import tools.tools as tools
-
-    available_functions = {"sql_query_func": sql_query_func,"table_generation_func":table_generation_func,
-                           "line_chart_generation_func":line_chart_generation_func,"bar_chart_generation_func":bar_chart_generation_func,
-                           "scatter_chart_generation_func":scatter_chart_generation_func, "pie_chart_generation_func":pie_chart_generation_func,
-                           "histogram_generation_func":histogram_generation_func,
-                           "regression_func":regression_func }
-
-    if message_dict[session_hash]['sql'] != None:
-        message_dict[session_hash]['sql'].append(ChatMessage.from_user(message))
-    else:
-        messages = [
-            ChatMessage.from_system(
-                f"""You are a helpful and knowledgeable agent who has access to an PostgreSQL database which has a series of tables called {db_tables}. 
-                You also have access to a function, called table_generation_func, that can take a query.csv file generated from our sql query and returns an iframe that we should display in our chat window.
-                You also have access to a scatter plot function, called scatter_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a scatter plot and returns an iframe that we should display in our chat window.
-                You also have access to a line chart function, called line_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a line chart and returns an iframe that we should display in our chat window.
-                You also have access to a bar graph function, called line_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a bar graph and returns an iframe that we should display in our chat window.
-                You also have access to a pie chart function, called pie_chart_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a pie chart and returns an iframe that we should display in our chat window.
-                You also have access to a histogram function, called histogram_generation_func, that can take a query.csv file generated from our sql query and uses plotly dictionaries to generate a histogram and returns an iframe that we should display in our chat window.
-                You also have access to a linear regression function, called regression_func, that can take a query.csv file generated from our sql query and a list of column names for our independent and dependent variables and return a regression data string and a regression chart which is returned as an iframe.
-                Could you please always display the generated charts, tables, and visualizations as part of your output?"""
-            )
-        ]
-        messages.append(ChatMessage.from_user(message))
-        message_dict[session_hash]['sql'] = messages
-
-    response = chat_generator.run(messages=message_dict[session_hash]['sql'], generation_kwargs={"tools": tools.sql_tools_call(db_tables)})
-
-    while True:
-        # if OpenAI response is a tool call
-        if response and response["replies"][0].meta["finish_reason"] == "tool_calls" or response["replies"][0].tool_calls:
-            function_calls = response["replies"][0].tool_calls
-            for function_call in function_calls:
-                message_dict[session_hash]['sql'].append(ChatMessage.from_assistant(tool_calls=[function_call]))
-                ## Parse function calling information
-                function_name = function_call.tool_name
-                function_args = function_call.arguments
-
-                ## Find the corresponding function and call it with the given arguments
-                function_to_call = available_functions[function_name]
-                function_response = function_to_call(**function_args, session_hash=session_hash, db_url=db_url, 
-                                                     db_port=db_port, db_user=db_user, db_pass=db_pass, db_name=db_name, session_folder='sql')
-                print(function_name)
-                ## Append function response to the messages list using `ChatMessage.from_tool`
-                message_dict[session_hash]['sql'].append(ChatMessage.from_tool(tool_result=function_response['reply'], origin=function_call))
-                response = chat_generator.run(messages=message_dict[session_hash]['sql'], generation_kwargs={"tools": tools.sql_tools_call(db_tables)})
-
-        # Regular Conversation
-        else:
-            message_dict[session_hash]['sql'].append(response["replies"][0])
-            break
-        
-    return response["replies"][0].text
-
-def doc_db_chatbot_with_fc(message, history, session_hash, db_connection_string, db_name, db_collections, db_schema):
-    from functions import doc_db_query_func, table_generation_func, regression_func, scatter_chart_generation_func, \
-        line_chart_generation_func,bar_chart_generation_func,pie_chart_generation_func,histogram_generation_func
-    import tools.tools as tools
-
-    available_functions = {"doc_db_query_func": doc_db_query_func,"table_generation_func":table_generation_func,
-                           "line_chart_generation_func":line_chart_generation_func,"bar_chart_generation_func":bar_chart_generation_func,
-                           "scatter_chart_generation_func":scatter_chart_generation_func, "pie_chart_generation_func":pie_chart_generation_func,
-                           "histogram_generation_func":histogram_generation_func,
-                           "regression_func":regression_func }
-
-    if message_dict[session_hash]['doc_db'] != None:
-        message_dict[session_hash]['doc_db'].append(ChatMessage.from_user(message))
+    if message_dict[session_hash][data_source] != None:
+        message_dict[session_hash][data_source].append(ChatMessage.from_user(message))
     else:
         messages = [
-            ChatMessage.from_system(
-                f"""You are a helpful and knowledgeable agent who has access to a NoSQL MongoDB Document database which has a series of collections called {db_collections}.
-                The schema of these collections is: {db_schema}.
-                You also have access to a function, called table_generation_func, that can take a query.csv file generated from our MongoDB query and returns an iframe that we should display in our chat window.
-                You also have access to a scatter plot function, called scatter_chart_generation_func, that can take a query.csv file generated from our MongoDB query and uses plotly dictionaries to generate a scatter plot and returns an iframe that we should display in our chat window.
-                You also have access to a line chart function, called line_chart_generation_func, that can take a query.csv file generated from our MongoDB query and uses plotly dictionaries to generate a line chart and returns an iframe that we should display in our chat window.
-                You also have access to a bar graph function, called line_chart_generation_func, that can take a query.csv file generated from our MongoDB query and uses plotly dictionaries to generate a bar graph and returns an iframe that we should display in our chat window.
-                You also have access to a pie chart function, called pie_chart_generation_func, that can take a query.csv file generated from our MongoDB query and uses plotly dictionaries to generate a pie chart and returns an iframe that we should display in our chat window.
-                You also have access to a histogram function, called histogram_generation_func, that can take a query.csv file generated from our MongoDB query and uses plotly dictionaries to generate a histogram and returns an iframe that we should display in our chat window.
-                You also have access to a linear regression function, called regression_func, that can take a query.csv file generated from our MongoDB query and a list of column names for our independent and dependent variables and return a regression data string and a regression chart which is returned as an iframe.
-                Could you please always display the generated charts, tables, and visualizations as part of your output?"""
-            )
+            ChatMessage.from_system(system_message(data_source, titles, schema))
         ]
         messages.append(ChatMessage.from_user(message))
-        message_dict[session_hash]['doc_db'] = messages
-
-    response = chat_generator.run(messages=message_dict[session_hash]['doc_db'], generation_kwargs={"tools": tools.doc_db_tools_call(db_collections)})
-
-    while True:
-        # if OpenAI response is a tool call
-        if response and response["replies"][0].meta["finish_reason"] == "tool_calls" or response["replies"][0].tool_calls:
-            function_calls = response["replies"][0].tool_calls
-            for function_call in function_calls:
-                message_dict[session_hash]['doc_db'].append(ChatMessage.from_assistant(tool_calls=[function_call]))
-                ## Parse function calling information
-                function_name = function_call.tool_name
-                function_args = function_call.arguments
-
-                ## Find the corresponding function and call it with the given arguments
-                function_to_call = available_functions[function_name]
-                function_response = function_to_call(**function_args, session_hash=session_hash, connection_string=db_connection_string, 
-                                                    doc_db_name=db_name, session_folder='doc_db')
-                print(function_name)
-                ## Append function response to the messages list using `ChatMessage.from_tool`
-                message_dict[session_hash]['doc_db'].append(ChatMessage.from_tool(tool_result=function_response['reply'], origin=function_call))
-                response = chat_generator.run(messages=message_dict[session_hash]['doc_db'], generation_kwargs={"tools": tools.doc_db_tools_call(db_collections)})
-
-        # Regular Conversation
-        else:
-            message_dict[session_hash]['doc_db'].append(response["replies"][0])
-            break
-        
-    return response["replies"][0].text
-
-def graphql_chatbot_with_fc(message, history, session_hash, graphql_api_string, graphql_api_token, graphql_token_header, graphql_types):
-    from functions import graphql_query_func, graphql_schema_query, graphql_csv_query, table_generation_func, regression_func, scatter_chart_generation_func, \
-        line_chart_generation_func,bar_chart_generation_func,pie_chart_generation_func,histogram_generation_func
-    import tools.tools as tools
+        message_dict[session_hash][data_source] = messages
     
-    available_functions = {"graphql_query_func": graphql_query_func,"graphql_schema_query": graphql_schema_query,"graphql_csv_query": graphql_csv_query,"table_generation_func":table_generation_func,
-                           "line_chart_generation_func":line_chart_generation_func,"bar_chart_generation_func":bar_chart_generation_func,
-                           "scatter_chart_generation_func":scatter_chart_generation_func, "pie_chart_generation_func":pie_chart_generation_func,
-                           "histogram_generation_func":histogram_generation_func,
-                           "regression_func":regression_func }
-
-    if message_dict[session_hash]['graphql'] != None:
-        message_dict[session_hash]['graphql'].append(ChatMessage.from_user(message))
-    else:
-        messages = [
-            ChatMessage.from_system(
-                f"""You are a helpful and knowledgeable agent who has access to a GraphQL API which has the following types: {graphql_types}.
-                We have also saved a schema.json file that contains the entire introspection query that we can use to find out more about each type before making a query.
-                You also have access to a function, called table_generation_func, that can take a query.csv file generated from our GraphQL API query and returns an iframe that we should display in our chat window.
-                You also have access to a scatter plot function, called scatter_chart_generation_func, that can take a query.csv file generated from our GraphQL API query and uses plotly dictionaries to generate a scatter plot and returns an iframe that we should display in our chat window.
-                You also have access to a line chart function, called line_chart_generation_func, that can take a query.csv file generated from our GraphQL API query and uses plotly dictionaries to generate a line chart and returns an iframe that we should display in our chat window.
-                You also have access to a bar graph function, called line_chart_generation_func, that can take a query.csv file generated from our GraphQL API query and uses plotly dictionaries to generate a bar graph and returns an iframe that we should display in our chat window.
-                You also have access to a pie chart function, called pie_chart_generation_func, that can take a query.csv file generated from our GraphQL API query and uses plotly dictionaries to generate a pie chart and returns an iframe that we should display in our chat window.
-                You also have access to a histogram function, called histogram_generation_func, that can take a query.csv file generated from our GraphQL API query and uses plotly dictionaries to generate a histogram and returns an iframe that we should display in our chat window.
-                You also have access to a linear regression function, called regression_func, that can take a query.csv file generated from our GraphQL API query and a list of column names for our independent and dependent variables and return a regression data string and a regression chart which is returned as an iframe.
-                Could you please always display the generated charts, tables, and visualizations as part of your output?"""
-            )
-        ]
-        messages.append(ChatMessage.from_user(message))
-        message_dict[session_hash]['graphql'] = messages
-
-    response = chat_generator.run(messages=message_dict[session_hash]['graphql'], generation_kwargs={"tools": tools.graphql_tools_call(graphql_types)})
+    response = chat_generator.run(messages=message_dict[session_hash][data_source], generation_kwargs={"tools": tools.tools_call(session_hash, data_source, titles)})
 
     while True:
         # if OpenAI response is a tool call
         if response and response["replies"][0].meta["finish_reason"] == "tool_calls" or response["replies"][0].tool_calls:
             function_calls = response["replies"][0].tool_calls
             for function_call in function_calls:
-                message_dict[session_hash]['graphql'].append(ChatMessage.from_assistant(tool_calls=[function_call]))
+                message_dict[session_hash][data_source].append(ChatMessage.from_assistant(tool_calls=[function_call]))
                 ## Parse function calling information
                 function_name = function_call.tool_name
                 function_args = function_call.arguments
 
                 ## Find the corresponding function and call it with the given arguments
                 function_to_call = available_functions[function_name]
-                function_response = function_to_call(**function_args, session_hash=session_hash, graphql_api_string=graphql_api_string, 
-                                                    graphql_api_token=graphql_api_token, graphql_token_header=graphql_token_header, session_folder='graphql')
+                function_response = function_to_call(**function_args, session_hash=session_hash, session_folder=data_source, args=args)
                 print(function_name)
                 ## Append function response to the messages list using `ChatMessage.from_tool`
-                message_dict[session_hash]['graphql'].append(ChatMessage.from_tool(tool_result=function_response['reply'], origin=function_call))
-                response = chat_generator.run(messages=message_dict[session_hash]['graphql'], generation_kwargs={"tools": tools.graphql_tools_call(graphql_types)})
+                message_dict[session_hash][data_source].append(ChatMessage.from_tool(tool_result=function_response['reply'], origin=function_call))
+                response = chat_generator.run(messages=message_dict[session_hash][data_source], generation_kwargs={"tools": tools.tools_call(session_hash, data_source, titles)})
 
         # Regular Conversation
         else:
-            message_dict[session_hash]['graphql'].append(response["replies"][0])
+            message_dict[session_hash][data_source].append(response["replies"][0])
             break
-        
+
     return response["replies"][0].text

functions/query_functions.py

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

templates/data_file.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from functions import example_question_generator, chatbot_with_fc
+from functions import example_question_generator, chatbot_func
 from data_sources import process_data_upload
 from utils import message_dict
 import ast
@@ -97,7 +97,7 @@ with gr.Blocks() as demo:
                                         ]
                 else:
                     try:
-                        generated_examples = ast.literal_eval(example_question_generator(request.session_hash))
+                        generated_examples = ast.literal_eval(example_question_generator(request.session_hash, 'file_upload', '', process_message[1], ''))
                         example_questions = [
                                                 ["Describe the dataset"]
                                             ]
@@ -111,16 +111,19 @@ with gr.Blocks() as demo:
                                             ["List the columns in the dataset"],
                                             ["What could this data be used for?"],
                                         ]
-                parameters = gr.Textbox(visible=False, value=request.session_hash)
+                session_hash = gr.Textbox(visible=False, value=request.session_hash)
+                data_source = gr.Textbox(visible=False, value='file_upload')
+                schema = gr.Textbox(visible=False, value='')
+                titles = gr.Textbox(value=process_message[1], interactive=False, visible=False)
                 bot = gr.Chatbot(type='messages', label="CSV Chat Window", render_markdown=True, sanitize_html=False, show_label=True, render=False, visible=True, elem_classes="chatbot")
                 chat = gr.ChatInterface(
-                                    fn=chatbot_with_fc,
+                                    fn=chatbot_func,
                                     type='messages',
                                     chatbot=bot,
                                     title="Chat with your data file",
                                     concurrency_limit=None,
                                     examples=example_questions,
-                                    additional_inputs=parameters
+                                    additional_inputs=[session_hash, data_source, titles, schema]
                                     )
                 
     def process_upload(upload_value, session_hash):

templates/doc_db.py

@@ -1,6 +1,6 @@
 import ast
 import gradio as gr
-from functions import doc_db_example_question_generator, doc_db_chatbot_with_fc
+from functions import example_question_generator, chatbot_func
 from data_sources import connect_doc_db
 from utils import message_dict
 
@@ -59,7 +59,7 @@ with gr.Blocks() as demo:
                                         ]
                 else:
                     try:
-                        generated_examples = ast.literal_eval(doc_db_example_question_generator(request.session_hash, process_message[2], doc_db_name, process_message[3]))
+                        generated_examples = ast.literal_eval(example_question_generator(request.session_hash, 'graphql', doc_db_name, process_message[2], process_message[3]))
                         example_questions = [
                                                 ["Describe the dataset"]
                                             ]
@@ -76,17 +76,18 @@ with gr.Blocks() as demo:
                 session_hash = gr.Textbox(visible=False, value=request.session_hash)
                 db_connection_string = gr.Textbox(visible=False, value=connection_login_value)
                 db_name = gr.Textbox(visible=False, value=doc_db_name)
-                db_collections = gr.Textbox(value=process_message[2], interactive=False, label="DB Collections")
-                db_schema = gr.Textbox(visible=False, value=process_message[3])
+                titles = gr.Textbox(value=process_message[2], interactive=False, label="DB Collections")
+                data_source = gr.Textbox(visible=False, value='doc_db')
+                schema = gr.Textbox(visible=False, value=process_message[3])
                 bot = gr.Chatbot(type='messages', label="DocDB Chat Window", render_markdown=True, sanitize_html=False, show_label=True, render=False, visible=True, elem_classes="chatbot")
                 chat = gr.ChatInterface(
-                                    fn=doc_db_chatbot_with_fc,
+                                    fn=chatbot_func,
                                     type='messages',
                                     chatbot=bot,
                                     title="Chat with your Database",
                                     examples=example_questions,
                                     concurrency_limit=None,
-                                    additional_inputs=[session_hash, db_connection_string, db_name, db_collections,db_schema]
+                                    additional_inputs=[session_hash, data_source, titles, schema, db_connection_string, db_name]
                                     )
             
     def process_doc_db(connection_string, nosql_db_name, session_hash):

templates/graphql.py

@@ -1,6 +1,6 @@
 import ast
 import gradio as gr
-from functions import graphql_example_question_generator, graphql_chatbot_with_fc
+from functions import example_question_generator, chatbot_func
 from data_sources import connect_graphql
 from utils import message_dict
 
@@ -69,7 +69,7 @@ with gr.Blocks() as demo:
                                         ]
                 else:
                     try:
-                        generated_examples = ast.literal_eval(graphql_example_question_generator(request.session_hash, graphql_url, process_message[2]))
+                        generated_examples = ast.literal_eval(example_question_generator(request.session_hash, 'graphql', graphql_url, process_message[2], ''))
                         example_questions = [
                                                 ["Describe the dataset"]
                                             ]
@@ -87,16 +87,18 @@ with gr.Blocks() as demo:
                 graphql_api_string = gr.Textbox(visible=False, value=graphql_url)
                 graphql_api_token = gr.Textbox(visible=False, value=api_token)
                 graphql_token_header = gr.Textbox(visible=False, value=api_token_header_name)
-                graphql_types = gr.Textbox(value=process_message[2], interactive=False, label="GraphQL Types")
+                titles = gr.Textbox(value=process_message[2], interactive=False, label="GraphQL Types")
+                data_source = gr.Textbox(visible=False, value='graphql')
+                schema = gr.Textbox(visible=False, value='')
                 bot = gr.Chatbot(type='messages', label="GraphQL Chat Window", render_markdown=True, sanitize_html=False, show_label=True, render=False, visible=True, elem_classes="chatbot")
                 chat = gr.ChatInterface(
-                                    fn=graphql_chatbot_with_fc,
+                                    fn=chatbot_func,
                                     type='messages',
                                     chatbot=bot,
                                     title="Chat with your Graphql API",
                                     examples=example_questions,
                                     concurrency_limit=None,
-                                    additional_inputs=[session_hash, graphql_api_string, graphql_api_token, graphql_token_header, graphql_types]
+                                    additional_inputs=[session_hash, data_source, titles, schema, graphql_api_string, graphql_api_token, graphql_token_header]
                                     )
             
     def process_graphql(graphql_url, api_token, api_token_header_name, session_hash):

templates/sql_db.py

@@ -1,6 +1,6 @@
 import ast
 import gradio as gr
-from functions import sql_example_question_generator, sql_chatbot_with_fc
+from functions import example_question_generator, chatbot_func
 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(sql_example_question_generator(request.session_hash, process_message[2], sql_db_name))
+                        generated_examples = ast.literal_eval(example_question_generator(request.session_hash, 'sql', sql_db_name, process_message[2], ""))
                         example_questions = [
                                                 ["Describe the dataset"]
                                             ]
@@ -75,16 +75,18 @@ 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)
-                db_tables = gr.Textbox(value=process_message[2], interactive=False, label="SQL Tables")
+                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='')
                 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=sql_chatbot_with_fc,
+                                    fn=chatbot_func,
                                     type='messages',
                                     chatbot=bot,
                                     title="Chat with your Database",
                                     examples=example_questions,
                                     concurrency_limit=None,
-                                    additional_inputs=[session_hash, db_url, db_port, db_user, db_pass, db_name, db_tables]
+                                    additional_inputs=[session_hash, data_source, titles, schema, db_url, db_port, db_user, db_pass, db_name]
                                     )
             
     def process_sql_db(url, sql_user, sql_port, sql_pass, sql_db_name, session_hash):

tools/tools.py

@@ -1,187 +1,149 @@
-import sqlite3
-import psycopg2
 from .stats_tools import stats_tools
 from .chart_tools import chart_tools
-from utils import TEMP_DIR
 
-def data_file_tools_call(session_hash):
-    dir_path = TEMP_DIR / str(session_hash)
-    connection = sqlite3.connect(f'{dir_path}/file_upload/data_source.db')
-    print("Querying Database in Tools.py");
-    cur=connection.execute('select * from data_source')
-    columns = [i[0] for i in cur.description]
-    print("COLUMNS 2")
-    print(columns)
-    cur.close()
-    connection.close()
+def tools_call(session_hash, data_source, titles):
 
-    column_string = (columns[:625] + '..') if len(columns) > 625 else columns
+    titles_string = (titles[:625] + '..') if len(titles) > 625 else titles
 
-    tools_calls = [
-        {
-            "type": "function",
-            "function": {
-                "name": "sql_query_func",
-                "description": f"""This is a tool useful to query a SQLite table called 'data_source' with the following Columns: {column_string}.
-                There may also be more columns in the table if the number of columns is too large to process. 
-                This function also saves the results of the query to csv file called query.csv.""",
-                "parameters": {
-                    "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",
-                            }
-                        }
+    tools_calls = {
+            'file_upload' : [
+                {
+                    "type": "function",
+                    "function": {
+                        "name": "sqlite_query_func",
+                        "description": f"""This is a tool useful to query a SQLite table called 'data_source' with the following Columns: {titles_string}.
+                        There may also be more columns in the table if the number of columns is too large to process. 
+                        This function also saves the results of the query to csv file called query.csv.""",
+                        "parameters": {
+                            "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",
+                                    }
+                                }
+                            },
+                            "required": ["queries"],
+                        },
                     },
-                    "required": ["queries"],
                 },
-            },
-        },
-    ]
-
-    tools_calls.extend(chart_tools)
-    tools_calls.extend(stats_tools)
-
-    return tools_calls
-
-def sql_tools_call(db_tables):
-    
-    table_string = (db_tables[:625] + '..') if len(db_tables) > 625 else db_tables
-        
-    tools_calls = [
-        {
-            "type": "function",
-            "function": {
-                "name": "sql_query_func",
-                "description": f"""This is a tool useful to query a PostgreSQL database with the following tables, {table_string}.
-                There may also be more tables in the database if the number of tables is too large to process.
-                This function also saves the results of the query to csv file called query.csv.""",
-                "parameters": {
-                    "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",
+            ],
+            'sql' : [
+            {
+                "type": "function",
+                "function": {
+                    "name": "sql_query_func",
+                    "description": f"""This is a tool useful to query a PostgreSQL database with the following tables, {titles_string}.
+                    There may also be more tables in the database if the number of tables is too large to process.
+                    This function also saves the results of the query to csv file called query.csv.""",
+                    "parameters": {
+                        "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",
+                                }
                             }
-                        }
+                        },
+                        "required": ["queries"],
                     },
-                    "required": ["queries"],
                 },
             },
-        },
-    ]
-
-    tools_calls.extend(chart_tools)
-    tools_calls.extend(stats_tools)
-
-    return tools_calls
-
-def doc_db_tools_call(db_collections):
-    
-    collection_string = (db_collections[:625] + '..') if len(db_collections) > 625 else db_collections
-        
-    tools_calls = [
-        {
-            "type": "function",
-            "function": {
-                "name": "doc_db_query_func",
-                "description": f"""This is a tool useful to build an aggregation pipeline to query a MongoDB NoSQL document database with the following collections, {collection_string}.
-                There may also be more collections in the database if the number of tables is too large to process.
-                This function also saves the results of the query to a csv file called query.csv.""",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "aggregation_pipeline": {
-                            "type": "string",
-                            "description": "The MongoDB aggregation pipeline to use in the search. Infer this from the user's message. It should be a question or a statement."
+        ],
+        'doc_db' : [
+            {
+                "type": "function",
+                "function": {
+                    "name": "doc_db_query_func",
+                    "description": f"""This is a tool useful to build an aggregation pipeline to query a MongoDB NoSQL document database with the following collections, {titles_string}.
+                    There may also be more collections in the database if the number of tables is too large to process.
+                    This function also saves the results of the query to a csv file called query.csv.""",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "aggregation_pipeline": {
+                                "type": "string",
+                                "description": "The MongoDB aggregation pipeline to use in the search. Infer this from the user's message. It should be a question or a statement."
+                            },
+                            "db_collection": {
+                                "type": "string",
+                                "description": "The MongoDB collection to use in the search. Infer this from the user's message. It should be a question or a statement.",
+                            }
                         },
-                        "db_collection": {
-                            "type": "string",
-                            "description": "The MongoDB collection to use in the search. Infer this from the user's message. It should be a question or a statement.",
-                        }
+                        "required": ["aggregation_pipeline","db_collection"],
                     },
-                    "required": ["aggregation_pipeline","db_collection"],
                 },
             },
-        },
-    ]
-
-    tools_calls.extend(chart_tools)
-    tools_calls.extend(stats_tools)
-
-    return tools_calls
-
-def graphql_tools_call(graphql_types):
-    
-    types_string = (graphql_types[:625] + '..') if len(graphql_types) > 625 else graphql_types
-        
-    tools_calls = [
-        {
-            "type": "function",
-            "function": {
-                "name": "graphql_query_func",
-                "description": f"""This is a tool useful to build a GraphQL query for a GraphQL API endpoint with the following types, {types_string}.
-                There may also be more types in the GraphQL endpoint if the number of types is too large to process. 
-                This function also saves the results of the query to a csv file called query.csv.""",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "graphql_query": {
-                            "type": "string",
-                            "description": "The GraphQL query to use in the search. Infer this from the user's message. It should be a question or a statement."
-                        }
+        ],
+        'graphql' : [
+            {
+                "type": "function",
+                "function": {
+                    "name": "graphql_query_func",
+                    "description": f"""This is a tool useful to build a GraphQL query for a GraphQL API endpoint with the following types, {titles_string}.
+                    There may also be more types in the GraphQL endpoint if the number of types is too large to process. 
+                    This function also saves the results of the query to a csv file called query.csv.""",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "graphql_query": {
+                                "type": "string",
+                                "description": "The GraphQL query to use in the search. Infer this from the user's message. It should be a question or a statement."
+                            }
+                        },
+                        "required": ["graphql_query"],
                     },
-                    "required": ["graphql_query"],
                 },
             },
-        },
-        {
-            "type": "function",
-            "function": {
-                "name": "graphql_schema_query",
-                "description": f"""This is a tool useful to query a GraphQL type and receive back information about its schema. This is useful because
-                the GraphQL introspection query is too large to be ingested all at once and this allows us to query the schema one type at a time to 
-                view it in manageable bites. You may realize after viewing the schema, that the type you selected was not appropriate for the question
-                you are attempting answer. You may then query additional types to find the appropriate types to use for your GraphQL API query.""",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "graphql_type": {
-                            "type": "string",
-                            "description": "The GraphQL type that we want to view the schema of in order to make the proper query with our graphql_query_func. Infer this from the user's message. It should be a question or a statement."
-                        }
+            {
+                "type": "function",
+                "function": {
+                    "name": "graphql_schema_query",
+                    "description": f"""This is a tool useful to query a GraphQL type and receive back information about its schema. This is useful because
+                    the GraphQL introspection query is too large to be ingested all at once and this allows us to query the schema one type at a time to 
+                    view it in manageable bites. You may realize after viewing the schema, that the type you selected was not appropriate for the question
+                    you are attempting answer. You may then query additional types to find the appropriate types to use for your GraphQL API query.""",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "graphql_type": {
+                                "type": "string",
+                                "description": "The GraphQL type that we want to view the schema of in order to make the proper query with our graphql_query_func. Infer this from the user's message. It should be a question or a statement."
+                            }
+                        },
+                        "required": ["graphql_type"],
                     },
-                    "required": ["graphql_type"],
                 },
             },
-        },
-        {
-            "type": "function",
-            "function": {
-                "name": "graphql_csv_query",
-                "description": f"""This is a tool useful to SQL query our query.csv file that is generated from our GraphQL query. This is useful in a situation
-                where the results of the GraphQL query need additional querying to answer the user question.  The query.csv file is converted to a Pandas dataframe
-                and we query that dataframe with SQL on a table called 'query' before converting it back to a csv file.""",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "csv_query": {
-                            "type": "string",
-                            "description": "The pandas dataframe SQL query to use in the search. The table that we query is named 'query'. Infer this from the user's message. It should be a question or a statement"
-                        }
+            {
+                "type": "function",
+                "function": {
+                    "name": "graphql_csv_query",
+                    "description": f"""This is a tool useful to SQL query our query.csv file that is generated from our GraphQL query. This is useful in a situation
+                    where the results of the GraphQL query need additional querying to answer the user question.  The query.csv file is converted to a Pandas dataframe
+                    and we query that dataframe with SQL on a table called 'query' before converting it back to a csv file.""",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {
+                            "csv_query": {
+                                "type": "string",
+                                "description": "The pandas dataframe SQL query to use in the search. The table that we query is named 'query'. Infer this from the user's message. It should be a question or a statement"
+                            }
+                        },
+                        "required": ["csv_query"],
                     },
-                    "required": ["csv_query"],
                 },
             },
-        },
-    ]
+        ]
+    }
+
+    tools = tools_calls[data_source]
 
-    tools_calls.extend(chart_tools)
-    tools_calls.extend(stats_tools)
+    tools.extend(chart_tools)
+    tools.extend(stats_tools)
 
-    return tools_calls
+    return tools