Format: application/octet-stream
1 {
2 "cells": [
3 {
4 "cell_type": "code",
5 "execution_count": 1,
6 "id": "a76b0ba4-1bc6-42f3-8f6c-1fb948bfef4b",
7 "metadata": {},
8 "outputs": [
9 {
10 "name": "stdout",
11 "output_type": "stream",
12 "text": [
13 "Defaulting to user installation because normal site-packages is not writeable\n",
14 "Collecting ollama\n",
15 " Downloading ollama-0.4.7-py3-none-any.whl.metadata (4.7 kB)\n",
16 "Requirement already satisfied: httpx<0.29,>=0.27 in /home/shubham/.local/lib/python3.8/site-packages (from ollama) (0.28.1)\n",
17 "Requirement already satisfied: pydantic<3.0.0,>=2.9.0 in /home/shubham/.local/lib/python3.8/site-packages (from ollama) (2.10.6)\n",
18 "Requirement already satisfied: anyio in /home/shubham/.local/lib/python3.8/site-packages (from httpx<0.29,>=0.27->ollama) (4.0.0)\n",
19 "Requirement already satisfied: certifi in /usr/lib/python3/dist-packages (from httpx<0.29,>=0.27->ollama) (2019.11.28)\n",
20 "Requirement already satisfied: httpcore==1.* in /home/shubham/.local/lib/python3.8/site-packages (from httpx<0.29,>=0.27->ollama) (1.0.8)\n",
21 "Requirement already satisfied: idna in /usr/lib/python3/dist-packages (from httpx<0.29,>=0.27->ollama) (2.8)\n",
22 "Requirement already satisfied: h11<0.15,>=0.13 in /home/shubham/.local/lib/python3.8/site-packages (from httpcore==1.*->httpx<0.29,>=0.27->ollama) (0.14.0)\n",
23 "Requirement already satisfied: annotated-types>=0.6.0 in /home/shubham/.local/lib/python3.8/site-packages (from pydantic<3.0.0,>=2.9.0->ollama) (0.7.0)\n",
24 "Requirement already satisfied: pydantic-core==2.27.2 in /home/shubham/.local/lib/python3.8/site-packages (from pydantic<3.0.0,>=2.9.0->ollama) (2.27.2)\n",
25 "Requirement already satisfied: typing-extensions>=4.12.2 in /home/shubham/.local/lib/python3.8/site-packages (from pydantic<3.0.0,>=2.9.0->ollama) (4.13.2)\n",
26 "Requirement already satisfied: sniffio>=1.1 in /home/shubham/.local/lib/python3.8/site-packages (from anyio->httpx<0.29,>=0.27->ollama) (1.3.0)\n",
27 "Requirement already satisfied: exceptiongroup>=1.0.2 in /home/shubham/.local/lib/python3.8/site-packages (from anyio->httpx<0.29,>=0.27->ollama) (1.1.3)\n",
28 "Downloading ollama-0.4.7-py3-none-any.whl (13 kB)\n",
29 "\u001b[33mDEPRECATION: distro-info 0.23ubuntu1 has a non-standard version number. pip 24.1 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of distro-info or contact the author to suggest that they release a version with a conforming version number. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
30 "\u001b[0m\u001b[33mDEPRECATION: python-debian 0.1.36ubuntu1 has a non-standard version number. pip 24.1 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of python-debian or contact the author to suggest that they release a version with a conforming version number. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
31 "\u001b[0mInstalling collected packages: ollama\n",
32 "Successfully installed ollama-0.4.7\n",
33 "\n",
34 "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.0\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.0.1\u001b[0m\n",
35 "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpython3 -m pip install --upgrade pip\u001b[0m\n"
36 ]
37 }
38 ],
39 "source": [
40 "!pip install ollama"
41 ]
42 },
43 {
44 "cell_type": "code",
45 "execution_count": 1,
46 "id": "842a38f7-00e1-4445-8e8b-c050da24829a",
47 "metadata": {},
48 "outputs": [],
49 "source": [
50 "import ollama\n",
51 "import json\n",
52 "from sqlalchemy import create_engine, MetaData\n",
53 "\n",
54 "# Example Schema Context\n",
55 "schema_context = \"\"\"\n",
56 "The database has the following tables:\n",
57 "- Orders (OrderID, CustomerID, TotalAmount, OrderDate, ProductID)\n",
58 "- Customers (CustomerID, FirstName, LastName)\n",
59 "- Products (ProductID, Name, Category, Price)\n",
60 "\"\"\"\n"
61 ]
62 },
63 {
64 "cell_type": "code",
65 "execution_count": 22,
66 "id": "7aca1a21-126d-48cd-93ee-06a6e92a07b9",
67 "metadata": {},
68 "outputs": [
69 {
70 "name": "stdout",
71 "output_type": "stream",
72 "text": [
73 "Generated SQL Query: model='deepseek-coder-v2' created_at='2025-04-15T06:56:23.993086105Z' done=True done_reason='stop' total_duration=3759103930 load_duration=17628895 prompt_eval_count=92 prompt_eval_duration=30000000 eval_count=367 eval_duration=3710000000 message=Message(role='assistant', content=\" To answer this question with a SQL query, we need to follow these steps:\\n1. Find the total order value for each customer and then find the customer with the highest total order value.\\n2. From that customer, find out which product they ordered most frequently (in terms of quantity).\\n\\nHere's how you can write this SQL query:\\n\\n```sql\\nSELECT p.Name \\nFROM Orders o\\nJOIN Customers c ON o.CustomerID = c.CustomerID\\nJOIN Products p ON o.ProductID = p.ProductID\\nWHERE c.CustomerID = (\\n SELECT CustomerID \\n FROM Orders \\n GROUP BY CustomerID \\n ORDER BY SUM(TotalAmount) DESC \\n LIMIT 1\\n)\\nGROUP BY p.Name \\nORDER BY COUNT(*) DESC \\nLIMIT 1;\\n```\\n\\nThis query does the following:\\n- The subquery in the `WHERE` clause finds the customer with the highest total order value by grouping orders by `CustomerID`, summing up their `TotalAmount`, and ordering them in descending order. It then limits the result to only the top row (the one with the highest total amount).\\n- From that specific customer, we join the Orders table with the Products table on `ProductID` to get the product names (`p.Name`) of all orders made by this customer.\\n- We group these results by `ProductID` and order them in descending order based on their count (i.e., which product was ordered most frequently).\\n- Finally, we limit our result to only the top row (the one with the highest frequency of orders), giving us the name of the product that this customer has ordered the most from among all products they've ordered.\", images=None, tool_calls=None)\n"
74 ]
75 }
76 ],
77 "source": [
78 "# Define the function to query Ollama with the schema context\n",
79 "def query_ollama_model(query):\n",
80 " # Prepare the full prompt with schema context\n",
81 " prompt = f\"{schema_context} User query: {query} {'only output sql query'} Response:\"\n",
82 " \n",
83 " # Use Ollama's API to generate the SQL query\n",
84 " response = ollama.chat(model=\"deepseek-coder-v2\", messages=[{\"role\": \"data engineer\", \"content\": prompt}], options ={'temperature':0.0})\n",
85 " response = ollama.chat(model=\"deepseek-coder-v2\", messages=[{\"role\": \"sql query extractor\", \"content\": prompt}], options ={'temperature':0.0})\n",
86 " \n",
87 " # Extract the generated SQL query from the response\n",
88 " sql_query = response\n",
89 " return sql_query\n",
90 "\n",
91 "# Test the function with a sample query\n",
92 "query = \"Which product is order the most by the customer who ordered the most [total order value is maximum]?\"\n",
93 "sql_response = query_ollama_model(query)\n",
94 "print(\"Generated SQL Query: \", sql_response)\n"
95 ]
96 },
97 {
98 "cell_type": "code",
99 "execution_count": 23,
100 "id": "5e8e0438-866c-435a-90cb-d05b218159f5",
101 "metadata": {},
102 "outputs": [
103 {
104 "data": {
105 "text/plain": [
106 "ollama._types.ChatResponse"
107 ]
108 },
109 "execution_count": 23,
110 "metadata": {},
111 "output_type": "execute_result"
112 }
113 ],
114 "source": [
115 "type(sql_response)"
116 ]
117 },
118 {
119 "cell_type": "code",
120 "execution_count": 24,
121 "id": "6f63f65b-a8e1-4fc1-b9fd-02b937463cb1",
122 "metadata": {},
123 "outputs": [
124 {
125 "name": "stdout",
126 "output_type": "stream",
127 "text": [
128 " To answer this question with a SQL query, we need to follow these steps:\n",
129 "1. Find the total order value for each customer and then find the customer with the highest total order value.\n",
130 "2. From that customer, find out which product they ordered most frequently (in terms of quantity).\n",
131 "\n",
132 "Here's how you can write this SQL query:\n",
133 "\n",
134 "```sql\n",
135 "SELECT p.Name \n",
136 "FROM Orders o\n",
137 "JOIN Customers c ON o.CustomerID = c.CustomerID\n",
138 "JOIN Products p ON o.ProductID = p.ProductID\n",
139 "WHERE c.CustomerID = (\n",
140 " SELECT CustomerID \n",
141 " FROM Orders \n",
142 " GROUP BY CustomerID \n",
143 " ORDER BY SUM(TotalAmount) DESC \n",
144 " LIMIT 1\n",
145 ")\n",
146 "GROUP BY p.Name \n",
147 "ORDER BY COUNT(*) DESC \n",
148 "LIMIT 1;\n",
149 "```\n",
150 "\n",
151 "This query does the following:\n",
152 "- The subquery in the `WHERE` clause finds the customer with the highest total order value by grouping orders by `CustomerID`, summing up their `TotalAmount`, and ordering them in descending order. It then limits the result to only the top row (the one with the highest total amount).\n",
153 "- From that specific customer, we join the Orders table with the Products table on `ProductID` to get the product names (`p.Name`) of all orders made by this customer.\n",
154 "- We group these results by `ProductID` and order them in descending order based on their count (i.e., which product was ordered most frequently).\n",
155 "- Finally, we limit our result to only the top row (the one with the highest frequency of orders), giving us the name of the product that this customer has ordered the most from among all products they've ordered.\n"
156 ]
157 }
158 ],
159 "source": [
160 "print(sql_response.message.content)"
161 ]
162 },
163 {
164 "cell_type": "code",
165 "execution_count": null,
166 "id": "ddc3b0f5-9d77-4db6-9582-52f0052c85d1",
167 "metadata": {},
168 "outputs": [],
169 "source": []
170 },
171 {
172 "cell_type": "code",
173 "execution_count": null,
174 "id": "69665fc3-014e-4452-a3f8-2da4ab805e7e",
175 "metadata": {},
176 "outputs": [],
177 "source": []
178 },
179 {
180 "cell_type": "markdown",
181 "id": "08601acb-00b9-449f-885f-16dc538f0b33",
182 "metadata": {},
183 "source": [
184 "VER2"
185 ]
186 },
187 {
188 "cell_type": "code",
189 "execution_count": 3,
190 "id": "17b2db33-8567-46db-bf87-bdaa7379cf95",
191 "metadata": {},
192 "outputs": [
193 {
194 "data": {
195 "text/plain": [
196 "'2025-04-15 08:44:21.456584'"
197 ]
198 },
199 "execution_count": 3,
200 "metadata": {},
201 "output_type": "execute_result"
202 }
203 ],
204 "source": [
205 "str(datetime.now())"
206 ]
207 },
208 {
209 "cell_type": "code",
210 "execution_count": 13,
211 "id": "aff42fc3-ba58-49a5-9424-e6261e653705",
212 "metadata": {},
213 "outputs": [],
214 "source": [
215 "import re\n",
216 "import ollama\n",
217 "from pydantic import BaseModel\n",
218 "from typing import Optional\n",
219 "from datetime import datetime\n",
220 "\n",
221 "\n",
222 "\n",
223 "\n",
224 "\n",
225 "schema_context = f\"\"\"\n",
226 "The database has the following tables:\n",
227 "- Orders (OrderID, CustomerID, TotalAmount, OrderDate, ProductID)\n",
228 "- Customers (CustomerID, FirstName, LastName)\n",
229 "- Products (ProductID, Name, Category, Price)\n",
230 "\n",
231 "You must output the response strictly in JSON format with the following keys:\n",
232 "{{\n",
233 " \"sql_query\": \"<SQL query string [Properly Formated]>\",\n",
234 " \"query_planning\": \"<Thought process, here you think before generating the query, how did plan the query.> [Optional]\",\n",
235 " \"explaination\": \"<Explain your query briefly> [Optional]\"\n",
236 "}}\n",
237 "\n",
238 "Below are some examples:\n",
239 "\n",
240 "Example 1:\n",
241 "{{\n",
242 " \"sql_query\": \"SELECT o.OrderID, c.FirstName, c.LastName, o.TotalAmount, o.OrderDate\\\\nFROM Orders AS o\\\\nJOIN Customers AS c ON o.CustomerID = c.CustomerID\\\\nWHERE c.FirstName = 'John' AND c.LastName = 'Doe';\",\n",
243 " \"query_planning\": \"Plan: Start by joining the Orders table with the Customers table using the CustomerID to link orders with the corresponding customer records. Next, filter the result set by the customer's first and last name to isolate the orders for 'John Doe'. Finally, select the relevant columns (OrderID, TotalAmount, OrderDate) along with the customer's first and last name.\",\n",
244 " \"explaination\": \"Retrieves order details (OrderID, TotalAmount, OrderDate) for the customer named John Doe.\"\n",
245 "}}\n",
246 "\n",
247 "Example 2:\n",
248 "{{\n",
249 " \"sql_query\": \"SELECT o.OrderID, p.Name AS ProductName, o.TotalAmount, o.OrderDate\\\\nFROM Orders AS o\\\\nJOIN Products AS p ON o.ProductID = p.ProductID\\\\nWHERE p.Category = 'Electronics';\",\n",
250 " \"query_planning\": \"Plan: Begin by joining the Orders table with the Products table on the ProductID field to associate each order with its product details. Then, filter the orders where the product category is 'Electronics'. Finally, select the order ID, product name, total amount, and order date to generate a complete view of electronics orders.\",\n",
251 " \"explaination\": \"Lists order details and product names for orders that include products from the 'Electronics' category.\"\n",
252 "}}\n",
253 "\n",
254 "Example 3:\n",
255 "{{\n",
256 " \"sql_query\": \"SELECT p.Category, SUM(o.TotalAmount) AS TotalSales\\\\nFROM Orders AS o\\\\nJOIN Products AS p ON o.ProductID = p.ProductID\\\\nGROUP BY p.Category;\",\n",
257 " \"query_planning\": \"Plan: Join the Orders table with the Products table using ProductID in order to map each order to its corresponding product. Then, group the results by product category. Finally, calculate the total sales for each category by summing up the TotalAmount for the grouped orders.\",\n",
258 " \"explaination\": \"Calculates and aggregates the total sales for each product category.\"\n",
259 "}}\n",
260 "\n",
261 "Example 4:\n",
262 "{{\n",
263 " \"sql_query\": \"SELECT DISTINCT c.CustomerID, c.FirstName, c.LastName\\\\nFROM Customers AS c\\\\nJOIN Orders AS o ON c.CustomerID = o.CustomerID\\\\nWHERE o.OrderDate >= CURRENT_DATE - INTERVAL '30' DAY;\",\n",
264 " \"query_planning\": \"Plan: Begin by joining the Customers table with the Orders table on the CustomerID to connect customer information with their orders. Next, filter out orders that were placed in the last 30 days by applying a date condition on the OrderDate. Use DISTINCT to avoid duplicate customer entries in case they placed multiple orders in that period.\",\n",
265 " \"explaination\": \"Retrieves a unique list of customers who have placed orders within the last 30 days.\"\n",
266 "}}\n",
267 "\n",
268 "Current Date and Time: {str(datetime.now())}\n",
269 "\"\"\"\n",
270 "\n",
271 "\n",
272 "\n",
273 "# Define a Pydantic model for the structured SQL query response.\n",
274 "class SQLQueryResponse(BaseModel):\n",
275 " sql_query: str\n",
276 " query_planning: Optional[str] = None\n",
277 " explaination: Optional[str] = None\n",
278 "\n",
279 "def extract_json_from_markdown(text: str) -> str:\n",
280 " \"\"\"\n",
281 " Extracts the JSON string from a markdown formatted code block, if present.\n",
282 " If not found, returns the stripped original text.\n",
283 " \"\"\"\n",
284 " import re\n",
285 " pattern = r\"```json\\s*(\\{.*\\})\\s*```\"\n",
286 " match = re.search(pattern, text, re.DOTALL)\n",
287 " if match:\n",
288 " return match.group(1)\n",
289 " return text.strip()\n",
290 "\n",
291 "# Function to query Ollama with the schema context and return a structured response.\n",
292 "def query_ollama_model(query: str) -> SQLQueryResponse:\n",
293 " # Prepare the full prompt with the modified schema context.\n",
294 " prompt = f\"{schema_context}\\nUser query: {query}\\nResponse (provide output in JSON format):\"\n",
295 " \n",
296 " # Use Ollama's API to generate the SQL query.\n",
297 " response = ollama.chat(\n",
298 " model=\"deepseek-coder-v2\",\n",
299 " messages=[{\"role\": \"user\", \"content\": prompt}],\n",
300 " options={'temperature': 0.0}\n",
301 " )\n",
302 " \n",
303 " # Extract the JSON from any markdown code fences.\n",
304 " json_str = extract_json_from_markdown(response.message.content)\n",
305 " \n",
306 " # Use the Pydantic V2 method model_validate_json to parse the JSON.\n",
307 " try:\n",
308 " structured_response = SQLQueryResponse.model_validate_json(json_str)\n",
309 " except Exception as e:\n",
310 " raise ValueError(f\"Failed to parse JSON response: {e}\")\n",
311 " \n",
312 " return structured_response\n",
313 "\n"
314 ]
315 },
316 {
317 "cell_type": "code",
318 "execution_count": 14,
319 "id": "5fc0b803-df12-4e2f-a83c-8aae87c980b0",
320 "metadata": {},
321 "outputs": [
322 {
323 "name": "stdout",
324 "output_type": "stream",
325 "text": [
326 "Generated SQL Query:\n",
327 "WITH MaxOrderValue AS (\n",
328 " SELECT c.CustomerID, SUM(o.TotalAmount) AS TotalOrderValue\n",
329 " FROM Orders AS o\n",
330 " JOIN Customers AS c ON o.CustomerID = c.CustomerID\n",
331 " WHERE o.OrderDate >= CURRENT_DATE - INTERVAL '10 weeks'\n",
332 " GROUP BY c.CustomerID\n",
333 "), \n",
334 "MostValuableCustomer AS (\n",
335 " SELECT CustomerID\n",
336 " FROM MaxOrderValue\n",
337 " ORDER BY TotalOrderValue DESC\n",
338 " LIMIT 1\n",
339 ")SELECT p.ProductID, p.Name, COUNT(o.OrderID) AS OrderCount\n",
340 "FROM Orders AS o\n",
341 "JOIN Products AS p ON o.ProductID = p.ProductID\n",
342 "WHERE o.CustomerID IN (SELECT CustomerID FROM MostValuableCustomer)\n",
343 "GROUP BY p.ProductID, p.Name\n",
344 "ORDER BY OrderCount DESC;\n",
345 "\n",
346 "Thinking Process:\n",
347 "Plan: First, we use a Common Table Expression (CTE) to find the customer with the maximum total order value for orders placed in the last 10 weeks. Then, using this result, we identify the most valuable customer's ID. Finally, we join the Orders table with the Products table on ProductID and filter by the identified CustomerID to count the number of times each product was ordered by the most valuable customer.\n",
348 "\n",
349 "Explanation:\n",
350 "This query identifies the most valuable customer based on total order value for orders placed in the last 10 weeks, then finds which products were ordered most frequently by this customer.\n",
351 "\n",
352 "Structured Response (JSON):\n",
353 "{\"sql_query\":\"WITH MaxOrderValue AS (\\n SELECT c.CustomerID, SUM(o.TotalAmount) AS TotalOrderValue\\n FROM Orders AS o\\n JOIN Customers AS c ON o.CustomerID = c.CustomerID\\n WHERE o.OrderDate >= CURRENT_DATE - INTERVAL '10 weeks'\\n GROUP BY c.CustomerID\\n), \\nMostValuableCustomer AS (\\n SELECT CustomerID\\n FROM MaxOrderValue\\n ORDER BY TotalOrderValue DESC\\n LIMIT 1\\n)SELECT p.ProductID, p.Name, COUNT(o.OrderID) AS OrderCount\\nFROM Orders AS o\\nJOIN Products AS p ON o.ProductID = p.ProductID\\nWHERE o.CustomerID IN (SELECT CustomerID FROM MostValuableCustomer)\\nGROUP BY p.ProductID, p.Name\\nORDER BY OrderCount DESC;\",\"query_planning\":\"Plan: First, we use a Common Table Expression (CTE) to find the customer with the maximum total order value for orders placed in the last 10 weeks. Then, using this result, we identify the most valuable customer's ID. Finally, we join the Orders table with the Products table on ProductID and filter by the identified CustomerID to count the number of times each product was ordered by the most valuable customer.\",\"explaination\":\"This query identifies the most valuable customer based on total order value for orders placed in the last 10 weeks, then finds which products were ordered most frequently by this customer.\"}\n"
354 ]
355 },
356 {
357 "name": "stderr",
358 "output_type": "stream",
359 "text": [
360 "<ipython-input-14-7a163988cecc>:16: PydanticDeprecatedSince20: The `json` method is deprecated; use `model_dump_json` instead. Deprecated in Pydantic V2.0 to be removed in V3.0. See Pydantic V2 Migration Guide at https://errors.pydantic.dev/2.10/migration/\n",
361 " print(response.json())\n"
362 ]
363 }
364 ],
365 "source": [
366 "\n",
367 "# Test the function with a sample query.\n",
368 "if __name__ == '__main__':\n",
369 " query = \"Which products were ordered the most by the customer who ordered the most [total order value is maximum]for last 10 weeks? [use CTE's please]\"\n",
370 " response = query_ollama_model(query)\n",
371 " \n",
372 " print(\"Generated SQL Query:\")\n",
373 " print(response.sql_query)\n",
374 " \n",
375 " print(\"\\nThinking Process:\")\n",
376 " print(response.query_planning)\n",
377 " \n",
378 " print(\"\\nExplanation:\")\n",
379 " print(response.explaination)\n",
380 " \n",
381 " print(\"\\nStructured Response (JSON):\")\n",
382 " print(response.json())"
383 ]
384 },
385 {
386 "cell_type": "code",
387 "execution_count": null,
388 "id": "2f85a5ac-6808-46fb-bb66-7768ccec6828",
389 "metadata": {},
390 "outputs": [],
391 "source": []
392 },
393 {
394 "cell_type": "code",
395 "execution_count": null,
396 "id": "e57e1e59-776f-41e6-bc82-526ff7ab49c5",
397 "metadata": {},
398 "outputs": [],
399 "source": []
400 },
401 {
402 "cell_type": "code",
403 "execution_count": null,
404 "id": "6168d107-45cf-4f6c-a55e-1a713708497f",
405 "metadata": {},
406 "outputs": [],
407 "source": []
408 },
409 {
410 "cell_type": "code",
411 "execution_count": null,
412 "id": "52ade50e-658b-44b1-a965-30b6e5d3750b",
413 "metadata": {},
414 "outputs": [],
415 "source": []
416 },
417 {
418 "cell_type": "code",
419 "execution_count": null,
420 "id": "2daa1991-5f62-42e5-ada5-e4699d015fd8",
421 "metadata": {},
422 "outputs": [],
423 "source": []
424 },
425 {
426 "cell_type": "markdown",
427 "id": "3755ee48-e271-4134-b512-513d7eb8dd51",
428 "metadata": {},
429 "source": [
430 "###VERSION 3"
431 ]
432 },
433 {
434 "cell_type": "code",
435 "execution_count": 29,
436 "id": "a09bedda-760e-493f-995c-015cdca70e06",
437 "metadata": {},
438 "outputs": [
439 {
440 "name": "stdout",
441 "output_type": "stream",
442 "text": [
443 "Response 1 Structured JSON:\n",
444 "Plan: First, identify the customer who has placed the highest total order value by joining the Orders table with itself on CustomerID and ordering by SUM(TotalAmount) in descending order, then limiting to the top result. Next, use this identified customer ID to filter orders for that specific customer. Finally, join these filtered orders with the Products table to get product names and sum their TotalAmounts to find out which product is ordered most. \n",
445 " SELECT p.Name AS ProductName, SUM(o.TotalAmount) AS TotalOrdered\n",
446 "FROM Orders AS o\n",
447 "JOIN Products AS p ON o.ProductID = p.ProductID\n",
448 "WHERE o.CustomerID IN (\n",
449 " SELECT CustomerID\n",
450 " FROM (SELECT CustomerID, SUM(TotalAmount) AS TotalOrderValue\n",
451 " FROM Orders\n",
452 " GROUP BY CustomerID\n",
453 " ORDER BY TotalOrderValue DESC\n",
454 " LIMIT 1)\n",
455 ") \n",
456 "GROUP BY p.Name\n",
457 "ORDER BY TotalOrdered DESC\n",
458 "LIMIT 1;\n",
459 "\n",
460 "Response 2 Structured JSON:\n",
461 "Plan: First, identify the customer who has placed the highest total order value by joining the Orders table with itself on CustomerID and ordering by SUM(TotalAmount) in descending order, then limiting to the top result. Next, use this identified customer ID to filter orders for that specific customer. Finally, join these filtered orders with the Products table to get product names and sum their TotalAmounts to find out which product is ordered most. \n",
462 " SELECT p.Name AS ProductName, SUM(o.TotalAmount) AS TotalOrdered\n",
463 "FROM Orders AS o\n",
464 "JOIN Products AS p ON o.ProductID = p.ProductID\n",
465 "WHERE o.CustomerID IN (\n",
466 " SELECT CustomerID\n",
467 " FROM (SELECT CustomerID, SUM(TotalAmount) AS TotalOrderValue\n",
468 " FROM Orders\n",
469 " GROUP BY CustomerID\n",
470 " ORDER BY TotalOrderValue DESC\n",
471 " LIMIT 1)\n",
472 ") \n",
473 "GROUP BY p.Name\n",
474 "ORDER BY TotalOrdered DESC\n",
475 "LIMIT 1;\n"
476 ]
477 }
478 ],
479 "source": [
480 "import re\n",
481 "import ollama\n",
482 "from pydantic import BaseModel\n",
483 "from typing import Optional, List\n",
484 "from datetime import datetime\n",
485 "\n",
486 "# Updated system context with role definition, detailed instructions, and conversation-style examples.\n",
487 "schema_context = f\"\"\"\n",
488 "The database has the following tables:\n",
489 "- Orders (OrderID, CustomerID, TotalAmount, OrderDate, ProductID)\n",
490 "- Customers (CustomerID, FirstName, LastName)\n",
491 "- Products (ProductID, Name, Category, Price)\n",
492 "\n",
493 "You are a Query Generator / Database Architect. You are highly experienced in crafting efficient and optimized SQL queries. When responding, always output your answer strictly in JSON format with the following keys:\n",
494 "{{\n",
495 " \"sql_query\": \"<SQL query string [Properly Formated with newlines and indentation for clarity]>\",\n",
496 " \"query_planning\": \"<A detailed, step-by-step explanation of your approach in planning the query. Include which tables are joined, what filters are applied, and why any aggregations or grouping are used.> [Optional]\",\n",
497 " \"explaination\": \"<A concise summary of what your query does and why it is optimal.> [Optional]\"\n",
498 "}}\n",
499 "\n",
500 "Your responses must be crafted in a professional tone, using technical language and best practices in SQL design.\n",
501 "\n",
502 "\n",
503 "Below are some conversation examples in a user/assistant format:\n",
504 "\n",
505 "Example 1:\n",
506 "User:\n",
507 " \"Generate a query to retrieve order details for a customer named John Doe.\"\n",
508 "Assistant:\n",
509 " {{\n",
510 " \"sql_query\": \"SELECT o.OrderID, c.FirstName, c.LastName, o.TotalAmount, o.OrderDate\\\\nFROM Orders AS o\\\\nJOIN Customers AS c ON o.CustomerID = c.CustomerID\\\\nWHERE c.FirstName = 'John' AND c.LastName = 'Doe';\",\n",
511 " \"query_planning\": \"Plan: Join the Orders table with the Customers table using CustomerID to link each order to its customer. Filter the results by checking that FirstName is 'John' and LastName is 'Doe'. Finally, select columns that provide essential order details.\",\n",
512 " \"explaination\": \"Retrieves detailed order information for customer John Doe.\"\n",
513 " }}\n",
514 "\n",
515 "Example 2:\n",
516 "User:\n",
517 " \"Generate a query that lists the orders with product details for all orders in the 'Electronics' category.\"\n",
518 "Assistant:\n",
519 " {{\n",
520 " \"sql_query\": \"SELECT o.OrderID, p.Name AS ProductName, o.TotalAmount, o.OrderDate\\\\nFROM Orders AS o\\\\nJOIN Products AS p ON o.ProductID = p.ProductID\\\\nWHERE p.Category = 'Electronics';\",\n",
521 " \"query_planning\": \"Plan: Join Orders with Products on ProductID to associate each order with its respective product. Filter the rows where the product category is 'Electronics'. Select order ID, product name, total amount, and order date for complete detail.\",\n",
522 " \"explaination\": \"Lists order information along with the product details for orders in the 'Electronics' category.\"\n",
523 " }}\n",
524 "\n",
525 "Example 3:\n",
526 "User:\n",
527 " \"Generate a query to calculate total sales per product category.\"\n",
528 "Assistant:\n",
529 " {{\n",
530 " \"sql_query\": \"SELECT p.Category, SUM(o.TotalAmount) AS TotalSales\\\\nFROM Orders AS o\\\\nJOIN Products AS p ON o.ProductID = p.ProductID\\\\nGROUP BY p.Category;\",\n",
531 " \"query_planning\": \"Plan: Join Orders with Products using the ProductID to map each order to its product. Group the orders by product category and sum the TotalAmount for each group to compute the total sales.\",\n",
532 " \"explaination\": \"Aggregates total sales by product category to provide revenue insights across product types.\"\n",
533 " }}\n",
534 "\n",
535 "Example 4:\n",
536 "User:\n",
537 " \"Generate a query to list unique customers who have placed orders within the last 30 days.\"\n",
538 "Assistant:\n",
539 " {{\n",
540 " \"sql_query\": \"SELECT DISTINCT c.CustomerID, c.FirstName, c.LastName\\\\nFROM Customers AS c\\\\nJOIN Orders AS o ON c.CustomerID = o.CustomerID\\\\nWHERE o.OrderDate >= CURRENT_DATE - INTERVAL '30' DAY;\",\n",
541 " \"query_planning\": \"Plan: Join the Customers table with the Orders table on CustomerID to combine customer details with their orders. Apply a filter on OrderDate to include only those orders placed in the last 30 days, and use DISTINCT to avoid duplicate customer entries.\",\n",
542 " \"explaination\": \"Retrieves a unique set of customers who placed orders within the last 30 days.\"\n",
543 " }}\n",
544 "\n",
545 "Current Date and Time: {str(datetime.now())}\n",
546 "\"\"\"\n",
547 "\n",
548 "# Define a Pydantic model for the structured SQL query response.\n",
549 "class SQLQueryResponse(BaseModel):\n",
550 " sql_query: str\n",
551 " query_planning: Optional[str] = None\n",
552 " explaination: Optional[str] = None\n",
553 "\n",
554 "def extract_json_from_markdown(text: str) -> str:\n",
555 " \"\"\"\n",
556 " Extracts the JSON string from a markdown formatted code block if present.\n",
557 " If not found, returns the stripped original text.\n",
558 " \"\"\"\n",
559 " pattern = r\"```json\\s*(\\{.*\\})\\s*```\"\n",
560 " match = re.search(pattern, text, re.DOTALL)\n",
561 " if match:\n",
562 " return match.group(1)\n",
563 " return text.strip()\n",
564 "\n",
565 "def build_messages(history: List[dict], new_query: str) -> List[dict]:\n",
566 " \"\"\"\n",
567 " Constructs the conversation messages list including:\n",
568 " - A system message with the schema context (only once)\n",
569 " - The conversation history (user and assistant messages)\n",
570 " - The new user query.\n",
571 " \"\"\"\n",
572 " messages = [{\"role\": \"system\", \"content\": schema_context}]\n",
573 " messages.extend(history)\n",
574 " messages.append({\"role\": \"user\", \"content\": f\"User query: {new_query}\\nResponse (provide output in JSON format):\"})\n",
575 " return messages\n",
576 "\n",
577 "def query_ollama_model_with_history(history: List[dict], query: str) -> SQLQueryResponse:\n",
578 " \"\"\"\n",
579 " Uses the conversation history along with the new query to build a prompt,\n",
580 " calls the Ollama API, extracts and validates the JSON response, and updates\n",
581 " the conversation history with the new messages.\n",
582 " \"\"\"\n",
583 " messages = build_messages(history, query)\n",
584 " response = ollama.chat(\n",
585 " model=\"deepseek-coder-v2\",\n",
586 " messages=messages,\n",
587 " options={'temperature': 0.0}\n",
588 " )\n",
589 " \n",
590 " json_str = extract_json_from_markdown(response.message.content)\n",
591 " \n",
592 " try:\n",
593 " structured_response = SQLQueryResponse.model_validate_json(json_str)\n",
594 " except Exception as e:\n",
595 " raise ValueError(f\"Failed to parse JSON response: {e}\")\n",
596 " \n",
597 " # Append the current turn to the conversation history.\n",
598 " history.append({\"role\": \"user\", \"content\": query})\n",
599 " history.append({\"role\": \"assistant\", \"content\": response.message.content})\n",
600 " return structured_response\n",
601 "\n",
602 "# Example usage:\n",
603 "if __name__ == '__main__':\n",
604 " # Initialize persistent conversation history.\n",
605 " conversation_history = []\n",
606 " \n",
607 " # First query.\n",
608 " query1 = \"Which product is ordered the most by the customer who ordered the most (based on total order value)?\"\n",
609 " response1 = query_ollama_model_with_history(conversation_history, query1)\n",
610 " print(\"Response 1 Structured JSON:\")\n",
611 " print(response1.query_planning,'\\n',response1.sql_query)\n",
612 " \n",
613 " # Second query within the same session.\n",
614 " query2 = \"What was your last output? in previous message?\"\n",
615 " response2 = query_ollama_model_with_history(conversation_history, query2)\n",
616 " print(\"\\nResponse 2 Structured JSON:\")\n",
617 " print(response2.query_planning,'\\n',response2.sql_query)\n"
618 ]
619 },
620 {
621 "cell_type": "code",
622 "execution_count": 25,
623 "id": "305c952e-2b6b-465e-8ccc-266fcea5ed41",
624 "metadata": {},
625 "outputs": [
626 {
627 "data": {
628 "text/plain": [
629 "SQLQueryResponse(sql_query='SELECT p.Name AS ProductName\\nFROM Orders o\\nJOIN Products p ON o.ProductID = p.ProductID\\nWHERE o.CustomerID IN (\\n SELECT CustomerID\\n FROM (SELECT CustomerID, SUM(TotalAmount) AS TotalOrderValue\\n FROM Orders\\n GROUP BY CustomerID\\n ORDER BY TotalOrderValue DESC\\n LIMIT 1)\\n) ORDER BY o.TotalAmount DESC LIMIT 1;', query_planning='Plan: First, identify the customer who has placed the highest total order value by joining the Orders table with itself to sum up the TotalAmount for each CustomerID and ordering them in descending order. Limit this selection to only the top result (the customer with the highest total order value). Next, within this subset of orders belonging to this single customer, find the product that has been ordered most frequently based on its TotalAmount. This is achieved by joining the Orders table with the Products table and ordering the results by TotalAmount in descending order, then limiting the output to just 1 record.', explaination='This query identifies the top-spending customer and then finds the product they have ordered most frequently based on total order value.')"
630 ]
631 },
632 "execution_count": 25,
633 "metadata": {},
634 "output_type": "execute_result"
635 }
636 ],
637 "source": [
638 "response1"
639 ]
640 },
641 {
642 "cell_type": "code",
643 "execution_count": 1,
644 "id": "fb13bf51-2f1d-439c-8118-8330f0ee7e9a",
645 "metadata": {},
646 "outputs": [
647 {
648 "name": "stdout",
649 "output_type": "stream",
650 "text": [
651 "model='hf.co/openerotica/writing-roleplay-20k-context-nemo-12b-v1.0-gguf:F16' modified_at=datetime.datetime(2025, 4, 17, 21, 20, 42, 976253, tzinfo=TzInfo(UTC)) digest='0f9f279647d1e5144473835c8354a9ae5d1adac46fb420907b3466885df39e5c' size=24504317245 details=ModelDetails(parent_model='', format='gguf', family='llama', families=['llama'], parameter_size='12.2B', quantization_level='unknown')\n",
652 "model='cogito:32b' modified_at=datetime.datetime(2025, 4, 17, 10, 15, 42, 710153, tzinfo=TzInfo(UTC)) digest='0b4aab772f57ddf3d9e46235ef6f142e42598fd121dc15588c7f06de786c510e' size=19848516836 details=ModelDetails(parent_model='', format='gguf', family='qwen2', families=['qwen2'], parameter_size='32.8B', quantization_level='Q4_K_M')\n",
653 "model='cogito:latest' modified_at=datetime.datetime(2025, 4, 15, 21, 39, 31, 7685, tzinfo=TzInfo(UTC)) digest='75b508ddece134205c621b91237310024ca983027fb7f724c2032d106419bd1b' size=4920747948 details=ModelDetails(parent_model='', format='gguf', family='llama', families=['llama'], parameter_size='8.0B', quantization_level='Q4_K_M')\n",
654 "model='artifish/llama3.2-uncensored:latest' modified_at=datetime.datetime(2025, 4, 15, 20, 28, 19, 718626, tzinfo=TzInfo(UTC)) digest='c73bea26e00483b58ba59729b74bfc1647971dafa1fe66c890d845ae63e2c34e' size=2241006824 details=ModelDetails(parent_model='', format='gguf', family='llama', families=['llama'], parameter_size='3.6B', quantization_level='Q4_K_M')\n",
655 "model='deepseek-coder-v2:latest' modified_at=datetime.datetime(2025, 4, 15, 6, 53, 32, 916369, tzinfo=TzInfo(UTC)) digest='63fb193b3a9b4322a18e8c6b250ca2e70a5ff531e962dbf95ba089b2566f2fa5' size=8905126121 details=ModelDetails(parent_model='', format='gguf', family='deepseek2', families=['deepseek2'], parameter_size='15.7B', quantization_level='Q4_0')\n",
656 "model='deepseek-v3:671b' modified_at=datetime.datetime(2025, 4, 15, 6, 7, 47, 388953, tzinfo=TzInfo(UTC)) digest='5da0e2d4a9e0e3cc2cac1e952a6808f1eccdc75acd9cf4387e41c8ccc77a929a' size=404430200769 details=ModelDetails(parent_model='', format='gguf', family='deepseek2', families=['deepseek2'], parameter_size='671.0B', quantization_level='Q4_K_M')\n",
657 "model='codellama:latest' modified_at=datetime.datetime(2025, 4, 15, 4, 47, 50, 580852, tzinfo=TzInfo(UTC)) digest='8fdf8f752f6e80de33e82f381aba784c025982752cd1ae9377add66449d2225f' size=3825910662 details=ModelDetails(parent_model='', format='gguf', family='llama', families=None, parameter_size='7B', quantization_level='Q4_0')\n",
658 "model='nomic-embed-text:latest' modified_at=datetime.datetime(2025, 1, 3, 9, 49, 49, 322519, tzinfo=TzInfo(UTC)) digest='0a109f422b47e3a30ba2b10eca18548e944e8a23073ee3f3e947efcf3c45e59f' size=274302450 details=ModelDetails(parent_model='', format='gguf', family='nomic-bert', families=['nomic-bert'], parameter_size='137M', quantization_level='F16')\n",
659 "model='llama3.2:latest' modified_at=datetime.datetime(2024, 12, 28, 13, 55, 10, 79145, tzinfo=TzInfo(UTC)) digest='a80c4f17acd55265feec403c7aef86be0c25983ab279d83f3bcd3abbcb5b8b72' size=2019393189 details=ModelDetails(parent_model='', format='gguf', family='llama', families=['llama'], parameter_size='3.2B', quantization_level='Q4_K_M')\n",
660 "model='llama2:latest' modified_at=datetime.datetime(2024, 2, 29, 0, 21, 2, 69015, tzinfo=TzInfo(UTC)) digest='78e26419b4469263f75331927a00a0284ef6544c1975b826b15abdaef17bb962' size=3826793677 details=ModelDetails(parent_model='', format='gguf', family='llama', families=['llama'], parameter_size='7B', quantization_level='Q4_0')\n",
661 "model='mistral:latest' modified_at=datetime.datetime(2024, 2, 26, 23, 49, 35, 566793, tzinfo=TzInfo(UTC)) digest='61e88e884507ba5e06c49b40e6226884b2a16e872382c2b44a42f2d119d804a5' size=4109865159 details=ModelDetails(parent_model='', format='gguf', family='llama', families=['llama'], parameter_size='7B', quantization_level='Q4_0')\n"
662 ]
663 }
664 ],
665 "source": [
666 "import ollama\n",
667 "\n",
668 "# List all available models\n",
669 "models = ollama.list()\n",
670 "\n",
671 "# Print model names\n",
672 "for model in models['models']:\n",
673 " print(model)\n"
674 ]
675 },
676 {
677 "cell_type": "markdown",
678 "id": "2bc37000-21fb-4192-84a0-613d3f5b133a",
679 "metadata": {},
680 "source": [
681 "###VERS 4"
682 ]
683 },
684 {
685 "cell_type": "code",
686 "execution_count": null,
687 "id": "30724ed5-34c3-457b-84ee-5f179a3f1a2d",
688 "metadata": {},
689 "outputs": [],
690 "source": []
691 },
692 {
693 "cell_type": "code",
694 "execution_count": null,
695 "id": "e613795f-cc74-4cf9-a7c9-85873f044b3f",
696 "metadata": {},
697 "outputs": [],
698 "source": []
699 },
700 {
701 "cell_type": "code",
702 "execution_count": 39,
703 "id": "618ca44f-823a-4a1c-8423-c66b8bfa1974",
704 "metadata": {},
705 "outputs": [
706 {
707 "name": "stdout",
708 "output_type": "stream",
709 "text": [
710 "Generated SQL Query:\n",
711 "SELECT DATE_TRUNC('week', o.OrderDate) AS Week, p.Name AS ProductName, SUM(o.TotalAmount) AS WeeklySales\n",
712 "FROM Orders AS o\n",
713 "JOIN Products AS p ON o.ProductID = p.ProductID\n",
714 "WHERE o.OrderDate >= CURRENT_DATE - INTERVAL '3 months'\n",
715 "GROUP BY DATE_TRUNC('week', o.OrderDate), p.Name\n",
716 "ORDER BY Week, ProductName;\n",
717 "Query Planning:\n",
718 "Plan: Join Orders with Products on ProductID to get product names and sales data. Filter orders from the last quarter (3 months) by OrderDate. Group results by week of the year for each order date and aggregate TotalAmount to calculate weekly sales per product.\n",
719 "Explanation:\n",
720 "This query retrieves weekly sales trends for products over the past three months, providing insights into which products are performing well during this period.\n",
721 "\n",
722 "Running SQL Query...\n",
723 "Error running query: SQL execution error: Execution failed on sql 'SELECT DATE_TRUNC('week', o.OrderDate) AS Week, p.Name AS ProductName, SUM(o.TotalAmount) AS WeeklySales\n",
724 "FROM Orders AS o\n",
725 "JOIN Products AS p ON o.ProductID = p.ProductID\n",
726 "WHERE o.OrderDate >= CURRENT_DATE - INTERVAL '3 months'\n",
727 "GROUP BY DATE_TRUNC('week', o.OrderDate), p.Name\n",
728 "ORDER BY Week, ProductName;': near \"'3 months'\": syntax error\n"
729 ]
730 }
731 ],
732 "source": [
733 "import re\n",
734 "import sqlite3\n",
735 "import ollama # Your API client for model calls\n",
736 "from pydantic import BaseModel\n",
737 "from typing import Optional, List\n",
738 "from datetime import datetime\n",
739 "import pandas as pd\n",
740 "import matplotlib.pyplot as plt\n",
741 "\n",
742 "#############################################################\n",
743 "# Agent 1: SQL Query Generator Agent (Using deepseek-coder-v2)\n",
744 "#############################################################\n",
745 "\n",
746 "# Updated system context with role definition and conversation examples (in user/assistant format).\n",
747 "schema_context = f\"\"\"\n",
748 "The database has the following tables:\n",
749 "- Orders (OrderID, CustomerID, TotalAmount, OrderDate, ProductID)\n",
750 "- Customers (CustomerID, FirstName, LastName)\n",
751 "- Products (ProductID, Name, Category, Price)\n",
752 "\n",
753 "You are a Query Generator / Database Architect. You are highly experienced in crafting efficient and optimized SQL queries. When responding, always output your answer strictly in JSON format with the following keys:\n",
754 "{{\n",
755 " \"sql_query\": \"<SQL query string [Properly Formated with newlines and indentation for clarity]>\",\n",
756 " \"query_planning\": \"<A detailed, step-by-step explanation of your approach in planning the query. Include which tables are joined, what filters are applied, and why any aggregations or grouping are used.> [Optional]\",\n",
757 " \"explaination\": \"<A concise summary of what your query does and why it is optimal.> [Optional]\"\n",
758 "}}\n",
759 "\n",
760 "Below are some conversation examples:\n",
761 "\n",
762 "Example 1:\n",
763 "User:\n",
764 " \"Generate a query to retrieve order details for a customer named John Doe.\"\n",
765 "Assistant:\n",
766 " {{\n",
767 " \"sql_query\": \"SELECT o.OrderID, c.FirstName, c.LastName, o.TotalAmount, o.OrderDate\\\\nFROM Orders AS o\\\\nJOIN Customers AS c ON o.CustomerID = c.CustomerID\\\\nWHERE c.FirstName = 'John' AND c.LastName = 'Doe';\",\n",
768 " \"query_planning\": \"Plan: Join Orders with Customers using CustomerID. Filter for FirstName 'John' and LastName 'Doe' to isolate the orders. Then select the necessary order columns.\",\n",
769 " \"explaination\": \"Retrieves detailed order information for customer John Doe.\"\n",
770 " }}\n",
771 "\n",
772 "Example 2:\n",
773 "User:\n",
774 " \"Generate a query that lists the orders with product details for all orders in the 'Electronics' category.\"\n",
775 "Assistant:\n",
776 " {{\n",
777 " \"sql_query\": \"SELECT o.OrderID, p.Name AS ProductName, o.TotalAmount, o.OrderDate\\\\nFROM Orders AS o\\\\nJOIN Products AS p ON o.ProductID = p.ProductID\\\\nWHERE p.Category = 'Electronics';\",\n",
778 " \"query_planning\": \"Plan: Join Orders with Products on ProductID to connect each order with its product information. Filter to only include rows where the Category is 'Electronics'.\",\n",
779 " \"explaination\": \"Lists order and product details for orders that involve electronics.\"\n",
780 " }}\n",
781 "\n",
782 "Example 3:\n",
783 "User:\n",
784 " \"Generate a query to calculate total sales per product category.\"\n",
785 "Assistant:\n",
786 " {{\n",
787 " \"sql_query\": \"SELECT p.Category, SUM(o.TotalAmount) AS TotalSales\\\\nFROM Orders AS o\\\\nJOIN Products AS p ON o.ProductID = p.ProductID\\\\nGROUP BY p.Category;\",\n",
788 " \"query_planning\": \"Plan: Join Orders with Products using ProductID. Group the orders by Category and aggregate TotalAmount to compute sales per category.\",\n",
789 " \"explaination\": \"Aggregates total sales by product category.\"\n",
790 " }}\n",
791 "\n",
792 "Example 4:\n",
793 "User:\n",
794 " \"Generate a query to list unique customers who have placed orders within the last 30 days.\"\n",
795 "Assistant:\n",
796 " {{\n",
797 " \"sql_query\": \"SELECT DISTINCT c.CustomerID, c.FirstName, c.LastName\\\\nFROM Customers AS c\\\\nJOIN Orders AS o ON c.CustomerID = o.CustomerID\\\\nWHERE o.OrderDate >= CURRENT_DATE - INTERVAL '30' DAY;\",\n",
798 " \"query_planning\": \"Plan: Join Customers with Orders on CustomerID. Filter orders from the last 30 days and use DISTINCT to remove duplicate customers.\",\n",
799 " \"explaination\": \"Returns a unique list of customers who placed orders within the last 30 days.\"\n",
800 " }}\n",
801 "\n",
802 "Current Date and Time: {str(datetime.now())}\n",
803 "\"\"\"\n",
804 "\n",
805 "# Pydantic model for SQL query responses.\n",
806 "class SQLQueryResponse(BaseModel):\n",
807 " sql_query: str\n",
808 " query_planning: Optional[str] = None\n",
809 " explaination: Optional[str] = None\n",
810 "\n",
811 "def extract_json_from_markdown(text: str) -> str:\n",
812 " pattern = r\"```json\\s*(\\{.*\\})\\s*```\"\n",
813 " match = re.search(pattern, text, re.DOTALL)\n",
814 " if match:\n",
815 " return match.group(1)\n",
816 " return text.strip()\n",
817 "\n",
818 "def build_messages(history: List[dict], new_query: str) -> List[dict]:\n",
819 " messages = [{\"role\": \"system\", \"content\": schema_context}]\n",
820 " messages.extend(history)\n",
821 " messages.append({\"role\": \"user\", \"content\": f\"User query: {new_query}\\nResponse (provide output in JSON format):\"})\n",
822 " return messages\n",
823 "\n",
824 "def sql_query_generator(history: List[dict], query: str) -> SQLQueryResponse:\n",
825 " messages = build_messages(history, query)\n",
826 " response = ollama.chat(\n",
827 " model=\"deepseek-coder-v2:latest\", # SQL Query generation model.\n",
828 " messages=messages,\n",
829 " options={'temperature': 0.0}\n",
830 " )\n",
831 " json_str = extract_json_from_markdown(response.message.content)\n",
832 " try:\n",
833 " structured_response = SQLQueryResponse.model_validate_json(json_str)\n",
834 " except Exception as e:\n",
835 " raise ValueError(f\"Failed to parse JSON response: {e}\")\n",
836 " history.append({\"role\": \"user\", \"content\": query})\n",
837 " history.append({\"role\": \"assistant\", \"content\": response.message.content})\n",
838 " return structured_response\n",
839 "\n",
840 "#############################################################\n",
841 "# Agent 2: SQL Query Runner Agent (Simulated with SQLite)\n",
842 "#############################################################\n",
843 "\n",
844 "def sql_query_runner(sql_query: str) -> pd.DataFrame:\n",
845 " \"\"\"\n",
846 " Simulate running the SQL query against an in-memory SQLite database populated with dummy data.\n",
847 " \"\"\"\n",
848 " conn = sqlite3.connect(\":memory:\")\n",
849 " \n",
850 " create_tables = [\n",
851 " \"\"\"\n",
852 " CREATE TABLE Customers (\n",
853 " CustomerID INTEGER PRIMARY KEY,\n",
854 " FirstName TEXT,\n",
855 " LastName TEXT\n",
856 " );\n",
857 " \"\"\",\n",
858 " \"\"\"\n",
859 " CREATE TABLE Products (\n",
860 " ProductID INTEGER PRIMARY KEY,\n",
861 " Name TEXT,\n",
862 " Category TEXT,\n",
863 " Price REAL\n",
864 " );\n",
865 " \"\"\",\n",
866 " \"\"\"\n",
867 " CREATE TABLE Orders (\n",
868 " OrderID INTEGER PRIMARY KEY,\n",
869 " CustomerID INTEGER,\n",
870 " TotalAmount REAL,\n",
871 " OrderDate TEXT,\n",
872 " ProductID INTEGER,\n",
873 " FOREIGN KEY (CustomerID) REFERENCES Customers(CustomerID),\n",
874 " FOREIGN KEY (ProductID) REFERENCES Products(ProductID)\n",
875 " );\n",
876 " \"\"\"\n",
877 " ]\n",
878 " \n",
879 " for stmt in create_tables:\n",
880 " conn.execute(stmt)\n",
881 " \n",
882 " dummy_data = {\n",
883 " \"Customers\": [\n",
884 " (1, \"John\", \"Doe\"),\n",
885 " (2, \"Jane\", \"Smith\"),\n",
886 " (3, \"Alice\", \"Brown\")\n",
887 " ],\n",
888 " \"Products\": [\n",
889 " (1, \"Laptop\", \"Electronics\", 1200.0),\n",
890 " (2, \"Smartphone\", \"Electronics\", 800.0),\n",
891 " (3, \"Desk Chair\", \"Furniture\", 150.0)\n",
892 " ],\n",
893 " \"Orders\": [\n",
894 " (1, 1, 1200.0, \"2025-04-01\", 1),\n",
895 " (2, 2, 800.0, \"2025-04-05\", 2),\n",
896 " (3, 1, 150.0, \"2025-04-10\", 3),\n",
897 " (4, 3, 800.0, \"2025-04-11\", 2)\n",
898 " ]\n",
899 " }\n",
900 " \n",
901 " conn.executemany(\"INSERT INTO Customers VALUES (?, ?, ?);\", dummy_data[\"Customers\"])\n",
902 " conn.executemany(\"INSERT INTO Products VALUES (?, ?, ?, ?);\", dummy_data[\"Products\"])\n",
903 " conn.executemany(\"INSERT INTO Orders VALUES (?, ?, ?, ?, ?);\", dummy_data[\"Orders\"])\n",
904 " conn.commit()\n",
905 " \n",
906 " try:\n",
907 " df = pd.read_sql_query(sql_query, conn)\n",
908 " except Exception as e:\n",
909 " raise ValueError(f\"SQL execution error: {e}\")\n",
910 " finally:\n",
911 " conn.close()\n",
912 " \n",
913 " return df\n",
914 "\n",
915 "#############################################################\n",
916 "# Agent 3: Data Analyzer and Visualization Expert Agent\n",
917 "#############################################################\n",
918 "\n",
919 "def data_analyzer_visualizer(data: pd.DataFrame) -> None:\n",
920 " \"\"\"\n",
921 " Analyze the DataFrame intelligently, decide which charts are most useful based on data types, and plot the charts.\n",
922 " \n",
923 " Steps:\n",
924 " - Detect numeric, categorical, and date columns.\n",
925 " - If a date column exists along with numeric data, attempt a time series analysis.\n",
926 " - For multiple numeric columns, display a correlation heatmap or scatter plots.\n",
927 " - For categorical columns with few unique values, display bar charts.\n",
928 " - Also print a summary of the data.\n",
929 " \"\"\"\n",
930 " if data.empty:\n",
931 " print(\"No data found.\")\n",
932 " return\n",
933 "\n",
934 " print(\"----- Data Analysis Summary -----\")\n",
935 " print(data.describe(include='all'))\n",
936 " \n",
937 " # Detect column types.\n",
938 " numeric_cols = data.select_dtypes(include=['number']).columns.tolist()\n",
939 " categorical_cols = data.select_dtypes(include=['object']).columns.tolist()\n",
940 " # Check for columns that might represent dates.\n",
941 " date_cols = [col for col in data.columns if 'date' in col.lower()]\n",
942 " \n",
943 " print(f\"Detected Numeric Columns: {numeric_cols}\")\n",
944 " print(f\"Detected Categorical Columns: {categorical_cols}\")\n",
945 " print(f\"Detected Date Columns: {date_cols}\")\n",
946 " \n",
947 " # Try converting detected date columns to datetime, if not already.\n",
948 " for col in date_cols:\n",
949 " try:\n",
950 " data[col] = pd.to_datetime(data[col])\n",
951 " except Exception as e:\n",
952 " print(f\"Could not convert column '{col}' to datetime: {e}\")\n",
953 " \n",
954 " # Decision logic for plotting:\n",
955 " if date_cols and numeric_cols:\n",
956 " # Plot a time series for the first date and numeric column.\n",
957 " date_col = date_cols[0]\n",
958 " num_col = numeric_cols[0]\n",
959 " try:\n",
960 " ts_data = data.groupby(date_col)[num_col].sum().reset_index().sort_values(by=date_col)\n",
961 " plt.figure(figsize=(8, 5))\n",
962 " plt.plot(ts_data[date_col], ts_data[num_col], marker='o')\n",
963 " plt.title(f\"Time Series: {num_col} over {date_col}\")\n",
964 " plt.xlabel(date_col)\n",
965 " plt.ylabel(num_col)\n",
966 " plt.xticks(rotation=45)\n",
967 " plt.tight_layout()\n",
968 " plt.show()\n",
969 " except Exception as e:\n",
970 " print(\"Time series plotting failed:\", e)\n",
971 " elif len(numeric_cols) >= 2:\n",
972 " # Display a correlation heatmap using matplotlib only.\n",
973 " try:\n",
974 " corr = data[numeric_cols].corr()\n",
975 " plt.figure(figsize=(8, 6))\n",
976 " plt.imshow(corr, cmap='coolwarm', interpolation='nearest')\n",
977 " plt.title(\"Correlation Heatmap\")\n",
978 " plt.colorbar()\n",
979 " tick_marks = range(len(numeric_cols))\n",
980 " plt.xticks(tick_marks, numeric_cols, rotation=45)\n",
981 " plt.yticks(tick_marks, numeric_cols)\n",
982 " for i in range(len(numeric_cols)):\n",
983 " for j in range(len(numeric_cols)):\n",
984 " plt.text(j, i, f\"{corr.iloc[i, j]:.2f}\", ha=\"center\", va=\"center\", color=\"black\")\n",
985 " plt.tight_layout()\n",
986 " plt.show()\n",
987 " except Exception as e:\n",
988 " print(\"Correlation heatmap plotting failed:\", e)\n",
989 " elif len(numeric_cols) == 1:\n",
990 " # Plot a histogram for the single numeric column.\n",
991 " plt.figure(figsize=(8, 5))\n",
992 " data[numeric_cols[0]].hist()\n",
993 " plt.title(f\"Histogram of {numeric_cols[0]}\")\n",
994 " plt.xlabel(numeric_cols[0])\n",
995 " plt.ylabel(\"Frequency\")\n",
996 " plt.tight_layout()\n",
997 " plt.show()\n",
998 " \n",
999 " # For categorical columns with few unique values, plot bar charts.\n",
1000 " for col in categorical_cols:\n",
1001 " if data[col].nunique() <= 10:\n",
1002 " counts = data[col].value_counts()\n",
1003 " plt.figure(figsize=(8, 5))\n",
1004 " counts.plot(kind='bar')\n",
1005 " plt.title(f\"Count of {col}\")\n",
1006 " plt.xlabel(col)\n",
1007 " plt.ylabel(\"Count\")\n",
1008 " plt.tight_layout()\n",
1009 " plt.show()\n",
1010 " \n",
1011 " # Optional: Get additional textual insights using another model.\n",
1012 " analysis_prompt = f\"Analyze the following data summary and provide insights:\\n{data.describe(include='all').to_string()}\"\n",
1013 " analysis_response = ollama.chat(\n",
1014 " model=\"codellama:latest\",\n",
1015 " messages=[{\"role\": \"user\", \"content\": analysis_prompt}],\n",
1016 " options={'temperature': 0.0}\n",
1017 " )\n",
1018 " print(\"Additional Model Analysis:\")\n",
1019 " print(analysis_response.message.content)\n",
1020 "\n",
1021 "#############################################################\n",
1022 "# Agent 4: Orchestrator Agent\n",
1023 "#############################################################\n",
1024 "\n",
1025 "def orchestrator(user_query: str) -> None:\n",
1026 " \"\"\"\n",
1027 " The orchestrator coordinates all agents:\n",
1028 " 1. Generates an SQL query from a natural language request.\n",
1029 " 2. Executes the SQL query on a simulated database.\n",
1030 " 3. Analyzes and visualizes the resulting data.\n",
1031 " \"\"\"\n",
1032 " conversation_history: List[dict] = [] # Persistent conversation history.\n",
1033 " \n",
1034 " # Step 1: Generate SQL query.\n",
1035 " sql_response = sql_query_generator(conversation_history, user_query)\n",
1036 " print(\"Generated SQL Query:\")\n",
1037 " print(sql_response.sql_query)\n",
1038 " print(\"Query Planning:\")\n",
1039 " print(sql_response.query_planning)\n",
1040 " print(\"Explanation:\")\n",
1041 " print(sql_response.explaination)\n",
1042 " \n",
1043 " # Step 2: Execute the SQL query.\n",
1044 " print(\"\\nRunning SQL Query...\")\n",
1045 " try:\n",
1046 " result_df = sql_query_runner(sql_response.sql_query)\n",
1047 " print(\"Query Result:\")\n",
1048 " print(result_df)\n",
1049 " except Exception as e:\n",
1050 " print(f\"Error running query: {e}\")\n",
1051 " return\n",
1052 " \n",
1053 " # Step 3: Analyze and visualize the result.\n",
1054 " print(\"\\nData Analysis & Visualization:\")\n",
1055 " data_analyzer_visualizer(result_df)\n",
1056 "\n",
1057 "#############################################################\n",
1058 "# Example Usage\n",
1059 "#############################################################\n",
1060 "\n",
1061 "if __name__ == '__main__':\n",
1062 " # Sample natural language query.\n",
1063 " user_nl_query = \"Find weekly product sales trends in last quarter?\"\n",
1064 " orchestrator(user_nl_query)\n"
1065 ]
1066 },
1067 {
1068 "cell_type": "code",
1069 "execution_count": null,
1070 "id": "c42a8b28-41bf-429c-893e-5a6dba14445e",
1071 "metadata": {},
1072 "outputs": [],
1073 "source": []
1074 },
1075 {
1076 "cell_type": "code",
1077 "execution_count": null,
1078 "id": "1888d2e6-76a6-4d83-ae74-b639bd98259d",
1079 "metadata": {},
1080 "outputs": [],
1081 "source": []
1082 },
1083 {
1084 "cell_type": "code",
1085 "execution_count": null,
1086 "id": "b2d2ebc6-e6bb-48f6-bcf6-dc55bee4d67b",
1087 "metadata": {},
1088 "outputs": [],
1089 "source": []
1090 },
1091 {
1092 "cell_type": "code",
1093 "execution_count": null,
1094 "id": "03184f1c-8e7a-455d-816f-a83075fadc0f",
1095 "metadata": {},
1096 "outputs": [],
1097 "source": []
1098 },
1099 {
1100 "cell_type": "code",
1101 "execution_count": null,
1102 "id": "493a493b-752d-4da5-93a7-4fbb016c7b87",
1103 "metadata": {},
1104 "outputs": [],
1105 "source": []
1106 },
1107 {
1108 "cell_type": "code",
1109 "execution_count": 3,
1110 "id": "87a96db6-b6eb-4387-9e2c-ac48d482ac2e",
1111 "metadata": {},
1112 "outputs": [
1113 {
1114 "name": "stdout",
1115 "output_type": "stream",
1116 "text": [
1117 "Generated SQL Query:\n",
1118 "SELECT p.*\n",
1119 "FROM Products p\n",
1120 "LEFT JOIN Orders o ON p.ProductID = o.ProductID\n",
1121 "WHERE o.OrderDate IS NULL AND o.OrderDate < DATE('now', '-2 weeks')\n",
1122 "Query Planning:\n",
1123 "1. Start by selecting all columns from the Products table (p). \n",
1124 "2. Use a LEFT JOIN with the Orders table on ProductID to find matching products that have been ordered.\n",
1125 "3. Filter the results where there is no corresponding order in the last 2 weeks using WHERE o.OrderDate IS NULL AND o.OrderDate < DATE('now', '-2 weeks').\n",
1126 "4. This approach ensures we only consider products not bought by any customer in the specified time frame.\n",
1127 "Explanation:\n",
1128 "This query efficiently identifies products that have never been ordered since the last 2 weeks by using a LEFT JOIN to check for orders and applying a date filter to exclude recent orders. It avoids retrieving unnecessary data, making it both space-efficient and performant.\n",
1129 "\n",
1130 "Running SQL Query...\n",
1131 "Query Result:\n",
1132 "Empty DataFrame\n",
1133 "Columns: [ProductID, Name, Category, Price]\n",
1134 "Index: []\n",
1135 "\n",
1136 "Data Analysis & Visualization:\n",
1137 "Data is empty.\n"
1138 ]
1139 }
1140 ],
1141 "source": [
1142 "import re\n",
1143 "import sqlite3\n",
1144 "import time\n",
1145 "import json\n",
1146 "import ollama # Your API client for model calls.\n",
1147 "from pydantic import BaseModel\n",
1148 "from typing import Optional, List, Union\n",
1149 "from datetime import datetime\n",
1150 "import pandas as pd\n",
1151 "import matplotlib.pyplot as plt\n",
1152 "\n",
1153 "#############################################################\n",
1154 "# Utility: Extract JSON from markdown output.\n",
1155 "#############################################################\n",
1156 "\n",
1157 "def extract_json_from_markdown(text: str) -> str:\n",
1158 " pattern = r\"```json\\s*(\\{.*\\})\\s*```\"\n",
1159 " match = re.search(pattern, text, re.DOTALL)\n",
1160 " if match:\n",
1161 " return match.group(1)\n",
1162 " return text.strip()\n",
1163 "\n",
1164 "#############################################################\n",
1165 "# Agent 1: SQL Query Generator Agent (Using deepseek-coder-v2:latest)\n",
1166 "#############################################################\n",
1167 "\n",
1168 "schema_context = f\"\"\"\n",
1169 "The database i'm using is SQLite [so use appropriate syntax]\n",
1170 "\n",
1171 "The database has the following tables:\n",
1172 "- Orders (OrderID, CustomerID, TotalAmount, OrderDate, ProductID)\n",
1173 "- Customers (CustomerID, FirstName, LastName)\n",
1174 "- Products (ProductID, Name, Category, Price)\n",
1175 "\"\"\"\n",
1176 "\n",
1177 "agent_role=\"You are a Query Generator / Database Architect. You are highly experienced in crafting efficient and optimized SQL queries.\"\n",
1178 "\n",
1179 "output_format = \"\"\"When responding, always output your answer strictly in JSON format with the following keys:\n",
1180 "{{\n",
1181 " \"sql_query\": \"<SQL query string [Properly Formated with newlines and indentation for clarity]>\",\n",
1182 " \"query_planning\": \"<A detailed, step-by-step explanation of your approach in planning the query. Include which tables are joined, what filters are applied, and why any aggregations or grouping are used.> [Optional]\",\n",
1183 " \"explaination\": \"<A concise summary of what your query does and why it is optimal.> [Optional]\"\n",
1184 "}}\n",
1185 "\"\"\"\n",
1186 "\n",
1187 "# Pydantic model for SQL query responses.\n",
1188 "# Updated: query_planning is now a union of string and dict.\n",
1189 "class SQLQueryResponse(BaseModel):\n",
1190 " sql_query: str\n",
1191 " query_planning: Optional[Union[str, dict]] = None\n",
1192 " explaination: Optional[str] = None\n",
1193 "\n",
1194 "def build_messages(history: List[dict], new_query: str) -> List[dict]:\n",
1195 " messages = [{\"role\": \"system\", \"content\": f\"Role: {agent_role} \\nDB Schema context:{schema_context}\\nOutput Format:{output_format}\" } ]\n",
1196 " messages.extend(history)\n",
1197 " messages.append({\"role\": \"user\", \"content\": f\"User query: {new_query}\\nResponse (provide output in JSON format):\"})\n",
1198 " return messages\n",
1199 "\n",
1200 "def sql_query_generator(history: List[dict], query: str) -> SQLQueryResponse:\n",
1201 " messages = build_messages(history, query)\n",
1202 " response = ollama.chat(\n",
1203 " model=\"deepseek-coder-v2:latest\", # SQL Query generation model.\n",
1204 " messages=messages,\n",
1205 " options={'temperature': 0.0}\n",
1206 " )\n",
1207 " json_str = extract_json_from_markdown(response.message.content)\n",
1208 " try:\n",
1209 " structured_response = SQLQueryResponse.model_validate_json(json_str)\n",
1210 " except Exception as e:\n",
1211 " raise ValueError(f\"Failed to parse JSON response: {e}\")\n",
1212 " history.append({\"role\": \"user\", \"content\": query})\n",
1213 " history.append({\"role\": \"assistant\", \"content\": response.message.content})\n",
1214 " return structured_response\n",
1215 "\n",
1216 "#############################################################\n",
1217 "# Agent 2: SQL Query Runner Agent (Simulated with SQLite)\n",
1218 "#############################################################\n",
1219 "\n",
1220 "def sql_query_runner(sql_query: str) -> pd.DataFrame:\n",
1221 " \"\"\"\n",
1222 " Simulate running the SQL query against an in-memory SQLite database populated with dummy data.\n",
1223 " \"\"\"\n",
1224 " conn = sqlite3.connect(\":memory:\")\n",
1225 " \n",
1226 " create_tables = [\n",
1227 " \"\"\"\n",
1228 " CREATE TABLE Customers (\n",
1229 " CustomerID INTEGER PRIMARY KEY,\n",
1230 " FirstName TEXT,\n",
1231 " LastName TEXT\n",
1232 " );\n",
1233 " \"\"\",\n",
1234 " \"\"\"\n",
1235 " CREATE TABLE Products (\n",
1236 " ProductID INTEGER PRIMARY KEY,\n",
1237 " Name TEXT,\n",
1238 " Category TEXT,\n",
1239 " Price REAL\n",
1240 " );\n",
1241 " \"\"\",\n",
1242 " \"\"\"\n",
1243 " CREATE TABLE Orders (\n",
1244 " OrderID INTEGER PRIMARY KEY,\n",
1245 " CustomerID INTEGER,\n",
1246 " TotalAmount REAL,\n",
1247 " OrderDate TEXT,\n",
1248 " ProductID INTEGER,\n",
1249 " FOREIGN KEY (CustomerID) REFERENCES Customers(CustomerID),\n",
1250 " FOREIGN KEY (ProductID) REFERENCES Products(ProductID)\n",
1251 " );\n",
1252 " \"\"\"\n",
1253 " ]\n",
1254 " \n",
1255 " for stmt in create_tables:\n",
1256 " conn.execute(stmt)\n",
1257 " \n",
1258 " dummy_data = {\n",
1259 " \"Customers\": [\n",
1260 " (1, \"John\", \"Doe\"),\n",
1261 " (2, \"Jane\", \"Smith\"),\n",
1262 " (3, \"Alice\", \"Brown\")\n",
1263 " ],\n",
1264 " \"Products\": [\n",
1265 " (1, \"Laptop\", \"Electronics\", 1200.0),\n",
1266 " (2, \"Smartphone\", \"Electronics\", 800.0),\n",
1267 " (3, \"Desk Chair\", \"Furniture\", 150.0)\n",
1268 " ],\n",
1269 " \"Orders\": [\n",
1270 " (1, 1, 1200.0, \"2025-04-01\", 1),\n",
1271 " (2, 2, 800.0, \"2025-04-05\", 2),\n",
1272 " (3, 1, 150.0, \"2025-04-10\", 3),\n",
1273 " (4, 3, 800.0, \"2025-04-11\", 2)\n",
1274 " ]\n",
1275 " }\n",
1276 " \n",
1277 " conn.executemany(\"INSERT INTO Customers VALUES (?, ?, ?);\", dummy_data[\"Customers\"])\n",
1278 " conn.executemany(\"INSERT INTO Products VALUES (?, ?, ?, ?);\", dummy_data[\"Products\"])\n",
1279 " conn.executemany(\"INSERT INTO Orders VALUES (?, ?, ?, ?, ?);\", dummy_data[\"Orders\"])\n",
1280 " conn.commit()\n",
1281 " \n",
1282 " try:\n",
1283 " df = pd.read_sql_query(sql_query, conn)\n",
1284 " except Exception as e:\n",
1285 " raise ValueError(f\"SQL execution error: {e}\")\n",
1286 " finally:\n",
1287 " conn.close()\n",
1288 " \n",
1289 " return df\n",
1290 "\n",
1291 "#############################################################\n",
1292 "# Agent 3: Data Analyzer and Visualization Expert Agent\n",
1293 "#############################################################\n",
1294 "\n",
1295 "def data_analyzer_visualizer_agent(data: pd.DataFrame, sql_response) -> None:\n",
1296 " \"\"\"\n",
1297 " This agent analyzes the DataFrame intelligently. It examines row count, schema, and sample data,\n",
1298 " then uses an LLM to decide if visualizations are needed. If not, it returns a business summary.\n",
1299 " Otherwise, it returns a set of visualization instructions which are then executed.\n",
1300 " \"\"\"\n",
1301 " if data.empty:\n",
1302 " print(\"Data is empty.\")\n",
1303 " return\n",
1304 "\n",
1305 " row_count = data.shape[0]\n",
1306 " schema_info = data.dtypes.to_string()\n",
1307 " sample_data = data.head(5).to_string()\n",
1308 "\n",
1309 " prompt = f\"\"\"\n",
1310 "You are a Data Analyzer and Visualization Expert and Customer Support Expert, which answers users in a friendly yet business way. Based on the following dataset details, decide whether visualizations\n",
1311 "would provide valuable insights or if a concise business summary is sufficient. Formulate your respose to best answer user asked question. We already have the final data just answer user question.\n",
1312 "\n",
1313 "User Asked Question: {user_nl_query}\n",
1314 "\n",
1315 "Thinking Steps already Executed:\n",
1316 "{sql_response}\n",
1317 " \n",
1318 "Data Details:\n",
1319 "- Row Count: {row_count}\n",
1320 "- Schema:\n",
1321 "{schema_info}\n",
1322 "- Sample Data (first 5 rows):\n",
1323 "{sample_data}\n",
1324 "\n",
1325 "Analyze the above information and decide:\n",
1326 "1. If visualizations are NOT necessary, provide a concise business summary in professional language.\n",
1327 "2. If visualizations would add value, decide which chart types to use (e.g., bar, histogram, time_series, scatter, correlation_heatmap)\n",
1328 " and specify the columns to use along with a brief description for each.\n",
1329 "\n",
1330 "Respond strictly in JSON format with the following keys:\n",
1331 "{{\n",
1332 " \"action\": \"<'visualize' or 'summary'>\",\n",
1333 " \"analysis\": \"<If action is 'summary', try to answer the question asked by User>\",\n",
1334 " \"visualizations\": [<If action is 'visualize', a list of JSON objects with keys 'chart_type', 'columns', 'description'>]\n",
1335 "}}\n",
1336 " \"\"\"\n",
1337 " response = ollama.chat(\n",
1338 " model=\"mistral:latest\",\n",
1339 " messages=[{\"role\": \"user\", \"content\": prompt}],\n",
1340 " options={'temperature': 0.0}\n",
1341 " )\n",
1342 " json_str = extract_json_from_markdown(response.message.content)\n",
1343 " try:\n",
1344 " analysis_output = json.loads(json_str)\n",
1345 " except Exception as e:\n",
1346 " print(f\"Error parsing JSON response: {e}\")\n",
1347 " print(\"Raw response:\")\n",
1348 " print(response.message.content)\n",
1349 " return\n",
1350 "\n",
1351 " action = analysis_output.get(\"action\", \"\").lower()\n",
1352 " if action == \"summary\":\n",
1353 " summary = analysis_output.get(\"analysis\", \"No summary provided.\")\n",
1354 " print(\"Business Summary:\")\n",
1355 " print(summary)\n",
1356 " elif action == \"visualize\":\n",
1357 " visualizations = analysis_output.get(\"visualizations\", [])\n",
1358 " if not visualizations:\n",
1359 " print(\"No visualization instructions provided.\")\n",
1360 " return\n",
1361 " for viz in visualizations:\n",
1362 " chart_type = viz.get(\"chart_type\", \"\").lower()\n",
1363 " columns = viz.get(\"columns\", [])\n",
1364 " description = viz.get(\"description\", \"Visualization\")\n",
1365 " print(f\"Generating visualization: {description}\")\n",
1366 " try:\n",
1367 " if chart_type == \"bar\":\n",
1368 " counts = data[columns[0]].value_counts()\n",
1369 " plt.figure(figsize=(8, 5))\n",
1370 " counts.plot(kind='bar')\n",
1371 " plt.title(description)\n",
1372 " plt.xlabel(columns[0])\n",
1373 " plt.ylabel(\"Count\")\n",
1374 " plt.tight_layout()\n",
1375 " plt.show()\n",
1376 " elif chart_type == \"histogram\":\n",
1377 " plt.figure(figsize=(8, 5))\n",
1378 " data[columns[0]].hist()\n",
1379 " plt.title(description)\n",
1380 " plt.xlabel(columns[0])\n",
1381 " plt.ylabel(\"Frequency\")\n",
1382 " plt.tight_layout()\n",
1383 " plt.show()\n",
1384 " elif chart_type == \"time_series\":\n",
1385 " ts_data = data.groupby(columns[0])[columns[1]].sum().reset_index().sort_values(by=columns[0])\n",
1386 " plt.figure(figsize=(8, 5))\n",
1387 " plt.plot(ts_data[columns[0]], ts_data[columns[1]], marker='o')\n",
1388 " plt.title(description)\n",
1389 " plt.xlabel(columns[0])\n",
1390 " plt.ylabel(columns[1])\n",
1391 " plt.xticks(rotation=45)\n",
1392 " plt.tight_layout()\n",
1393 " plt.show()\n",
1394 " elif chart_type == \"scatter\":\n",
1395 " plt.figure(figsize=(8, 5))\n",
1396 " plt.scatter(data[columns[0]], data[columns[1]])\n",
1397 " plt.title(description)\n",
1398 " plt.xlabel(columns[0])\n",
1399 " plt.ylabel(columns[1])\n",
1400 " plt.tight_layout()\n",
1401 " plt.show()\n",
1402 " elif chart_type == \"correlation_heatmap\":\n",
1403 " numeric_cols = data.select_dtypes(include=['number']).columns.tolist()\n",
1404 " if numeric_cols:\n",
1405 " corr = data[numeric_cols].corr()\n",
1406 " plt.figure(figsize=(8, 6))\n",
1407 " plt.imshow(corr, cmap='coolwarm', interpolation='nearest')\n",
1408 " plt.title(\"Correlation Heatmap\")\n",
1409 " plt.colorbar()\n",
1410 " tick_marks = range(len(numeric_cols))\n",
1411 " plt.xticks(tick_marks, numeric_cols, rotation=45)\n",
1412 " plt.yticks(tick_marks, numeric_cols)\n",
1413 " for i in range(len(numeric_cols)):\n",
1414 " for j in range(len(numeric_cols)):\n",
1415 " plt.text(j, i, f\"{corr.iloc[i, j]:.2f}\", ha=\"center\", va=\"center\", color=\"black\")\n",
1416 " plt.tight_layout()\n",
1417 " plt.show()\n",
1418 " else:\n",
1419 " print(\"No numeric columns available for correlation heatmap.\")\n",
1420 " else:\n",
1421 " print(f\"Chart type '{chart_type}' is not implemented.\")\n",
1422 " except Exception as e:\n",
1423 " print(f\"Error generating visualization '{chart_type}': {e}\")\n",
1424 " else:\n",
1425 " print(\"The AI agent did not return a valid action.\")\n",
1426 "\n",
1427 "#############################################################\n",
1428 "# Agent 4: Query Debugger Agent\n",
1429 "#############################################################\n",
1430 "\n",
1431 "def query_debugger_agent(original_query: str, error_message: str, history: List[dict]) -> SQLQueryResponse:\n",
1432 " \"\"\"\n",
1433 " This agent debugs a failing SQL query. It takes the original query and error message,\n",
1434 " then uses the LLM (deepseek-coder-v2:latest) to generate a corrected query with a detailed debugging plan.\n",
1435 " A retry mechanism is implemented to handle transient errors.\n",
1436 " \"\"\"\n",
1437 " debug_prompt = f\"\"\"\n",
1438 "The following SQL query produced an error during execution:\n",
1439 "Query:\n",
1440 "{original_query}\n",
1441 "\n",
1442 "Error Message:\n",
1443 "{error_message}\n",
1444 "\n",
1445 "As a Query Generator / Database Architect, analyze the error and explain what went wrong.\n",
1446 "Then, provide a corrected SQL query along with a detailed debugging plan describing the changes made.\n",
1447 "\n",
1448 "Respond strictly in JSON format with these keys:\n",
1449 "{{\n",
1450 " \"sql_query\": \"<Revised SQL query string [Properly Formated with newlines and indentation for clarity]>\",\n",
1451 " \"query_planning\": \"<Detailed debugging and reasoning process>\",\n",
1452 " \"explaination\": \"<A concise summary of the debugging process and the corrections made>\"\n",
1453 "}}\n",
1454 "\"\"\"\n",
1455 " messages = [{\"role\": \"system\", \"content\": schema_context},\n",
1456 " {\"role\": \"user\", \"content\": debug_prompt}]\n",
1457 " max_attempts = 5\n",
1458 " attempt = 0\n",
1459 " while attempt < max_attempts:\n",
1460 " try:\n",
1461 " response = ollama.chat(\n",
1462 " model=\"deepseek-coder-v2:latest\",\n",
1463 " messages=messages,\n",
1464 " options={'temperature': 0.9}\n",
1465 " )\n",
1466 " json_str = extract_json_from_markdown(response.message.content)\n",
1467 " debug_response = SQLQueryResponse.model_validate_json(json_str)\n",
1468 " history.append({\"role\": \"user\", \"content\": debug_prompt})\n",
1469 " history.append({\"role\": \"assistant\", \"content\": response.message.content})\n",
1470 " return debug_response\n",
1471 " except Exception as e:\n",
1472 " attempt += 1\n",
1473 " print(f\"Query Debugger attempt {attempt} failed with error: {str(e)}\")\n",
1474 " time.sleep(2)\n",
1475 " raise ValueError(f\"Query Debugger failed after {max_attempts} attempts.\")\n",
1476 "\n",
1477 "#############################################################\n",
1478 "# Agent 5: Orchestrator Agent\n",
1479 "#############################################################\n",
1480 "\n",
1481 "def orchestrator(user_query: str) -> None:\n",
1482 " \"\"\"\n",
1483 " The orchestrator coordinates all agents:\n",
1484 " 1. Generate an SQL query from a natural language request.\n",
1485 " 2. Execute the query on a simulated database.\n",
1486 " 3. If the query fails, invoke the Query Debugger Agent to fix it.\n",
1487 " 4. Analyze and visualize the resulting data.\n",
1488 " \"\"\"\n",
1489 " conversation_history: List[dict] = [] # Persistent history for the session.\n",
1490 " \n",
1491 " # Step 1: Generate SQL query.\n",
1492 " sql_response = sql_query_generator(conversation_history, user_query)\n",
1493 " print(\"Generated SQL Query:\")\n",
1494 " print(sql_response.sql_query)\n",
1495 " print(\"Query Planning:\")\n",
1496 " print(sql_response.query_planning)\n",
1497 " print(\"Explanation:\")\n",
1498 " print(sql_response.explaination)\n",
1499 " \n",
1500 " # Step 2: Execute the SQL query.\n",
1501 " print(\"\\nRunning SQL Query...\")\n",
1502 " try:\n",
1503 " result_df = sql_query_runner(sql_response.sql_query)\n",
1504 " print(\"Query Result:\")\n",
1505 " print(result_df)\n",
1506 " except Exception as e:\n",
1507 " error_message = str(e)\n",
1508 " print(f\"Original query failed with error: {error_message}\")\n",
1509 " # Invoke Query Debugger Agent.\n",
1510 " debug_response = query_debugger_agent(sql_response.sql_query, error_message, conversation_history)\n",
1511 " print(\"\\nDebugged SQL Query:\")\n",
1512 " print(debug_response.sql_query)\n",
1513 " print(\"Debug Query Planning:\")\n",
1514 " print(debug_response.query_planning)\n",
1515 " print(\"Debug Explanation:\")\n",
1516 " print(debug_response.explaination)\n",
1517 " # Attempt to run the debugged query.\n",
1518 " try:\n",
1519 " result_df = sql_query_runner(debug_response.sql_query)\n",
1520 " print(\"Query Result from Debugged Query:\")\n",
1521 " print(result_df)\n",
1522 " except Exception as e2:\n",
1523 " print(f\"Debugged query also failed: {str(e2)}\")\n",
1524 " return\n",
1525 " \n",
1526 " # Step 3: Analyze and visualize the result.\n",
1527 " print(\"\\nData Analysis & Visualization:\")\n",
1528 " data_analyzer_visualizer_agent(result_df, sql_response)\n",
1529 "\n",
1530 "#############################################################\n",
1531 "# Example Usage\n",
1532 "#############################################################\n",
1533 "\n",
1534 "if __name__ == '__main__':\n",
1535 " # Example natural language query.\n",
1536 " user_nl_query = \"Products which are never bought by any customer since last 2 weeks \"\n",
1537 " orchestrator(user_nl_query)\n"
1538 ]
1539 },
1540 {
1541 "cell_type": "code",
1542 "execution_count": null,
1543 "id": "43799253-8a85-40ad-8ff2-e88dd0349c7e",
1544 "metadata": {},
1545 "outputs": [],
1546 "source": []
1547 }
1548 ],
1549 "metadata": {
1550 "kernelspec": {
1551 "display_name": "Python 3",
1552 "language": "python",
1553 "name": "python3"
1554 },
1555 "language_info": {
1556 "codemirror_mode": {
1557 "name": "ipython",
1558 "version": 3
1559 },
1560 "file_extension": ".py",
1561 "mimetype": "text/x-python",
1562 "name": "python",
1563 "nbconvert_exporter": "python",
1564 "pygments_lexer": "ipython3",
1565 "version": "3.8.10"
1566 }
1567 },
1568 "nbformat": 4,
1569 "nbformat_minor": 5
1570 }
1571