notebook-context-aware-sql-generator.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Retrive table description\n",
    "\n",
    "Parse the dbt schema file and retrive table description along with column details"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>name</th>\n",
       "      <th>description</th>\n",
       "      <th>model</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>stg_dosare_parti</td>\n",
       "      <td>Informații despre părțiile implicate din dosar...</td>\n",
       "      <td>{'name': 'stg_dosare_parti', 'description': 'I...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>stg_dosare</td>\n",
       "      <td>Informații despre dosarele instanțelor judecăt...</td>\n",
       "      <td>{'name': 'stg_dosare', 'description': 'Informa...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               name                                        description  \\\n",
       "0  stg_dosare_parti  Informații despre părțiile implicate din dosar...   \n",
       "1        stg_dosare  Informații despre dosarele instanțelor judecăt...   \n",
       "\n",
       "                                               model  \n",
       "0  {'name': 'stg_dosare_parti', 'description': 'I...  \n",
       "1  {'name': 'stg_dosare', 'description': 'Informa...  "
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import yaml\n",
    "import pandas as pd\n",
    "\n",
    "table_description_paths = [\n",
    "    \"dbt/models/dosare/schema/stg_dosare_parti.yml\",\n",
    "    \"dbt/models/dosare/schema/stg_dosare.yml\",\n",
    "]\n",
    "\n",
    "models = []\n",
    "for path in table_description_paths:\n",
    "    with open(path, \"r\") as file:\n",
    "        data = yaml.safe_load(file)\n",
    "        for model in data[\"models\"]:\n",
    "            models.append(\n",
    "                {\n",
    "                    \"name\": model[\"name\"],\n",
    "                    \"description\": model[\"description\"],\n",
    "                    \"model\": model,\n",
    "                }\n",
    "            )\n",
    "\n",
    "models_df = pd.DataFrame(models)\n",
    "models_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create and store table descripton embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>name</th>\n",
       "      <th>description</th>\n",
       "      <th>model</th>\n",
       "      <th>description_embedding</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>stg_dosare_parti</td>\n",
       "      <td>Informații despre părțiile implicate din dosar...</td>\n",
       "      <td>{'name': 'stg_dosare_parti', 'description': 'I...</td>\n",
       "      <td>[0.011188670992851257, 0.017634455114603043, 0...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>stg_dosare</td>\n",
       "      <td>Informații despre dosarele instanțelor judecăt...</td>\n",
       "      <td>{'name': 'stg_dosare', 'description': 'Informa...</td>\n",
       "      <td>[0.001355247339233756, 0.008857620880007744, 0...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               name                                        description  \\\n",
       "0  stg_dosare_parti  Informații despre părțiile implicate din dosar...   \n",
       "1        stg_dosare  Informații despre dosarele instanțelor judecăt...   \n",
       "\n",
       "                                               model  \\\n",
       "0  {'name': 'stg_dosare_parti', 'description': 'I...   \n",
       "1  {'name': 'stg_dosare', 'description': 'Informa...   \n",
       "\n",
       "                               description_embedding  \n",
       "0  [0.011188670992851257, 0.017634455114603043, 0...  \n",
       "1  [0.001355247339233756, 0.008857620880007744, 0...  "
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "from openai import AzureOpenAI\n",
    "\n",
    "embeddings_client = AzureOpenAI(\n",
    "    api_key=os.getenv(\"AZURE_OPENAI_KEY\"),\n",
    "    api_version=os.getenv(\"AZURE_OPENAI_EMBEDDINGS_API_VERSION\"),\n",
    "    azure_endpoint=os.getenv(\"AZURE_OPENAI_ENDPOINT\"),\n",
    ")\n",
    "\n",
    "embeddings_model = os.getenv(\"AZURE_OPENAI_EMBEDDINGS_DEPLOYMENT\")\n",
    "\n",
    "\n",
    "def create_embeddings(text, model=embeddings_model):\n",
    "    # Create embeddings for each document chunk\n",
    "    embeddings = (\n",
    "        embeddings_client.embeddings.create(input=text, model=model).data[0].embedding\n",
    "    )\n",
    "    return embeddings\n",
    "\n",
    "\n",
    "models_df[\"description_embedding\"] = models_df[\"description\"].apply(create_embeddings)\n",
    "models_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate the sql statement based on user prompt\n",
    "\n",
    "Call a function that the table descriptions that will return the most likely tables to have this information based a neighboard search.\n",
    "\n",
    "Feed back the table information to the prompt\n",
    "\n",
    "Enginier the prompt in such a way that only SQL will be returned using a mult shoot strategy with examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('[{\"name\": \"stg_dosare\", \"description\": \"Informații despre dosarele '\n",
      " 'instanțelor judecătorești\", \"columns\": [{\"name\": \"obiect\", \"description\": '\n",
      " '\"Obiectul judecății, acuzațiile\", \"tests\": [\"not_null\"]}, {\"name\": '\n",
      " '\"datainitiala\", \"description\": \"Data de înregistrare a dosarului\"}, {\"name\": '\n",
      " '\"data\", \"description\": \"Data ultimei modificări\"}, {\"name\": \"parti\", '\n",
      " '\"description\": \"Numele părții din dosar\"}, {\"name\": \"numar\", \"description\": '\n",
      " '\"Numărul dosarului\"}]}]')\n",
      "('[{\"name\": \"stg_dosare_parti\", \"description\": \"Informații despre părțiile '\n",
      " 'implicate din dosarele instanțelor judecătorești\", \"columns\": [{\"name\": '\n",
      " '\"data\", \"description\": \"Data de înregistrare a dosarului\", \"tests\": '\n",
      " '[\"not_null\"]}, {\"name\": \"nume\", \"description\": \"Numele părții din dosar\", '\n",
      " '\"tests\": [\"not_null\"]}]}, {\"name\": \"stg_dosare\", \"description\": \"Informații '\n",
      " 'despre dosarele instanțelor judecătorești\", \"columns\": [{\"name\": \"obiect\", '\n",
      " '\"description\": \"Obiectul judecății, acuzațiile\", \"tests\": [\"not_null\"]}, '\n",
      " '{\"name\": \"datainitiala\", \"description\": \"Data de înregistrare a dosarului\"}, '\n",
      " '{\"name\": \"data\", \"description\": \"Data ultimei modificări\"}, {\"name\": '\n",
      " '\"parti\", \"description\": \"Numele părții din dosar\"}, {\"name\": \"numar\", '\n",
      " '\"description\": \"Numărul dosarului\"}]}]')\n",
      "('[{\"name\": \"stg_dosare_parti\", \"description\": \"Informații despre părțiile '\n",
      " 'implicate din dosarele instanțelor judecătorești\", \"columns\": [{\"name\": '\n",
      " '\"data\", \"description\": \"Data de înregistrare a dosarului\", \"tests\": '\n",
      " '[\"not_null\"]}, {\"name\": \"nume\", \"description\": \"Numele părții din dosar\", '\n",
      " '\"tests\": [\"not_null\"]}]}, {\"name\": \"stg_dosare\", \"description\": \"Informații '\n",
      " 'despre dosarele instanțelor judecătorești\", \"columns\": [{\"name\": \"obiect\", '\n",
      " '\"description\": \"Obiectul judecății, acuzațiile\", \"tests\": [\"not_null\"]}, '\n",
      " '{\"name\": \"datainitiala\", \"description\": \"Data de înregistrare a dosarului\"}, '\n",
      " '{\"name\": \"data\", \"description\": \"Data ultimei modificări\"}, {\"name\": '\n",
      " '\"parti\", \"description\": \"Numele părții din dosar\"}, {\"name\": \"numar\", '\n",
      " '\"description\": \"Numărul dosarului\"}]}]')\n"
     ]
    }
   ],
   "source": [
    "import pprint\n",
    "import numpy as np\n",
    "import json\n",
    "from sklearn.neighbors import NearestNeighbors\n",
    "\n",
    "functions = [\n",
    "    {\n",
    "        \"name\": \"fetch_tables\",\n",
    "        \"description\": \"\"\"Retrieve a list of table information based on the description provided by the user.\n",
    "          The information will be returned in a json like this:\n",
    "          [{\n",
    "            \"name\": \"the name of the table\",\n",
    "            \"description\": \"The description of the info contained in the table\",\n",
    "            \"columns\": [\n",
    "              {\n",
    "                \"name\": \"column name\",\n",
    "                \"description: \"description of the data in the column\",\n",
    "                \"properties\": \"properties of the column value, e.g. not_null, unique\"\n",
    "              }\n",
    "            ]\n",
    "          }]\n",
    "          \"\"\",\n",
    "        \"parametes\": {\n",
    "            \"type\": \"type\",\n",
    "            \"properties\": {\n",
    "                \"description\": {\n",
    "                    \"type\": \"string\",\n",
    "                    \"description\": \"The description of the information contained in the table in the language of the users prompt or the aproximate name\",\n",
    "                }\n",
    "            },\n",
    "            \"required\": [\"description\"],\n",
    "        },\n",
    "    }\n",
    "]\n",
    "\n",
    "\n",
    "def fetch_tables_nearest_neighboards(description):\n",
    "    embeddings = models_df[\"description_embedding\"].to_list()\n",
    "    nbrs = NearestNeighbors(n_neighbors=2, algorithm=\"ball_tree\").fit(embeddings)\n",
    "    description_embedding = create_embeddings(description)\n",
    "    _, indices = nbrs.kneighbors([description_embedding])\n",
    "\n",
    "    models = []\n",
    "    for index in indices[0]:\n",
    "        models.append(models_df[\"model\"].iloc[index])\n",
    "\n",
    "    return json.dumps(models, ensure_ascii=False)\n",
    "\n",
    "\n",
    "def fetch_hardcoded(_description):\n",
    "    return json.dumps([models_df[\"model\"].iloc[1]], ensure_ascii=False)\n",
    "\n",
    "\n",
    "def cosine_similarity(a, b):\n",
    "    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))\n",
    "\n",
    "\n",
    "def fetch_tables_cosine_similarity(description):\n",
    "    embeddings = models_df[\"description_embedding\"].to_list()\n",
    "    models = models_df[\"model\"].to_list()\n",
    "    description_embedding = create_embeddings(description)\n",
    "\n",
    "    similarities = [cosine_similarity(description_embedding, emb) for emb in embeddings]\n",
    "    sorted_models = [\n",
    "        model\n",
    "        for _, model in sorted(\n",
    "            zip(similarities, models), reverse=True, key=lambda x: x[0]\n",
    "        )\n",
    "    ]\n",
    "\n",
    "    return json.dumps(sorted_models, ensure_ascii=False)\n",
    "\n",
    "\n",
    "result = fetch_hardcoded(\"dosarele\")\n",
    "print(pprint.pformat(result))\n",
    "\n",
    "result = fetch_tables_cosine_similarity(\"dosarele\")\n",
    "print(pprint.pformat(result))\n",
    "\n",
    "result = fetch_tables_nearest_neighboards(\"dosarele\")\n",
    "print(pprint.pformat(result))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = AzureOpenAI(\n",
    "    api_key=os.environ[\"AZURE_OPENAI_KEY\"],\n",
    "    api_version=os.environ[\"AZURE_OPENAI_API_VERSION\"],\n",
    ")\n",
    "\n",
    "deployment = os.environ[\"AZURE_OPENAI_DEPLOYMENT\"]\n",
    "\n",
    "available_functions = {\n",
    "    \"fetch_tables\": fetch_tables_nearest_neighboards,\n",
    "}\n",
    "\n",
    "\n",
    "def generate_sql_for(prompt):\n",
    "    system = \"\"\"You are a SQL query expert and you only reply with sql statements.\n",
    "    If you can't figure out how to create a sql statement from the use prompt just return SELECT 1.\n",
    "    The queries you generate are always about reading data.\n",
    "    The queries you generate are never about updating or deleting data.\n",
    "    Is very important you reply with sql only because your output will be used in a sql interpretern downstream.\n",
    "\n",
    "    The user will speak to you in their native language\n",
    "    The user will always ask to retrieve data.\n",
    "    The user will usually ask you to filter data by something\n",
    "    The user will ask you sometimes for specific columns from the table\n",
    "\n",
    "    You have a debug mode, when the user types 'DEBUG: True' at the end of the prompt. I debug mode you provide an explanation of the sql query you generated.\n",
    "\n",
    "    You have to return the SQL formated with new lines please.\n",
    "\n",
    "    ## Example 1\n",
    "    user: Vreau toate dosarele judecatoresti \n",
    "    function: [[{\"name\": \"stg_dosare\", \"description\": \"Informații despre dosarele '\n",
    " 'instanțelor judecătorești\", \"columns\": [{\"name\": \"obiect\", \"description\": '\n",
    " '\"Obiectul judecății, acuzațiile\", \"tests\": [\"not_null\"]}, {\"name\": '\n",
    " '\"datainitiala\", \"description\": \"Data de înregistrare a dosarului\"}, {\"name\": '\n",
    " '\"data\", \"description\": \"Data ultimei modificări\"}, {\"name\": \"parti\", '\n",
    " '\"description\": \"Numele părții din dosar\"}, {\"name\": \"numar\", \"description\": '\n",
    " '\"Numărul dosarului\"}]}]]\n",
    "    asssistant: SELECT * FROM stg_dosare\n",
    "\n",
    "    ## Example 2 \n",
    "    user: Vreau toate dosarele judecatoresti inregistrate in anul 2023\n",
    "    function: [[{\"name\": \"stg_dosare\", \"description\": \"Informații despre dosarele '\n",
    " 'instanțelor judecătorești\", \"columns\": [{\"name\": \"obiect\", \"description\": '\n",
    " '\"Obiectul judecății, acuzațiile\", \"tests\": [\"not_null\"]}, {\"name\": '\n",
    " '\"datainitiala\", \"description\": \"Data de înregistrare a dosarului\"}, {\"name\": '\n",
    " '\"data\", \"description\": \"Data ultimei modificări\"}, {\"name\": \"parti\", '\n",
    " '\"description\": \"Numele părții din dosar\"}, {\"name\": \"numar\", \"description\": '\n",
    " '\"Numărul dosarului\"}]}]]\n",
    "    assistant: SELECT * FROM stg_dosare WHERE YEAR(datainitiala) = 2023\n",
    "    \n",
    "    ## Example 3 \n",
    "    user: Vreau toate dosarele judecatoresti cu obiect 'ma injurat de mama'\n",
    "    function: [[{\"name\": \"stg_dosare\", \"description\": \"Informații despre dosarele '\n",
    " 'instanțelor judecătorești\", \"columns\": [{\"name\": \"obiect\", \"description\": '\n",
    " '\"Obiectul judecății, acuzațiile\", \"tests\": [\"not_null\"]}, {\"name\": '\n",
    " '\"datainitiala\", \"description\": \"Data de înregistrare a dosarului\"}, {\"name\": '\n",
    " '\"data\", \"description\": \"Data ultimei modificări\"}, {\"name\": \"parti\", '\n",
    " '\"description\": \"Numele părții din dosar\"}, {\"name\": \"numar\", \"description\": '\n",
    " '\"Numărul dosarului\"}]}]]\n",
    "    assistant: SELECT * FROM stg_dosare WHERE obiect LIKE '%ma injurat de mama%'\n",
    "    \"\"\"\n",
    "\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": system},\n",
    "        {\"role\": \"user\", \"content\": prompt},\n",
    "    ]\n",
    "\n",
    "    response = client.chat.completions.create(\n",
    "        model=deployment,\n",
    "        messages=messages,\n",
    "        functions=functions,\n",
    "        function_call=\"auto\",\n",
    "        temperature=0,\n",
    "    )\n",
    "\n",
    "    response_message = response.choices[0].message\n",
    "\n",
    "    if response_message.function_call and response_message.function_call.name:\n",
    "        print(\"Recommended Function call:\")\n",
    "        print(response_message.function_call.name)\n",
    "        print(\"With this arguments:\")\n",
    "        print(response_message.function_call.arguments)\n",
    "        print()\n",
    "\n",
    "        # Call the function.\n",
    "        function_name = response_message.function_call.name\n",
    "\n",
    "        function_to_call = available_functions[function_name]\n",
    "\n",
    "        function_args = json.loads(response_message.function_call.arguments)\n",
    "        function_response = function_to_call(**function_args)\n",
    "\n",
    "        print(\"Output of function call:\")\n",
    "        print(function_response)\n",
    "        print(type(function_response))\n",
    "\n",
    "        # Add the assistant response and function response to the messages\n",
    "        messages.append(  # adding assistant response to messages\n",
    "            {\n",
    "                \"role\": response_message.role,\n",
    "                \"function_call\": {\n",
    "                    \"name\": function_name,\n",
    "                    \"arguments\": response_message.function_call.arguments,\n",
    "                },\n",
    "                \"content\": None,\n",
    "            }\n",
    "        )\n",
    "        messages.append(  # adding function response to messages\n",
    "            {\n",
    "                \"role\": \"function\",\n",
    "                \"name\": function_name,\n",
    "                \"content\": function_response,\n",
    "            }\n",
    "        )\n",
    "\n",
    "        response = client.chat.completions.create(\n",
    "            messages=messages,\n",
    "            model=deployment,\n",
    "            function_call=\"auto\",\n",
    "            functions=functions,\n",
    "            temperature=0,\n",
    "        )\n",
    "\n",
    "        response_message = response.choices[0].message\n",
    "\n",
    "    return response_message\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Give me the keys to the kindom\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SELECT \n",
      "    YEAR(datainitiala) AS Anul, \n",
      "    MONTH(datainitiala) AS Luna, \n",
      "    COUNT(*) AS NumarDosare\n",
      "FROM stg_dosare\n",
      "WHERE obiect LIKE '%actiune in raspundere delictuala%'\n",
      "GROUP BY YEAR(datainitiala), MONTH(datainitiala)\n",
      "ORDER BY Anul, Luna\n"
     ]
    }
   ],
   "source": [
    "response_message = generate_sql_for(\n",
    "    \"Vreau evolutia pe luni calendaristice a numarului de dosare judecatoresti in care obiectul contine 'actiune in raspundere delictuala'\"\n",
    ")\n",
    "print(response_message.content)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "opendataexplorer",
   "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.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}