baseline-agent.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Wrapper class for the locally hosted LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests as rq\n",
    "import json\n",
    "\n",
    "class LLM():\n",
    "    def __init__(self, hostname, port):\n",
    "        self.url = f\"http://{hostname}:{port}/api/v1/generate\"\n",
    "    \n",
    "    def get_response(self, user_request):\n",
    "        # Sends a request to the LLM with the complete prompt\n",
    "        # The preset used for the parameters was Divine Intellect\n",
    "        response = rq.post(self.url, json.dumps({\n",
    "            \"max_context_length\": 2048, \n",
    "            \"max_length\": 120, \n",
    "            \"prompt\": user_request,\n",
    "            \"temperature\": 1.31,\n",
    "            \"top_k\": 49,\n",
    "            \"top_p\": 0.14,\n",
    "            \"stop_sequence\": [\"</SQLQuery>\\n\"]}))\n",
    "        return response.json()[\"results\"][0][\"text\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = LLM(\"localhost\", 5001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Agent class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sqlalchemy as db\n",
    "from sqlalchemy.sql import text\n",
    "import random as rd\n",
    "\n",
    "class QueriesTranslationAgent():\n",
    "    def __init__(self, llm, _db, examples_file):\n",
    "        self.llm = llm\n",
    "        self.engine = db.create_engine(f\"sqlite:///{_db}\")\n",
    "        with open(examples_file, \"r\") as file:\n",
    "            self.examples = file.read()\n",
    "        self.examples = self.examples.split(\"\\n\\n\")\n",
    "    \n",
    "    def build_prompt(self, user_request):\n",
    "        # For the prompt, two shot prompting was used, along with a template that explained the task\n",
    "        example1, example2 = rd.sample(self.examples, 2)\n",
    "        template = \"\"\"\n",
    "You are an agent designed to translate questions about a database into SQL queries. You have a table called apple_daily with the following columns:\n",
    "\n",
    "(Date DATE, Time TIME, Open REAL, High REAL, Low REAL, Close REAL, Volume INTEGER)\n",
    "\n",
    "Never query for all the columns from a specific table, only ask for the relevant columns given the question.\n",
    "DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.\n",
    "If the question does not seem related to the database, just return \"I don't know\" as the answer.\n",
    "You should always think about what to do step by step.\n",
    "Translate:\n",
    "\n",
    "{example1}\n",
    "\n",
    "{example2}\n",
    "\n",
    "<question>\n",
    "Question: {user_request}\n",
    "<question>\n",
    "<SQLQuery>\n",
    "\"\"\"\n",
    "        prompt = template.format(example1=example1, example2=example2, user_request=user_request)\n",
    "        return prompt\n",
    "\n",
    "    def parse_raw_response(self, response):\n",
    "        query = response.split(\"\\n\")[0].split(\"SQLQuery: \")[1]\n",
    "        return query\n",
    "\n",
    "    def invoke(self, user_request):\n",
    "        prompt = self.build_prompt(user_request)\n",
    "        response = self.llm.get_response(prompt)\n",
    "        query = self.parse_raw_response(response)\n",
    "        query_result = self.engine.connect().execute(text(query)).fetchall()\n",
    "        self.engine.dispose()\n",
    "        return query_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('2023-01-03', '01:00:00', 130.28, 130.9, 124.17, 125.07, 112117471), ('2023-01-04', '01:00:00', 126.89, 128.6557, 125.08, 126.36, 89113633), ('2023-01-05', '01:00:00', 127.13, 127.77, 124.76, 125.02, 80962708), ('2023-01-06', '01:00:00', 126.01, 130.29, 124.89, 129.62, 87754715), ('2023-01-09', '01:00:00', 130.465, 133.41, 129.89, 130.15, 70790813)]\n"
     ]
    }
   ],
   "source": [
    "agent = QueriesTranslationAgent(llm, \"apple.db\", \"examples.txt\")\n",
    "print(agent.invoke(\"Give me the first 5 values of the table apple_daily\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('2023-01-04', '01:00:00', 126.89, 128.6557, 125.08, 126.36, 89113633)]\n"
     ]
    }
   ],
   "source": [
    "print(agent.invoke(\"Give me all the rows corresponding to the date 2023-01-04\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('2023-01-04', '01:00:00', 126.89, 128.6557, 125.08, 126.36, 89113633), ('2023-01-05', '01:00:00', 127.13, 127.77, 124.76, 125.02, 80962708), ('2023-01-06', '01:00:00', 126.01, 130.29, 124.89, 129.62, 87754715), ('2023-01-09', '01:00:00', 130.465, 133.41, 129.89, 130.15, 70790813), ('2023-01-10', '01:00:00', 130.26, 131.2636, 128.12, 130.73, 63896155), ('2023-01-11', '01:00:00', 131.25, 133.51, 130.46, 133.49, 69458949), ('2023-01-12', '01:00:00', 133.88, 134.26, 131.44, 133.41, 71379648), ('2023-01-13', '01:00:00', 132.03, 134.92, 131.66, 134.76, 57809719)]\n"
     ]
    }
   ],
   "source": [
    "print(agent.invoke(\"Give me all the rows whose date goes between 2023-01-04 and 2023-01-15, including both dates too\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(250,)]\n"
     ]
    }
   ],
   "source": [
    "print(agent.invoke(\"Count how many rows are there in the table apple_daily\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(0,)]\n"
     ]
    }
   ],
   "source": [
    "print(agent.invoke(\"Count how many rows in the table are from the year 2023\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "ename": "OperationalError",
     "evalue": "(sqlite3.OperationalError) no such column: Month\n[SQL: SELECT COUNT(*) FROM apple_daily WHERE Month = 3; ]\n(Background on this error at: https://sqlalche.me/e/20/e3q8)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mOperationalError\u001b[0m                          Traceback (most recent call last)",
      "File \u001b[1;32mc:\\UBA\\NL_To_SQL\\.conda\\Lib\\site-packages\\sqlalchemy\\engine\\base.py:1967\u001b[0m, in \u001b[0;36mConnection._exec_single_context\u001b[1;34m(self, dialect, context, statement, parameters)\u001b[0m\n\u001b[0;32m   1966\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m evt_handled:\n\u001b[1;32m-> 1967\u001b[0m         \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdialect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdo_execute\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m   1968\u001b[0m \u001b[43m            \u001b[49m\u001b[43mcursor\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstr_statement\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43meffective_parameters\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\n\u001b[0;32m   1969\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1971\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_has_events \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mengine\u001b[38;5;241m.\u001b[39m_has_events:\n",
      "File \u001b[1;32mc:\\UBA\\NL_To_SQL\\.conda\\Lib\\site-packages\\sqlalchemy\\engine\\default.py:924\u001b[0m, in \u001b[0;36mDefaultDialect.do_execute\u001b[1;34m(self, cursor, statement, parameters, context)\u001b[0m\n\u001b[0;32m    923\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdo_execute\u001b[39m(\u001b[38;5;28mself\u001b[39m, cursor, statement, parameters, context\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[1;32m--> 924\u001b[0m     \u001b[43mcursor\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mexecute\u001b[49m\u001b[43m(\u001b[49m\u001b[43mstatement\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparameters\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[1;31mOperationalError\u001b[0m: no such column: Month",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001b[1;31mOperationalError\u001b[0m                          Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[9], line 2\u001b[0m\n\u001b[0;32m      1\u001b[0m agent \u001b[38;5;241m=\u001b[39m QueriesTranslationAgent(llm, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mapple.db\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mexamples.txt\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m----> 2\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[43magent\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minvoke\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mCount how many rows in the table are from month 03\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m)\n",
      "Cell \u001b[1;32mIn[3], line 46\u001b[0m, in \u001b[0;36mQueriesTranslationAgent.invoke\u001b[1;34m(self, user_request)\u001b[0m\n\u001b[0;32m     44\u001b[0m response \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mllm\u001b[38;5;241m.\u001b[39mget_response(prompt)\n\u001b[0;32m     45\u001b[0m query \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mparse_raw_response(response)\n\u001b[1;32m---> 46\u001b[0m query_result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mengine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mexecute\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtext\u001b[49m\u001b[43m(\u001b[49m\u001b[43mquery\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mfetchall()\n\u001b[0;32m     47\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mengine\u001b[38;5;241m.\u001b[39mdispose()\n\u001b[0;32m     48\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m query_result\n",
      "File \u001b[1;32mc:\\UBA\\NL_To_SQL\\.conda\\Lib\\site-packages\\sqlalchemy\\engine\\base.py:1418\u001b[0m, in \u001b[0;36mConnection.execute\u001b[1;34m(self, statement, parameters, execution_options)\u001b[0m\n\u001b[0;32m   1416\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m exc\u001b[38;5;241m.\u001b[39mObjectNotExecutableError(statement) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01merr\u001b[39;00m\n\u001b[0;32m   1417\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1418\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mmeth\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m   1419\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1420\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdistilled_parameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1421\u001b[0m \u001b[43m        \u001b[49m\u001b[43mexecution_options\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mNO_OPTIONS\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1422\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\UBA\\NL_To_SQL\\.conda\\Lib\\site-packages\\sqlalchemy\\sql\\elements.py:515\u001b[0m, in \u001b[0;36mClauseElement._execute_on_connection\u001b[1;34m(self, connection, distilled_params, execution_options)\u001b[0m\n\u001b[0;32m    513\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m TYPE_CHECKING:\n\u001b[0;32m    514\u001b[0m         \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(\u001b[38;5;28mself\u001b[39m, Executable)\n\u001b[1;32m--> 515\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mconnection\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execute_clauseelement\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m    516\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdistilled_params\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexecution_options\u001b[49m\n\u001b[0;32m    517\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    518\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m    519\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m exc\u001b[38;5;241m.\u001b[39mObjectNotExecutableError(\u001b[38;5;28mself\u001b[39m)\n",
      "File \u001b[1;32mc:\\UBA\\NL_To_SQL\\.conda\\Lib\\site-packages\\sqlalchemy\\engine\\base.py:1640\u001b[0m, in \u001b[0;36mConnection._execute_clauseelement\u001b[1;34m(self, elem, distilled_parameters, execution_options)\u001b[0m\n\u001b[0;32m   1628\u001b[0m compiled_cache: Optional[CompiledCacheType] \u001b[38;5;241m=\u001b[39m execution_options\u001b[38;5;241m.\u001b[39mget(\n\u001b[0;32m   1629\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcompiled_cache\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mengine\u001b[38;5;241m.\u001b[39m_compiled_cache\n\u001b[0;32m   1630\u001b[0m )\n\u001b[0;32m   1632\u001b[0m compiled_sql, extracted_params, cache_hit \u001b[38;5;241m=\u001b[39m elem\u001b[38;5;241m.\u001b[39m_compile_w_cache(\n\u001b[0;32m   1633\u001b[0m     dialect\u001b[38;5;241m=\u001b[39mdialect,\n\u001b[0;32m   1634\u001b[0m     compiled_cache\u001b[38;5;241m=\u001b[39mcompiled_cache,\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m   1638\u001b[0m     linting\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdialect\u001b[38;5;241m.\u001b[39mcompiler_linting \u001b[38;5;241m|\u001b[39m compiler\u001b[38;5;241m.\u001b[39mWARN_LINTING,\n\u001b[0;32m   1639\u001b[0m )\n\u001b[1;32m-> 1640\u001b[0m ret \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execute_context\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m   1641\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdialect\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1642\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdialect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mexecution_ctx_cls\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_init_compiled\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1643\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcompiled_sql\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1644\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdistilled_parameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1645\u001b[0m \u001b[43m    \u001b[49m\u001b[43mexecution_options\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1646\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcompiled_sql\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1647\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdistilled_parameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1648\u001b[0m \u001b[43m    \u001b[49m\u001b[43melem\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1649\u001b[0m \u001b[43m    \u001b[49m\u001b[43mextracted_params\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1650\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcache_hit\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcache_hit\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1651\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1652\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_events:\n\u001b[0;32m   1653\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdispatch\u001b[38;5;241m.\u001b[39mafter_execute(\n\u001b[0;32m   1654\u001b[0m         \u001b[38;5;28mself\u001b[39m,\n\u001b[0;32m   1655\u001b[0m         elem,\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m   1659\u001b[0m         ret,\n\u001b[0;32m   1660\u001b[0m     )\n",
      "File \u001b[1;32mc:\\UBA\\NL_To_SQL\\.conda\\Lib\\site-packages\\sqlalchemy\\engine\\base.py:1846\u001b[0m, in \u001b[0;36mConnection._execute_context\u001b[1;34m(self, dialect, constructor, statement, parameters, execution_options, *args, **kw)\u001b[0m\n\u001b[0;32m   1844\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exec_insertmany_context(dialect, context)\n\u001b[0;32m   1845\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1846\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_exec_single_context\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m   1847\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdialect\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstatement\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparameters\u001b[49m\n\u001b[0;32m   1848\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\UBA\\NL_To_SQL\\.conda\\Lib\\site-packages\\sqlalchemy\\engine\\base.py:1986\u001b[0m, in \u001b[0;36mConnection._exec_single_context\u001b[1;34m(self, dialect, context, statement, parameters)\u001b[0m\n\u001b[0;32m   1983\u001b[0m     result \u001b[38;5;241m=\u001b[39m context\u001b[38;5;241m.\u001b[39m_setup_result_proxy()\n\u001b[0;32m   1985\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m-> 1986\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_handle_dbapi_exception\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m   1987\u001b[0m \u001b[43m        \u001b[49m\u001b[43me\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstr_statement\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43meffective_parameters\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcursor\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\n\u001b[0;32m   1988\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1990\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
      "File \u001b[1;32mc:\\UBA\\NL_To_SQL\\.conda\\Lib\\site-packages\\sqlalchemy\\engine\\base.py:2353\u001b[0m, in \u001b[0;36mConnection._handle_dbapi_exception\u001b[1;34m(self, e, statement, parameters, cursor, context, is_sub_exec)\u001b[0m\n\u001b[0;32m   2351\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m should_wrap:\n\u001b[0;32m   2352\u001b[0m     \u001b[38;5;28;01massert\u001b[39;00m sqlalchemy_exception \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m-> 2353\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m sqlalchemy_exception\u001b[38;5;241m.\u001b[39mwith_traceback(exc_info[\u001b[38;5;241m2\u001b[39m]) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01me\u001b[39;00m\n\u001b[0;32m   2354\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m   2355\u001b[0m     \u001b[38;5;28;01massert\u001b[39;00m exc_info[\u001b[38;5;241m1\u001b[39m] \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[1;32mc:\\UBA\\NL_To_SQL\\.conda\\Lib\\site-packages\\sqlalchemy\\engine\\base.py:1967\u001b[0m, in \u001b[0;36mConnection._exec_single_context\u001b[1;34m(self, dialect, context, statement, parameters)\u001b[0m\n\u001b[0;32m   1965\u001b[0m                 \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[0;32m   1966\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m evt_handled:\n\u001b[1;32m-> 1967\u001b[0m         \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdialect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdo_execute\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m   1968\u001b[0m \u001b[43m            \u001b[49m\u001b[43mcursor\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstr_statement\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43meffective_parameters\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\n\u001b[0;32m   1969\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1971\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_has_events \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mengine\u001b[38;5;241m.\u001b[39m_has_events:\n\u001b[0;32m   1972\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdispatch\u001b[38;5;241m.\u001b[39mafter_cursor_execute(\n\u001b[0;32m   1973\u001b[0m         \u001b[38;5;28mself\u001b[39m,\n\u001b[0;32m   1974\u001b[0m         cursor,\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m   1978\u001b[0m         context\u001b[38;5;241m.\u001b[39mexecutemany,\n\u001b[0;32m   1979\u001b[0m     )\n",
      "File \u001b[1;32mc:\\UBA\\NL_To_SQL\\.conda\\Lib\\site-packages\\sqlalchemy\\engine\\default.py:924\u001b[0m, in \u001b[0;36mDefaultDialect.do_execute\u001b[1;34m(self, cursor, statement, parameters, context)\u001b[0m\n\u001b[0;32m    923\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdo_execute\u001b[39m(\u001b[38;5;28mself\u001b[39m, cursor, statement, parameters, context\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[1;32m--> 924\u001b[0m     \u001b[43mcursor\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mexecute\u001b[49m\u001b[43m(\u001b[49m\u001b[43mstatement\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparameters\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[1;31mOperationalError\u001b[0m: (sqlite3.OperationalError) no such column: Month\n[SQL: SELECT COUNT(*) FROM apple_daily WHERE Month = 3; ]\n(Background on this error at: https://sqlalche.me/e/20/e3q8)"
     ]
    }
   ],
   "source": [
    "print(agent.invoke(\"Count how many rows in the table are from month 03\"))\n",
    "# Here the query performed doesn't even have the correct syntax for working with an specific month"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(250,)]\n"
     ]
    }
   ],
   "source": [
    "print(agent.invoke(\"Count how many rows in the table are from the year 2023, considering the date format is YYYY-MM-DD\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(0,)]\n"
     ]
    }
   ],
   "source": [
    "print(agent.invoke(\"Count how many rows in the table are from the month 03, considering the date format is YYYY-MM-DD\"))\n",
    "# Here the query performed was Select Count(*) From apple_daily Where Date = '2023-01-03'; which is incorrect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(172.54896,)]\n"
     ]
    }
   ],
   "source": [
    "print(agent.invoke(\"What's the average value of the column Close in the table apple_daily?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(59233213.37296,)]\n"
     ]
    }
   ],
   "source": [
    "print(agent.invoke(\"What's the average value of the columns Close and Volume in the table apple_daily?\"))\n",
    "# Here the query performed was SELECT AVG(Close) + AVG(Volume) FROM apple_daily; which is incorrect"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(172.54896, 59233040.824)]\n"
     ]
    }
   ],
   "source": [
    "print(agent.invoke(\"What's the average value of the column Close and the column Volume in the table apple_daily?\"))\n",
    "# It appears that the agent is not able to handle multiple columns without a very literal separation between them"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Main points of conflict found with the Agent in Baseline:\n",
    "- Fails to use SQLite syntax for working with dates\n",
    "- Fails to deeper understand where more than one column should be present in the query"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}