feat: user-supplied API keys, model selector, and Anthropic support

- Replace hardwired OpenAI key with a per-session API key form at the top
of the UI; key is stored in memory only and cleared on page unload
- Add model dropdown that dynamically switches between OpenAI and Anthropic
model lists based on detected key prefix (sk-ant- vs sk-)
- Add anthropic-haystack dependency for AnthropicChatGenerator support
- Refactor tools/chart_tools.py, tools/stats_tools.py, and tools/tools.py
to use Haystack's provider-agnostic Tool class instead of raw OpenAI-format
dicts, enabling tool calling to work with both OpenAI and Anthropic models
- Fix invalid JSON Schema (items field on non-array properties) in tool
parameter definitions for stricter Anthropic compatibility
- Pass tools via tools= parameter to chat_generator.run() rather than
generation_kwargs so both generators handle format conversion natively

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

app.py

@@ -1,9 +1,8 @@
-from utils import TEMP_DIR, message_dict
+from utils import TEMP_DIR, message_dict, api_key_store, model_store
 import gradio as gr
 import templates.data_file as data_file, templates.sql_db as sql_db, templates.doc_db as doc_db, templates.graphql as graphql
 
 import os
-from getpass import getpass
 from dotenv import load_dotenv
 
 load_dotenv()
@@ -14,9 +13,16 @@ def delete_db(req: gr.Request):
     if os.path.exists(dir_path):
         shutil.rmtree(dir_path)
         message_dict[req.session_hash] = {}
+    api_key_store.pop(req.session_hash, None)
+    model_store.pop(req.session_hash, None)
 
-if "OPENAI_API_KEY" not in os.environ:
-    os.environ["OPENAI_API_KEY"] = getpass("Enter OpenAI API key:")
+def set_api_key(api_key, model, request: gr.Request):
+    api_key = api_key.strip()
+    if not api_key:
+        return gr.update(visible=True), "<p style='color:#b91c1c;text-align:center;'>Please enter your OpenAI API key.</p>"
+    api_key_store[request.session_hash] = api_key
+    model_store[request.session_hash] = model
+    return gr.update(visible=False), f"<p style='color:#15803d;text-align:center;font-weight:500;'>&#10003; API key saved &mdash; using <strong>{model}</strong>. You can now use all features below.</p>"
 
 css= ".file_marker .large{min-height:50px !important;} .padding{padding:0;} .description_component{overflow:visible !important;}"
 head = """<meta charset="UTF-8">
@@ -40,7 +46,56 @@ theme = gr.themes.Base(primary_hue="sky", secondary_hue="slate",font=[gr.themes.
 from pathlib import Path
 gr.set_static_paths(paths=[Path.cwd().absolute()/"assets"])
 
+_env_api_key = os.getenv("OPENAI_API_KEY", "")
+
+OPENAI_MODELS = [
+    "gpt-4o", "gpt-4o-mini",
+    "gpt-4.1", "gpt-4.1-mini", "gpt-4.1-nano",
+    "o3-mini", "o4-mini",
+    "gpt-5.4-mini", "gpt-5.4", "gpt-5.5",
+]
+ANTHROPIC_MODELS = [
+    "claude-sonnet-4-6",
+    "claude-opus-4-8",
+    "claude-haiku-4-5-20251001",
+]
+
+def update_models(api_key):
+    if api_key.strip().startswith("sk-ant-"):
+        return gr.update(choices=ANTHROPIC_MODELS, value=ANTHROPIC_MODELS[0])
+    return gr.update(choices=OPENAI_MODELS, value=OPENAI_MODELS[0])
+
 with gr.Blocks(theme=theme, css=css, head=head, delete_cache=(3600,3600)) as demo:
+    with gr.Column(visible=True) as api_key_section:
+        gr.HTML("""
+            <div style="max-width:600px;margin:24px auto 8px;padding:16px 20px;background:#fffbeb;border:1px solid #f59e0b;border-radius:8px;">
+                <h3 style="color:#92400e;margin:0 0 6px;font-size:16px;font-weight:600;">
+                    <i class="fas fa-key"></i>&nbsp; OpenAI API Key Required
+                </h3>
+                <p style="color:#78350f;font-size:14px;margin:0;">
+                    Enter your OpenAI (<code>sk-...</code>) or Anthropic (<code>sk-ant-...</code>) API key. The model list updates automatically based on your key. Your key is held only in memory for your current session and is never saved to disk or shared.
+                </p>
+            </div>
+        """)
+        with gr.Row(equal_height=True):
+            api_key_input = gr.Textbox(
+                label="OpenAI API Key",
+                placeholder="sk-proj-...",
+                type="password",
+                value=_env_api_key,
+                scale=4
+            )
+            model_dropdown = gr.Dropdown(
+                label="Model",
+                choices=OPENAI_MODELS,
+                value=OPENAI_MODELS[0],
+                scale=2
+            )
+            api_key_btn = gr.Button("Set API Key", variant="primary", scale=1, min_width=120)
+    api_key_msg = gr.HTML("")
+    api_key_input.change(fn=update_models, inputs=api_key_input, outputs=model_dropdown)
+    api_key_btn.click(fn=set_api_key, inputs=[api_key_input, model_dropdown], outputs=[api_key_section, api_key_msg])
+
     header = gr.HTML("""
                         <!-- Header -->
                         <header class="max-w-4xl mx-auto mb-12 text-center">

assets/styles.css

@@ -174,4 +174,14 @@
         grid-template-columns: 1fr 2fr;
         align-items: baseline;
     }
-  }
+  }
+  
+dialog {
+  margin: 10% auto;
+  width: 80%;
+  max-width: 350px;
+  background-color: #fff;
+  padding: 34px;
+  border: 0;
+  border-radius: 5px;
+}  

functions/chat_functions.py

@@ -1,9 +1,19 @@
-from utils import message_dict
+from utils import message_dict, api_key_store, model_store
 
 from haystack.dataclasses import ChatMessage
 from haystack.components.generators.chat import OpenAIChatGenerator
+from haystack.utils import Secret
+
+def _get_generator(session_hash):
+    api_key = api_key_store.get(session_hash)
+    if not api_key:
+        raise ValueError("No API key found for this session. Please enter your API key at the top of the page.")
+    model = model_store.get(session_hash, "gpt-4o")
+    if api_key.startswith("sk-ant-"):
+        from haystack_integrations.components.generators.chat import AnthropicChatGenerator
+        return AnthropicChatGenerator(model=model, api_key=Secret.from_token(api_key))
+    return OpenAIChatGenerator(model=model, api_key=Secret.from_token(api_key))
 
-chat_generator = OpenAIChatGenerator(model="gpt-4o")
 response = None
 
 def example_question_message(data_source, name, titles, schema):
@@ -13,7 +23,7 @@ def example_question_message(data_source, name, titles, schema):
                          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
+                        that we can suggest that you believe will be insightful or helpful to a data analyst 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."""],
 
@@ -21,7 +31,7 @@ def example_question_message(data_source, name, titles, schema):
                  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
+                        that we can suggest that you believe will be insightful or helpful to a data analyst 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."""],
 
@@ -30,7 +40,7 @@ def example_question_message(data_source, name, titles, schema):
                         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
+                        that we can suggest that you believe will be insightful or helpful to a data analyst 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."""],
 
@@ -38,7 +48,7 @@ def example_question_message(data_source, name, titles, schema):
                      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
+                        that we can suggest that you believe will be insightful or helpful to a data analyst 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."""]         
 
@@ -57,9 +67,16 @@ def example_question_generator(session_hash, data_source, name, titles, schema):
 
     example_messages.append(ChatMessage.from_user(text=example_message_list[1]))
 
-    example_response = chat_generator.run(messages=example_messages)
+    example_response = _get_generator(session_hash).run(messages=example_messages)
 
-    return example_response["replies"][0].text
+    response_text = example_response["replies"][0].text
+    start = response_text.index("[") + 1
+    end = response_text.index("]")
+    response_content = response_text[start:end]
+    response_list = '[' + response_content + ']'
+    print(response_list)
+
+    return response_list
 
 def system_message(data_source, titles, schema=""):
     print("TITLES")
@@ -112,6 +129,11 @@ def system_message(data_source, titles, schema=""):
     return system_message_dict[data_source]
 
 def chatbot_func(message, history, session_hash, data_source, titles, schema, *args):
+    try:
+        chat_generator = _get_generator(session_hash)
+    except ValueError as e:
+        return str(e)
+
     from functions import table_generation_func, regression_func, scatter_chart_generation_func, \
         query_func, graphql_schema_query, graphql_csv_query, \
         line_chart_generation_func,bar_chart_generation_func,pie_chart_generation_func,histogram_generation_func
@@ -133,10 +155,11 @@ def chatbot_func(message, history, session_hash, data_source, titles, schema, *a
         messages.append(ChatMessage.from_user(message))
         message_dict[session_hash][data_source] = messages
     
-    response = chat_generator.run(messages=message_dict[session_hash][data_source], generation_kwargs={"tools": tools.tools_call(session_hash, data_source, titles)})
+    active_tools = tools.tools_call(session_hash, data_source, titles)
+    response = chat_generator.run(messages=message_dict[session_hash][data_source], tools=active_tools)
 
     while True:
-        # if OpenAI response is a tool call
+        # if the 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:
@@ -151,7 +174,7 @@ def chatbot_func(message, history, session_hash, data_source, titles, schema, *a
                 print(function_name)
                 ## Append function response to the messages list using `ChatMessage.from_tool`
                 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)})
+                response = chat_generator.run(messages=message_dict[session_hash][data_source], tools=active_tools)
 
         # Regular Conversation
         else:

requirements.txt

@@ -1,4 +1,5 @@
 haystack-ai
+anthropic-haystack
 python-dotenv
 gradio
 pandas

templates/data_file.py

@@ -97,7 +97,7 @@ with gr.Blocks() as demo:
                                         ]
                 else:
                     try:
-                        generated_examples = ast.literal_eval(example_question_generator(request.session_hash, 'file_upload', '', process_message[1], ''))
+                        generated_examples = ast.literal_eval(example_question_generator(request.session_hash, 'file_upload', '', process_message[2], ''))
                         example_questions = [
                                                 ["Describe the dataset"]
                                             ]
@@ -111,6 +111,30 @@ with gr.Blocks() as demo:
                                             ["List the columns in the dataset"],
                                             ["What could this data be used for?"],
                                         ]
+                data_summary = gr.HTML("""
+                        <script>                                        
+                        const modal = document.querySelector('.modal');
+                        const openButton = document.querySelector('.open-button');
+                        const closeButton = document.querySelector('.close-button');
+                                                            console.log("JS DETECTED")
+                        openButton.addEventListener('click', () => {
+                        modal.showModal();
+                        });
+
+                        closeButton.addEventListener('click', () => {
+                        modal.close();
+                        });
+                        </script>                        
+                        <button class="open-button">
+                        Open modal 🙈
+                        </button>
+                        <dialog class="modal">
+                        <p>Udohgram</p>
+                        <p>Insert pretty picture</p>
+                        <button class="close-button">Close modal</button>
+                        <p>Sorry, you can't login. You can't contact support either.</p>
+                        </dialog>                                   
+                """)
                 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='')

tools/chart_tools.py

@@ -1,371 +1,276 @@
-chart_tools = [
-    {
-        "type": "function",
-            "function": {
-                "name": "scatter_chart_generation_func",
-                "description": f"""This is a scatter plot generation tool useful to generate scatter plots from queried data from our data source that we are querying.
-                The data values will come from the columns of our query.csv (the 'x' and 'y' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
-                Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned 
-                from the scatter_chart_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary 
-                to it for context if desired.""",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "data": {
-                            "type": "array",
-                            "description": """The array containing a dictionary that contains the 'data' portion of the plotly chart generation and will include the options requested by the user. 
-                            The array must contain a json formatted dictionary with outer brackets included, any other format will not work. 
-                            Do not include the 'x' or 'y' portions of the object as this will come from the query.csv file generated by our SQLite query. 
-                            Infer this from the user's message.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "x_column": {
-                            "type": "array",
-                            "description": f"""An array of strings that correspond to the the column names in our query.csv file that contain the x values of the graph. There can be more than one column
-                            that can each be plotted against the y_column, if needed.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "y_column": {
-                            "type": "string",
-                            "description": f"""The column name in our query.csv file that contain the y values of the graph.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "category": {
-                            "type": "string",
-                            "description": f"""An optional column name in our query.csv file that contain a parameter that will define the category for the data.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "trendline": {
-                            "type": "string",
-                            "description": f"""An optional field to specify the type of plotly trendline we wish to use in the scatter plot.
-                                This trendline value can be one of ['ols','lowess','rolling','ewm','expanding']. 
-                                Do not send any values outside of this array as the function will fail. 
-                                Infer this from the user's message.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "trendline_options": {
-                            "type": "array",
-                            "description": """An array containing a dictionary that contains the 'trendline_options' portion of the plotly chart generation.
-                            The 'lowess', 'rolling', and 'ewm' options require trendline_options to be included.
-                            The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "marginal_x": {
-                            "type": "string",
-                            "description": f"""The type of marginal distribution plot we'd like to specify for the plotly scatter plot for the x axis.
-                                This marginal_x value can be one of ['histogram','rug','box','violin']. 
-                                Do not send any values outside of this array as the function will fail. 
-                                Infer this from the user's message.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },  
-                        "marginal_y": {
-                            "type": "string",
-                            "description": f"""The type of marginal distribution plot we'd like to specify for the plotly scatter plot for the y axis.
-                                This marginal_y value can be one of ['histogram','rug','box','violin']. 
-                                Do not send any values outside of this array as the function will fail. 
-                                Infer this from the user's message.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },                                               
-                        "layout": {
-                            "type": "array",
-                            "description": """An array containing a dictionary that contains the 'layout' portion of the plotly chart generation. 
-                            The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "size": {
-                            "type": "string",
-                            "description": f"""An optional column in our query.csv file that contain a parameter that will define the size of each plot point.
-                            This is useful for a bubble chart where another value in our query can be represented by the size of the plotted point.
-                            Values must be greater than or equal to 0 and so in our query, all values less than 0 should be set equal to zero.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        }
-                    },
-                    "required": ["x_column","y_column"],
-                },
-            }, 
-        }, 
-        {
-        "type": "function",
-            "function": {
-                "name": "line_chart_generation_func",
-                "description": f"""This is a line chart generation tool useful to generate line charts from queried data from our data source that we are querying.
-                The data values will come from the columns of our query.csv (the 'x' and 'y' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
-                Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned 
-                from the line_chart_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary 
-                to it for context if desired.""",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "data": {
-                            "type": "array",
-                            "description": """The array containing a dictionary that contains the 'data' portion of the plotly chart generation and will include the options requested by the user. 
-                            The array must contain a json formatted dictionary with outer brackets included, any other format will not work. 
-                            Do not include the 'x' or 'y' portions of the object as this will come from the query.csv file generated by our SQLite query. 
-                            Infer this from the user's message.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "x_column": {
-                            "type": "string",
-                            "description": f"""The column name in our query.csv file that contain the x values of the graph.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "y_column": {
-                            "type": "string",
-                            "description": f"""The column name in our query.csv file that contain the y values of the graph.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "category": {
-                            "type": "string",
-                            "description": f"""An optional column name in our query.csv file that contain a parameter that will define the category for the data.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },                                               
-                        "layout": {
-                            "type": "array",
-                            "description": """An array containing a dictionary that contains the 'layout' portion of the plotly chart generation. 
-                            The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        }
-                    },
-                    "required": ["x_column","y_column","layout"],
-                },
-            }, 
-        }, 
-        {
-        "type": "function",
-            "function": {
-                "name": "bar_chart_generation_func",
-                "description": f"""This is a bar chart generation tool useful to generate line charts from queried data from our data source that we are querying.
-                The data values will come from the columns of our query.csv (the 'x' and 'y' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
-                Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned 
-                from the bar_chart_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary 
-                to it for context if desired.""",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "data": {
-                            "type": "array",
-                            "description": """The array containing a dictionary that contains the 'data' portion of the plotly chart generation and will include the options requested by the user. 
-                            The array must contain a json formatted dictionary with outer brackets included, any other format will not work. 
-                            Do not include the 'x' or 'y' portions of the object as this will come from the query.csv file generated by our SQLite query. 
-                            Infer this from the user's message.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "x_column": {
-                            "type": "string",
-                            "description": f"""The column name in our query.csv file that contains the x values of the graph.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "y_column": {
-                            "type": "string",
-                            "description": f"""The column name in our query.csv file that contains the y values of the graph.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "category": {
-                            "type": "string",
-                            "description": f"""An optional column name in our query.csv file that contains a parameter that will define the category for the data.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "facet_row": {
-                            "type": "string",
-                            "description": f"""An optional column name in our query.csv file that contains a parameter that will define a faceted subplot, where different rows
-                            correspond to different values of the query specified in this parameter.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "facet_col": {
-                            "type": "string",
-                            "description": f"""An optional column name in our query.csv file that contain a parameter that will define the faceted column, corresponding to 
-                            different values of our query specified in this parameter.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },                                            
-                        "layout": {
-                            "type": "array",
-                            "description": """An array containing a dictionary that contains the 'layout' portion of the plotly chart generation. 
-                            The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        }
-                    },
-                    "required": ["x_column","y_column","layout"],
-                },
-            }, 
-        },  
-        {
-        "type": "function",
-            "function": {
-                "name": "pie_chart_generation_func",
-                "description": f"""This is a pie chart generation tool useful to generate pie charts from queried data from our data source that we are querying.
-                The data values will come from the columns of our query.csv (the 'values' and 'names' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
-                Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned 
-                from the pie_chart_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary 
-                to it for context if desired.""",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "data": {
-                            "type": "array",
-                            "description": """The array containing a dictionary that contains the 'data' portion of the plotly chart generation and will include the options requested by the user. 
-                            The array must contain a json formatted dictionary with outer brackets included, any other format will not work. 
-                            Do not include the 'x' or 'y' portions of the object as this will come from the query.csv file generated by our SQLite query. 
-                            Infer this from the user's message.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "values": {
-                            "type": "string",
-                            "description": f"""The column name in our query.csv file that contain the values of the pie chart.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "names": {
-                            "type": "string",
-                            "description": f"""The column name in our query.csv file that contain the label or section of each piece of the pie graph and allow us to know what each piece of the pie chart represents.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },                                           
-                        "layout": {
-                            "type": "array",
-                            "description": """An array containing a dictionary that contains the 'layout' portion of the plotly chart generation. 
-                            The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        }
-                    },
-                    "required": ["values","names","layout"],
-                },
-            }, 
-        },
-        {
-        "type": "function",
-            "function": {
-                "name": "histogram_generation_func",
-                "description": f"""This is a histogram generation tool useful to generate histograms from queried data from our data source that we are querying.
-                The data values will come from the columns of our query.csv (the 'values' and 'names' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
-                Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned 
-                from the histogram_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary 
-                to it for context if desired.""",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "data": {
-                            "type": "array",
-                            "description": """The array containing a dictionary that contains the 'data' portion of the plotly chart generation and will include the options requested by the user. 
-                            The array must contain a json formatted dictionary with outer brackets included, any other format will not work. 
-                            Do not include the 'x' or 'y' portions of the object as this will come from the query.csv file generated by our SQLite query. 
-                            Infer this from the user's message.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "x_column": {
-                            "type": "string",
-                            "description": f"""The column name in our query.csv file that contains the x values of the histogram. 
-                            This would correspond to the counts that would be distributed in the histogram.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "y_column": {
-                            "type": "string",
-                            "description": f"""An optional column name in our query.csv file that contains the y values of the histogram.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "histnorm": {
-                            "type": "string",
-                            "description": f"""An optional argument to specify the type of normalization if the default isn't used.
-                            This histnorm value can be one of ['percent','probability','density','probability density'].
-                            Do not send any values outside of this array as the function will fail.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        }, 
-                        "category": {
-                            "type": "string",
-                            "description": f"""An optional column name in our query.csv file that contains a parameter that will define the category for the data.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "histfunc": {
-                            "type": "string",
-                            "description": f"""An optional value that represents the function of data to compute the function which is used on the optional y column.
-                            This histfunc value can be one of ['avg','sum','count']. 
-                                Do not send any values outside of this array as the function will fail.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },                                         
-                        "layout": {
-                            "type": "array",
-                            "description": """An array containing a dictionary that contains the 'layout' portion of the plotly chart generation. 
-                            The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        }
-                    },
-                    "required": ["x_column"],
-                },
-            }, 
-        },                     
-        {
-        "type": "function",
-            "function": {
-                "name": "table_generation_func",
-                "description": f"""This an table generation tool useful to format data as a table from queried data from our data source that we are querying.
-                Takes no parameters as it uses data queried in our query.csv file to build the table.
-                Call this function after running our SQLite query and generating query.csv.
-                Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned 
-                from the table_generation_func function in any way and always display the iframe fully to the user in the chat window.""",
-                "parameters": {},
-            }, 
-        }
-]
+from haystack.tools import Tool
+
+_noop = lambda **kwargs: None
+
+chart_tools = [
+    Tool(
+        name="scatter_chart_generation_func",
+        description="""This is a scatter plot generation tool useful to generate scatter plots from queried data from our data source that we are querying.
+        The data values will come from the columns of our query.csv (the 'x' and 'y' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
+        Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned
+        from the scatter_chart_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary
+        to it for context if desired.""",
+        parameters={
+            "type": "object",
+            "properties": {
+                "data": {
+                    "type": "array",
+                    "description": """The array containing a dictionary that contains the 'data' portion of the plotly chart generation and will include the options requested by the user.
+                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.
+                    Do not include the 'x' or 'y' portions of the object as this will come from the query.csv file generated by our SQLite query.
+                    Infer this from the user's message.""",
+                    "items": {"type": "string"}
+                },
+                "x_column": {
+                    "type": "array",
+                    "description": """An array of strings that correspond to the column names in our query.csv file that contain the x values of the graph. There can be more than one column
+                    that can each be plotted against the y_column, if needed.""",
+                    "items": {"type": "string"}
+                },
+                "y_column": {
+                    "type": "string",
+                    "description": """The column name in our query.csv file that contain the y values of the graph."""
+                },
+                "category": {
+                    "type": "string",
+                    "description": """An optional column name in our query.csv file that contain a parameter that will define the category for the data."""
+                },
+                "trendline": {
+                    "type": "string",
+                    "description": """An optional field to specify the type of plotly trendline we wish to use in the scatter plot.
+                        This trendline value can be one of ['ols','lowess','rolling','ewm','expanding'].
+                        Do not send any values outside of this array as the function will fail.
+                        Infer this from the user's message."""
+                },
+                "trendline_options": {
+                    "type": "array",
+                    "description": """An array containing a dictionary that contains the 'trendline_options' portion of the plotly chart generation.
+                    The 'lowess', 'rolling', and 'ewm' options require trendline_options to be included.
+                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
+                    "items": {"type": "string"}
+                },
+                "marginal_x": {
+                    "type": "string",
+                    "description": """The type of marginal distribution plot we'd like to specify for the plotly scatter plot for the x axis.
+                        This marginal_x value can be one of ['histogram','rug','box','violin'].
+                        Do not send any values outside of this array as the function will fail.
+                        Infer this from the user's message."""
+                },
+                "marginal_y": {
+                    "type": "string",
+                    "description": """The type of marginal distribution plot we'd like to specify for the plotly scatter plot for the y axis.
+                        This marginal_y value can be one of ['histogram','rug','box','violin'].
+                        Do not send any values outside of this array as the function will fail.
+                        Infer this from the user's message."""
+                },
+                "layout": {
+                    "type": "array",
+                    "description": """An array containing a dictionary that contains the 'layout' portion of the plotly chart generation.
+                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
+                    "items": {"type": "string"}
+                },
+                "size": {
+                    "type": "string",
+                    "description": """An optional column in our query.csv file that contain a parameter that will define the size of each plot point.
+                    This is useful for a bubble chart where another value in our query can be represented by the size of the plotted point.
+                    Values must be greater than or equal to 0 and so in our query, all values less than 0 should be set equal to zero."""
+                }
+            },
+            "required": ["x_column", "y_column"]
+        },
+        function=_noop
+    ),
+    Tool(
+        name="line_chart_generation_func",
+        description="""This is a line chart generation tool useful to generate line charts from queried data from our data source that we are querying.
+        The data values will come from the columns of our query.csv (the 'x' and 'y' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
+        Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned
+        from the line_chart_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary
+        to it for context if desired.""",
+        parameters={
+            "type": "object",
+            "properties": {
+                "data": {
+                    "type": "array",
+                    "description": """The array containing a dictionary that contains the 'data' portion of the plotly chart generation and will include the options requested by the user.
+                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.
+                    Do not include the 'x' or 'y' portions of the object as this will come from the query.csv file generated by our SQLite query.
+                    Infer this from the user's message.""",
+                    "items": {"type": "string"}
+                },
+                "x_column": {
+                    "type": "string",
+                    "description": """The column name in our query.csv file that contain the x values of the graph."""
+                },
+                "y_column": {
+                    "type": "string",
+                    "description": """The column name in our query.csv file that contain the y values of the graph."""
+                },
+                "category": {
+                    "type": "string",
+                    "description": """An optional column name in our query.csv file that contain a parameter that will define the category for the data."""
+                },
+                "layout": {
+                    "type": "array",
+                    "description": """An array containing a dictionary that contains the 'layout' portion of the plotly chart generation.
+                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
+                    "items": {"type": "string"}
+                }
+            },
+            "required": ["x_column", "y_column", "layout"]
+        },
+        function=_noop
+    ),
+    Tool(
+        name="bar_chart_generation_func",
+        description="""This is a bar chart generation tool useful to generate bar charts from queried data from our data source that we are querying.
+        The data values will come from the columns of our query.csv (the 'x' and 'y' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
+        Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned
+        from the bar_chart_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary
+        to it for context if desired.""",
+        parameters={
+            "type": "object",
+            "properties": {
+                "data": {
+                    "type": "array",
+                    "description": """The array containing a dictionary that contains the 'data' portion of the plotly chart generation and will include the options requested by the user.
+                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.
+                    Do not include the 'x' or 'y' portions of the object as this will come from the query.csv file generated by our SQLite query.
+                    Infer this from the user's message.""",
+                    "items": {"type": "string"}
+                },
+                "x_column": {
+                    "type": "string",
+                    "description": """The column name in our query.csv file that contains the x values of the graph."""
+                },
+                "y_column": {
+                    "type": "string",
+                    "description": """The column name in our query.csv file that contains the y values of the graph."""
+                },
+                "category": {
+                    "type": "string",
+                    "description": """An optional column name in our query.csv file that contains a parameter that will define the category for the data."""
+                },
+                "facet_row": {
+                    "type": "string",
+                    "description": """An optional column name in our query.csv file that contains a parameter that will define a faceted subplot, where different rows
+                    correspond to different values of the query specified in this parameter."""
+                },
+                "facet_col": {
+                    "type": "string",
+                    "description": """An optional column name in our query.csv file that contain a parameter that will define the faceted column, corresponding to
+                    different values of our query specified in this parameter."""
+                },
+                "layout": {
+                    "type": "array",
+                    "description": """An array containing a dictionary that contains the 'layout' portion of the plotly chart generation.
+                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
+                    "items": {"type": "string"}
+                }
+            },
+            "required": ["x_column", "y_column", "layout"]
+        },
+        function=_noop
+    ),
+    Tool(
+        name="pie_chart_generation_func",
+        description="""This is a pie chart generation tool useful to generate pie charts from queried data from our data source that we are querying.
+        The data values will come from the columns of our query.csv (the 'values' and 'names' values of each graph) file but the layout section of the plotly dictionary objects will be generated by you.
+        Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned
+        from the pie_chart_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary
+        to it for context if desired.""",
+        parameters={
+            "type": "object",
+            "properties": {
+                "data": {
+                    "type": "array",
+                    "description": """The array containing a dictionary that contains the 'data' portion of the plotly chart generation and will include the options requested by the user.
+                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.
+                    Do not include the 'x' or 'y' portions of the object as this will come from the query.csv file generated by our SQLite query.
+                    Infer this from the user's message.""",
+                    "items": {"type": "string"}
+                },
+                "values": {
+                    "type": "string",
+                    "description": """The column name in our query.csv file that contain the values of the pie chart."""
+                },
+                "names": {
+                    "type": "string",
+                    "description": """The column name in our query.csv file that contain the label or section of each piece of the pie graph and allow us to know what each piece of the pie chart represents."""
+                },
+                "layout": {
+                    "type": "array",
+                    "description": """An array containing a dictionary that contains the 'layout' portion of the plotly chart generation.
+                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
+                    "items": {"type": "string"}
+                }
+            },
+            "required": ["values", "names", "layout"]
+        },
+        function=_noop
+    ),
+    Tool(
+        name="histogram_generation_func",
+        description="""This is a histogram generation tool useful to generate histograms from queried data from our data source that we are querying.
+        The data values will come from the columns of our query.csv file but the layout section of the plotly dictionary objects will be generated by you.
+        Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned
+        from the histogram_generation_func function in any way and always display the iframe fully to the user in the chat window. You can add your own text supplementary
+        to it for context if desired.""",
+        parameters={
+            "type": "object",
+            "properties": {
+                "data": {
+                    "type": "array",
+                    "description": """The array containing a dictionary that contains the 'data' portion of the plotly chart generation and will include the options requested by the user.
+                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.
+                    Do not include the 'x' or 'y' portions of the object as this will come from the query.csv file generated by our SQLite query.
+                    Infer this from the user's message.""",
+                    "items": {"type": "string"}
+                },
+                "x_column": {
+                    "type": "string",
+                    "description": """The column name in our query.csv file that contains the x values of the histogram.
+                    This would correspond to the counts that would be distributed in the histogram."""
+                },
+                "y_column": {
+                    "type": "string",
+                    "description": """An optional column name in our query.csv file that contains the y values of the histogram."""
+                },
+                "histnorm": {
+                    "type": "string",
+                    "description": """An optional argument to specify the type of normalization if the default isn't used.
+                    This histnorm value can be one of ['percent','probability','density','probability density'].
+                    Do not send any values outside of this array as the function will fail."""
+                },
+                "category": {
+                    "type": "string",
+                    "description": """An optional column name in our query.csv file that contains a parameter that will define the category for the data."""
+                },
+                "histfunc": {
+                    "type": "string",
+                    "description": """An optional value that represents the function of data to compute the function which is used on the optional y column.
+                    This histfunc value can be one of ['avg','sum','count'].
+                    Do not send any values outside of this array as the function will fail."""
+                },
+                "layout": {
+                    "type": "array",
+                    "description": """An array containing a dictionary that contains the 'layout' portion of the plotly chart generation.
+                    The array must contain a json formatted dictionary with outer brackets included, any other format will not work.""",
+                    "items": {"type": "string"}
+                }
+            },
+            "required": ["x_column"]
+        },
+        function=_noop
+    ),
+    Tool(
+        name="table_generation_func",
+        description="""This is a table generation tool useful to format data as a table from queried data from our data source that we are querying.
+        Takes no parameters as it uses data queried in our query.csv file to build the table.
+        Call this function after running our query and generating query.csv.
+        Returns an iframe string which will be displayed inline in our chat window. Do not edit the iframe string returned
+        from the table_generation_func function in any way and always display the iframe fully to the user in the chat window.""",
+        parameters={"type": "object", "properties": {}},
+        function=_noop
+    ),
+]

tools/stats_tools.py

@@ -1,44 +1,38 @@
-stats_tools = [
-        {
-        "type": "function",
-            "function": {
-                "name": "regression_func",
-                "description": f"""This a tool to calculate regressions on our data source that we are querying.
-                We can run queries with our 'sql_query_func' function and they will be available to use in this function via the query.csv file that is generated.
-                Returns a dictionary of values that includes a regression_summary and a regression chart (which is an iframe displaying the
-                linear regression in chart form and should be shown to the user).""",
-                "parameters": {
-                    "type": "object",
-                    "properties": {
-                        "independent_variables": {
-                            "type": "array",
-                            "description": f"""An array of strings that states the independent variables in our data set which should be column names in our query.csv file that is generated
-                            in the 'sql_query_func' function. This will allow us to identify the data to use for our independent variables.
-                            Infer this from the user's message.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "dependent_variable": {
-                            "type": "string",
-                            "description": f"""A string that states the dependent variables in our data set which should be a column name in our query.csv file that is generated
-                            in the 'sql_query_func' function. This will allow us to identify the data to use for our dependent variables.
-                            Infer this from the user's message.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        },
-                        "category": {
-                            "type": "string",
-                            "description": f"""An optional column name in our query.csv file that contain a parameter that will define the category for the data. 
-                            Do not send value if no category is needed or specified. This category must be present in our query.csv file to be valid.""",
-                            "items": {
-                                "type": "string",
-                            }
-                        }
-                    },
-                    "required": ["independent_variables","dependent_variable"],
-                },
-            }, 
-        }    
-]
+from haystack.tools import Tool
+
+_noop = lambda **kwargs: None
+
+stats_tools = [
+    Tool(
+        name="regression_func",
+        description="""This a tool to calculate regressions on our data source that we are querying.
+        We can run queries with our 'sql_query_func' function and they will be available to use in this function via the query.csv file that is generated.
+        Returns a dictionary of values that includes a regression_summary and a regression chart (which is an iframe displaying the
+        linear regression in chart form and should be shown to the user).""",
+        parameters={
+            "type": "object",
+            "properties": {
+                "independent_variables": {
+                    "type": "array",
+                    "description": """An array of strings that states the independent variables in our data set which should be column names in our query.csv file that is generated
+                    in the 'sql_query_func' function. This will allow us to identify the data to use for our independent variables.
+                    Infer this from the user's message.""",
+                    "items": {"type": "string"}
+                },
+                "dependent_variable": {
+                    "type": "string",
+                    "description": """A string that states the dependent variable in our data set which should be a column name in our query.csv file that is generated
+                    in the 'sql_query_func' function. This will allow us to identify the data to use for our dependent variable.
+                    Infer this from the user's message."""
+                },
+                "category": {
+                    "type": "string",
+                    "description": """An optional column name in our query.csv file that contain a parameter that will define the category for the data.
+                    Do not send value if no category is needed or specified. This category must be present in our query.csv file to be valid."""
+                }
+            },
+            "required": ["independent_variables", "dependent_variable"]
+        },
+        function=_noop
+    )
+]

tools/tools.py

@@ -1,143 +1,127 @@
-from .stats_tools import stats_tools
-from .chart_tools import chart_tools
-
-def tools_call(session_hash, data_source, titles):
-
-    titles_string = (titles[:625] + '..') if len(titles) > 625 else titles
-
-    tools_calls = {
-            'file_upload' : [
-                {
-                    "type": "function",
-                    "function": {
-                        "name": "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": "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"],
-                        },
-                    },
-                },
-            ],
-            'sql' : [
-            {
-                "type": "function",
-                "function": {
-                    "name": "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": "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"],
-                    },
-                },
-            },
-        ],
-        'doc_db' : [
-            {
-                "type": "function",
-                "function": {
-                    "name": "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": {
-                            "queries": {
-                                "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.",
-                            }
-                        },
-                        "required": ["queries","db_collection"],
-                    },
-                },
-            },
-        ],
-        'graphql' : [
-            {
-                "type": "function",
-                "function": {
-                    "name": "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": {
-                            "queries": {
-                                "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": ["queries"],
-                    },
-                },
-            },
-            {
-                "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"],
-                    },
-                },
-            },
-            {
-                "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"],
-                    },
-                },
-            },
-        ]
-    }
-
-    tools = tools_calls[data_source]
-
-    tools.extend(chart_tools)
-    tools.extend(stats_tools)
-
-    return tools
+from haystack.tools import Tool
+from .stats_tools import stats_tools
+from .chart_tools import chart_tools
+
+_noop = lambda **kwargs: None
+
+def tools_call(session_hash, data_source, titles):
+
+    titles_string = (titles[:625] + '..') if len(titles) > 625 else titles
+
+    query_tools = {
+        'file_upload': Tool(
+            name="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": "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"]
+            },
+            function=_noop
+        ),
+        'sql': Tool(
+            name="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": "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"]
+            },
+            function=_noop
+        ),
+        'doc_db': Tool(
+            name="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 collections 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": {
+                    "queries": {
+                        "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."
+                    }
+                },
+                "required": ["queries", "db_collection"]
+            },
+            function=_noop
+        ),
+        'graphql': [
+            Tool(
+                name="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": {
+                        "queries": {
+                            "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": ["queries"]
+                },
+                function=_noop
+            ),
+            Tool(
+                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"]
+                },
+                function=_noop
+            ),
+            Tool(
+                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"]
+                },
+                function=_noop
+            ),
+        ]
+    }
+
+    source_tools = query_tools[data_source]
+    tools = source_tools if isinstance(source_tools, list) else [source_tools]
+    tools = tools + chart_tools + stats_tools
+
+    return tools

utils.py

@@ -4,4 +4,6 @@ current_dir = Path(__file__).parent
 
 TEMP_DIR = current_dir / 'temp'
 
-message_dict = {}
+message_dict = {}
+api_key_store = {}
+model_store = {}