{ "cells": [ { "cell_type": "markdown", "id": "104f2b97-f9bb-4dcc-a4c8-099710768851", "metadata": {}, "source": [ "# Using JSON Mode and Function Calling for SQL Querying" ] }, { "cell_type": "markdown", "id": "f946b4f9-3993-41a5-8e82-77623e9d546f", "metadata": {}, "source": [ "With the rise of Large Language Models (LLMs), one of the first practical applications has been the \"chat with X\" app. In this notebook we will explore methods for building \"chat with my database\" tools with Groq API, exploring benefits and drawbacks of each and leveraging Groq API's [JSON mode](https://console.groq.com/docs/text-chat#json-mode) and [tool use](https://console.groq.com/docs/tool-use) feature for function calling.\n", "\n", "We will show two methods for using Groq API to query a database, and how leveraging tool use for function calling can improve the predictability and reliability of our outputs. We will use the [DuckDB](https://duckdb.org/) query language on local CSV files in this example, but the general framework could also work against standard data warehouse platforms like [BigQuery](https://cloud.google.com/bigquery)." ] }, { "cell_type": "markdown", "id": "f8dc57b6-2c48-4ee3-bb2c-25441274ed2f", "metadata": {}, "source": [ "### Setup" ] }, { "cell_type": "code", "execution_count": 1, "id": "822dae2b-ddd9-48ac-9356-e7b63a06d190", "metadata": {}, "outputs": [], "source": [ "from groq import Groq\n", "import os \n", "import json\n", "import sqlparse\n", "from IPython.display import Markdown\n", "import duckdb\n", "import glob\n", "import yaml" ] }, { "cell_type": "markdown", "id": "7f7c9c55-e925-4cc1-89f2-58237acf14a4", "metadata": {}, "source": [ "We will use the ```llama3-70b-8192``` model in this demo. Note that you will need a Groq API Key to proceed and can create an account [here](https://console.groq.com/) to generate one for free." ] }, { "cell_type": "code", "execution_count": 2, "id": "0cca781b-1950-4167-b36a-c1099d6b3b00", "metadata": {}, "outputs": [], "source": [ "client = Groq(api_key = os.getenv('GROQ_API_KEY'))\n", "model = 'llama3-70b-8192'" ] }, { "cell_type": "markdown", "id": "af6b1c5a-8940-4e9a-804c-6f52cb1e1f0e", "metadata": {}, "source": [ "### Text-To-SQL\n", "\n", "The first method is a standard **Text-To-SQL** implementation. With Text-To-SQL, we describe the database schema to the LLM, ask it to answer a question, and let it write an on-the-fly SQL query against the database to answer that question. Let's see how we can use the Groq API to build a Text-To-SQL pipeline:" ] }, { "cell_type": "markdown", "id": "e1f8d353-95cf-48d6-9a81-ac71a819abcd", "metadata": {}, "source": [ "First, we have our system prompt. A system prompt is an initial input or instruction given to the model, setting the context or specifying the task it needs to perform, essentially guiding the model's response generation. In our case, our system prompt will serve 3 purposes:\n", "\n", "1. Provide the metadata schemas for our database tables\n", "2. Indicate any relevant context or tips for querying the DuckDB language or our database schema specifically\n", "3. Define our desired JSON output (note that to use JSON mode, we must include 'JSON' in the prompt)" ] }, { "cell_type": "code", "execution_count": 3, "id": "a537a3c2-8fa3-49b4-9115-6ec5142d5630", "metadata": {}, "outputs": [], "source": [ "system_prompt = '''\n", "You are Groq Advisor, and you are tasked with generating SQL queries for DuckDB based on user questions about data stored in two tables derived from CSV files:\n", "\n", "Table: employees.csv\n", "Columns:\n", "employee_id (INTEGER): A unique identifier for each employee.\n", "name (VARCHAR): The full name of the employee.\n", "email (VARCHAR): employee's email address\n", "\n", "Table: purchases.csv\n", "Columns:\n", "purchase_id (INTEGER): A unique identifier for each purchase.\n", "purchase_date (DATE): Date of purchase\n", "employee_id (INTEGER): References the employee_id from the employees table, indicating which employee made the purchase.\n", "amount (FLOAT): The monetary value of the purchase.\n", "product_name (STRING): The name of the product purchased\n", "\n", "Given a user's question about this data, write a valid DuckDB SQL query that accurately extracts or calculates the requested information from these tables and adheres to SQL best practices for DuckDB, optimizing for readability and performance where applicable.\n", "\n", "Here are some tips for writing DuckDB queries:\n", "* DuckDB syntax requires querying from the .csv file itself, i.e. employees.csv and purchases.csv. For example: SELECT * FROM employees.csv as employees\n", "* All tables referenced MUST be aliased\n", "* DuckDB does not implicitly include a GROUP BY clause\n", "* CURRENT_DATE gets today's date\n", "* Aggregated fields like COUNT(*) must be appropriately named\n", "\n", "And some rules for querying the dataset:\n", "* Never include employee_id in the output - show employee name instead\n", "\n", "Also note that:\n", "* Valid values for product_name include 'Tesla','iPhone' and 'Humane pin'\n", "\n", "\n", "Question:\n", "--------\n", "{user_question}\n", "--------\n", "Reminder: Generate a DuckDB SQL to answer to the question:\n", "* respond as a valid JSON Document\n", "* [Best] If the question can be answered with the available tables: {\"sql\": }\n", "* If the question cannot be answered with the available tables: {\"error\": }\n", "* Ensure that the entire output is returned on only one single line\n", "* Keep your query as simple and straightforward as possible; do not use subqueries\n", "'''" ] }, { "cell_type": "markdown", "id": "d67ce3f8-b682-4e63-8a15-ebf8851bd676", "metadata": {}, "source": [ "Now we will define a ```text_to_sql``` function which takes in the system prompt and the user's question and outputs the LLM-generated DuckDB SQL query. Note that since we are using Groq API's [JSON mode](https://console.groq.com/docs/text-chat#json-mode-object-object) to format our output, we must indicate our expected JSON output format in either the system or user prompt." ] }, { "cell_type": "code", "execution_count": 4, "id": "705a0f56-828d-426c-8921-24913612f289", "metadata": {}, "outputs": [], "source": [ "def text_to_sql(client,system_prompt,user_question,model):\n", "\n", " completion = client.chat.completions.create(\n", " model = model,\n", " response_format = {\"type\": \"json_object\"},\n", " messages=[\n", " {\n", " \"role\": \"system\",\n", " \"content\": system_prompt\n", " },\n", " {\n", " \"role\": \"user\",\n", " \"content\": user_question\n", " }\n", " ]\n", " )\n", " \n", " return completion.choices[0].message.content" ] }, { "cell_type": "markdown", "id": "f9bdcb40-0587-4d70-ab29-1ff1c56450e4", "metadata": {}, "source": [ "...and a function for executing the DuckDB query that was generated:" ] }, { "cell_type": "code", "execution_count": 5, "id": "fd33a970-d7c8-4aba-a5e5-1017fc5d867e", "metadata": {}, "outputs": [], "source": [ "def execute_duckdb_query(query):\n", " original_cwd = os.getcwd()\n", " os.chdir('data')\n", " \n", " try:\n", " conn = duckdb.connect(database=':memory:', read_only=False)\n", " query_result = conn.execute(query).fetchdf().reset_index(drop=True)\n", " finally:\n", " os.chdir(original_cwd)\n", "\n", "\n", " return query_result" ] }, { "cell_type": "markdown", "id": "4c639e0d-8d09-47a1-ad55-5d8e118a6643", "metadata": {}, "source": [ "Now, we can query our database just by asking a question about the data. Here, the LLM generates a valid SQL query that reasonably answers the question:" ] }, { "cell_type": "code", "execution_count": 6, "id": "255be29d-d1b5-4a98-a2a3-f6ad513d7161", "metadata": {}, "outputs": [ { "data": { "text/markdown": [ "```sql\n", "SELECT e.name,\n", " p.purchase_date,\n", " p.product_name,\n", " p.amount\n", "FROM purchases.csv AS p\n", "JOIN employees.csv AS e ON p.employee_id = e.employee_id\n", "ORDER BY p.purchase_date DESC\n", "LIMIT 10;\n", "```" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
namepurchase_dateproduct_nameamount
0Jared Dunn2024-02-05Tesla75000
1Bertram Gilfoyle2024-02-04iPhone700
2Dinesh Chugtai2024-02-03Humane pin500
3Erlich Bachman2024-02-02Tesla70000
4Richard Hendricks2024-02-01iPhone750
\n", "
" ], "text/plain": [ " name purchase_date product_name amount\n", "0 Jared Dunn 2024-02-05 Tesla 75000\n", "1 Bertram Gilfoyle 2024-02-04 iPhone 700\n", "2 Dinesh Chugtai 2024-02-03 Humane pin 500\n", "3 Erlich Bachman 2024-02-02 Tesla 70000\n", "4 Richard Hendricks 2024-02-01 iPhone 750" ] }, "execution_count": 6, "metadata": {}, "output_type": "execute_result" } ], "source": [ "user_question = 'What are the most recent purchases?'\n", "\n", "\n", "llm_response = text_to_sql(client,system_prompt,user_question,model)\n", "sql_json = json.loads(llm_response)\n", "parsed_sql = sqlparse.format(sql_json['sql'], reindent=True, keyword_case='upper')\n", "formatted_sql = f\"```sql\\n{parsed_sql.strip()}\\n```\"\n", "display(Markdown(formatted_sql)) \n", "\n", "execute_duckdb_query(parsed_sql)" ] }, { "cell_type": "markdown", "id": "c633f9a3-00a5-4767-b3dc-da3b0686a4f2", "metadata": {}, "source": [ "Note, however, that due to the non-deterministic nature of LLMs, we cannot guarantee a reliable or consistent result every time. I might get a different result than you, and I might get a totally different query tomorrow. How should \"most recent purchases\" be defined? Which fields should be returned?\n", "\n", "Obviously, this is not ideal for making any kind of data-driven decisions. It's hard enough to land on a reliable source-of-truth data model, and even harder when your AI analyst cannot give you a consistent result. While text-to-SQL can be great for generating ad-hoc insights, the non-determinism feature of LLMs makes raw text-to-SQL an impractical solution for a production environment." ] }, { "cell_type": "markdown", "id": "cb3918e7-2f0b-4d95-8377-8742c3833a4f", "metadata": {}, "source": [ "### Function Calling for Verified Queries" ] }, { "cell_type": "markdown", "id": "eb224954-e17a-408c-954e-30757ade35f0", "metadata": {}, "source": [ "A different approach is to leverage the LLM to call on pre-vetted queries that can answer a set of questions. Since you wouldn't trust a traditional business intelligence tool without rigorously developed and validated SQL, a \"chat with your data\" app should be no different. For this example, we will use the verified queries stored [here](https://github.com/groq/groq-api-cookbook/tree/main/function-calling-sql/verified-queries)." ] }, { "cell_type": "code", "execution_count": 7, "id": "4ed69656-1c59-4180-a26d-b83c22944590", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "{'most-recent-purchases': {'description': 'Five most recent purchases',\n", " 'sql': 'SELECT \\n purchases.purchase_date,\\n purchases.product_name,\\n purchases.amount,\\n employees.name\\nFROM purchases.csv AS purchases\\nJOIN employees.csv AS employees ON purchases.employee_id = employees.employee_id\\nORDER BY purchases.purchase_date DESC\\nLIMIT 5;\\n'},\n", " 'most-expensive-purchase': {'description': 'Employee with the most expensive purchase',\n", " 'sql': 'SELECT employees.name AS employee_name,\\n MAX(amount) AS max_purchase_amount\\nFROM purchases.csv AS purchases\\nJOIN employees.csv AS employees ON purchases.employee_id = employees.employee_id\\nGROUP BY employees.name\\nORDER BY max_purchase_amount DESC\\nLIMIT 1\\n'},\n", " 'number-of-teslas': {'description': 'Number of Teslas purchased',\n", " 'sql': \"SELECT COUNT(*) as number_of_teslas\\nFROM purchases.csv AS p\\nJOIN employees.csv AS e ON e.employee_id = p.employee_id\\nWHERE p.product_name = 'Tesla'\\n\"},\n", " 'employees-without-purchases': {'description': 'Employees without a purchase since Feb 1, 2024',\n", " 'sql': \"SELECT employees.name as employees_without_purchases\\nFROM employees.csv AS employees\\nLEFT JOIN purchases.csv AS purchases ON employees.employee_id = purchases.employee_id\\nAND purchases.purchase_date > '2024-02-01'\\nWHERE purchases.purchase_id IS NULL\\n\"}}" ] }, "execution_count": 7, "metadata": {}, "output_type": "execute_result" } ], "source": [ "def get_verified_queries(directory_path):\n", " verified_queries_yaml_files = glob.glob(os.path.join(directory_path, '*.yaml'))\n", " verified_queries_dict = {}\n", " for file in verified_queries_yaml_files:\n", " with open(file, 'r') as stream:\n", " try:\n", " file_name = file[len(directory_path):-5]\n", " verified_queries_dict[file_name] = yaml.safe_load(stream)\n", " except yaml.YAMLError as exc:\n", " continue\n", " \n", " return verified_queries_dict\n", "\n", "directory_path = 'verified-queries/'\n", "verified_queries_dict = get_verified_queries(directory_path)\n", "verified_queries_dict" ] }, { "cell_type": "markdown", "id": "5ae230c2-220c-4ab7-83a2-2af799bd0a7b", "metadata": {}, "source": [ "Note that each of these queries are stored in ```yaml``` files with some additional metadata, like a description. This metadata is important for when the LLM needs to select the most appropriate query for the question at hand.\n", "\n", "Now, let's define a new function for executing SQL - this one is tweaked slightly to extract the SQL query from ```verified_queries_dict``` inside the function, given a query name:" ] }, { "cell_type": "code", "execution_count": 8, "id": "7e4b1e52-60bc-469b-8784-fedb8cb742d4", "metadata": {}, "outputs": [], "source": [ "def execute_duckdb_query_function_calling(query_name,verified_queries_dict):\n", " \n", " original_cwd = os.getcwd()\n", " os.chdir('data')\n", "\n", " query = verified_queries_dict[query_name]['sql']\n", " \n", " try:\n", " conn = duckdb.connect(database=':memory:', read_only=False)\n", " query_result = conn.execute(query).fetchdf().reset_index(drop=True)\n", " finally:\n", " os.chdir(original_cwd)\n", "\n", " return query_result" ] }, { "cell_type": "markdown", "id": "24a49b1d-7454-4955-b7b9-cb2480a68d9f", "metadata": {}, "source": [ "Finally, we will write a function to utilize Groq API's [Tool Use](https://console.groq.com/docs/tool-use) functionality to call the ```execute_duckdb_query_function_calling``` with the appropriate query name. We will provide the query/description mappings from ```verified_queries_dict``` in the system prompt so that the LLM can determine which query most appropriately answers the user's question:" ] }, { "cell_type": "code", "execution_count": 9, "id": "6b454910-4352-40cc-b9b2-cc79edabd7c1", "metadata": {}, "outputs": [], "source": [ "def call_verified_sql(user_question,verified_queries_dict,model):\n", " \n", " #Simplify verified_queries_dict to just show query name and description\n", " query_description_mapping = {key: subdict['description'] for key, subdict in verified_queries_dict.items()}\n", " \n", " # Step 1: send the conversation and available functions to the model\n", " messages=[\n", " {\n", " \"role\": \"system\",\n", " \"content\": '''You are a function calling LLM that uses the data extracted from the execute_duckdb_query_function_calling function to answer questions around a DuckDB dataset.\n", " \n", " Extract the query_name parameter from this mapping by finding the one whose description best matches the user's question: \n", " {query_description_mapping}\n", " '''.format(query_description_mapping=query_description_mapping)\n", " },\n", " {\n", " \"role\": \"user\",\n", " \"content\": user_question,\n", " }\n", " ]\n", " tools = [\n", " {\n", " \"type\": \"function\",\n", " \"function\": {\n", " \"name\": \"execute_duckdb_query_function_calling\",\n", " \"description\": \"Executes a verified DuckDB SQL Query\",\n", " \"parameters\": {\n", " \"type\": \"object\",\n", " \"properties\": {\n", " \"query_name\": {\n", " \"type\": \"string\",\n", " \"description\": \"The name of the verified query (i.e. 'most-recent-purchases')\",\n", " }\n", " },\n", " \"required\": [\"query_name\"],\n", " },\n", " },\n", " }\n", " ]\n", " response = client.chat.completions.create(\n", " model=model,\n", " messages=messages,\n", " tools=tools,\n", " tool_choice=\"auto\", \n", " max_tokens=4096\n", " )\n", " \n", " response_message = response.choices[0].message\n", " tool_calls = response_message.tool_calls\n", " \n", " available_functions = {\n", " \"execute_duckdb_query_function_calling\": execute_duckdb_query_function_calling,\n", " }\n", " for tool_call in tool_calls:\n", " function_name = tool_call.function.name\n", " function_to_call = available_functions[function_name]\n", " function_args = json.loads(tool_call.function.arguments)\n", " print('Query found: ',function_args.get(\"query_name\"))\n", " function_response = function_to_call(\n", " query_name=function_args.get(\"query_name\"),\n", " verified_queries_dict=verified_queries_dict\n", " )\n", " \n", " return function_response" ] }, { "cell_type": "markdown", "id": "19bd48ae-4091-4abb-ae65-ac89719e0146", "metadata": {}, "source": [ "Now, when we ask the LLM \"What were the most recent purchases?\", we will get the same logic every time:" ] }, { "cell_type": "code", "execution_count": 10, "id": "f828b583-291d-4aef-80b0-ce520b010ba5", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Query found: most-recent-purchases\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
purchase_dateproduct_nameamountname
02024-02-05Tesla75000Jared Dunn
12024-02-04iPhone700Bertram Gilfoyle
22024-02-03Humane pin500Dinesh Chugtai
32024-02-02Tesla70000Erlich Bachman
42024-02-01iPhone750Richard Hendricks
\n", "
" ], "text/plain": [ " purchase_date product_name amount name\n", "0 2024-02-05 Tesla 75000 Jared Dunn\n", "1 2024-02-04 iPhone 700 Bertram Gilfoyle\n", "2 2024-02-03 Humane pin 500 Dinesh Chugtai\n", "3 2024-02-02 Tesla 70000 Erlich Bachman\n", "4 2024-02-01 iPhone 750 Richard Hendricks" ] }, "execution_count": 10, "metadata": {}, "output_type": "execute_result" } ], "source": [ "user_prompt = 'What were the most recent purchases?'\n", "call_verified_sql(user_prompt,verified_queries_dict,model)" ] }, { "cell_type": "markdown", "id": "4b57679e-8828-46c7-9ae9-d420f0d794af", "metadata": {}, "source": [ "The downside with using verified queries, of course, is having to write and verify them, which takes away from the magic of watching an LLM generate a SQL query on the fly. But in an environment where reliability is critical, function calling against verified queries is a much more consistent way to chat with your data than Text-To-SQL." ] }, { "cell_type": "markdown", "id": "632b4242-9a51-4d04-9a11-b0f48bf9f177", "metadata": {}, "source": [ "This is a simple example, but you could even take it a step further by defining parameters for each query (that you might find in a WHERE clause), and doing another function call once the verified query is found to find the parameter(s) to inject in the query from the user prompt. Go ahead and try it out!" ] }, { "cell_type": "markdown", "id": "74c9774e-6ef2-4b90-9f62-f69bedef1e61", "metadata": {}, "source": [ "### Conclusion" ] }, { "cell_type": "markdown", "id": "5cfd0fb7-6606-40ac-9936-2fc5c880be9d", "metadata": {}, "source": [ "In this notebook we've explored two techniques for writing and executing SQL with LLMs using Groq API: Text-to-SQL (where the LLM generates SQL in the moment) and Verified Queries (where the LLM determines which verified query is most appropriate for your question and executes it). But perhaps the best approach is a blend - for ad-hoc reporting, there is still a lot of power in Text-to-SQL for quick answers. For user questions where there is no good verified query, you could default to using Text-To-SQL and then add that query to your dictionary of verified queries if it looks good. Either way, using LLMs on top of your data will lead to better and faster insights - just be sure to follow good data governance practices." ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.3" } }, "nbformat": 4, "nbformat_minor": 5 }