wip(AI_session_overview_generator): add Part 1

notebooks/AI_session_overview_generator.ipynb

@@ -11,11 +11,24 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "from dotenv import load_dotenv\n",
+    "from IPython.display import Markdown, display\n",
+    "\n",
     "load_dotenv()"
    ]
   },
@@ -30,14 +43,22 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "WARNI [opentelemetry.instrumentation.instrumentor] Attempting to instrument while already instrumented\n"
+      "WARNI [phoenix.session.session] Existing running Phoenix instance detected! Shutting it down and starting a new instance...\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "🌍 To view the Phoenix app in your browser, visit http://localhost:6006/\n",
+      "📖 For more information on how to use Phoenix, check out https://docs.arize.com/phoenix\n"
      ]
     }
    ],
@@ -46,6 +67,10 @@
     "from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter\n",
     "from opentelemetry.sdk import trace as trace_sdk\n",
     "from opentelemetry.sdk.trace.export import SimpleSpanProcessor\n",
+    "import phoenix as px\n",
+    "\n",
+    "px.launch_app()\n",
+    "\n",
     "\n",
     "endpoint = \"http://127.0.0.1:6006/v1/traces\"  # Phoenix receiver address\n",
     "\n",
@@ -56,27 +81,6 @@
     "LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Start SQL Database\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from sqlalchemy import create_engine\n",
-    "from llama_index.core import SQLDatabase\n",
-    "\n",
-    "engine = create_engine('sqlite:///spain_practice_1.db')\n",
-    "\n",
-    "sql_database = SQLDatabase(engine, include_tables=[\"mercedes\"])"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -86,7 +90,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -102,109 +106,125 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Define the Retriever and the Query Engine\n"
+    "### Start SQL Database\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
-    "from llama_index.core.retrievers import NLSQLRetriever\n",
-    "from llama_index.core.query_engine import RetrieverQueryEngine\n",
+    "from sqlalchemy import create_engine\n",
+    "from llama_index.core import SQLDatabase\n",
     "\n",
-    "nl_sql_retriever = NLSQLRetriever(\n",
-    "    sql_database, tables=[\"mercedes\"], return_raw=True\n",
-    ")\n",
+    "engine = create_engine('sqlite:///Spain_2024_FP1.db')\n",
     "\n",
-    "query_engine = RetrieverQueryEngine.from_args(nl_sql_retriever)"
+    "sql_database = SQLDatabase(engine)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Start retrieving information to construct the report\n"
+    "# Part 1: Text-to-SQL Query Engine\n",
+    "\n",
+    "Once we have constructed our SQL database, we can use the `NLSQLTableQueryEngine` to construct natural language queries that are synthesized into SQL queries.\n",
+    "\n",
+    "Note that we need to specify the tables we want to use with this query engine. If we don't the query engine will pull all the schema context, which could overflow the context window of the LLM.\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [
     {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
-      "I0000 00:00:1723244252.910416 2380678 fork_posix.cc:77] Other threads are currently calling into gRPC, skipping fork() handlers\n"
-     ]
+     "data": {
+      "text/markdown": [
+       "In the first free practice session, driver HAM started on the HARD compound. His speed varied throughout the session, with his fastest lap being 77.632 seconds on lap 2. He showed consistent speed on the straights, reaching up to 316 km/h, but struggled with lower speeds in some corners. Overall, his performance on the HARD compound seemed solid, with room for improvement in certain areas."
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
     }
    ],
    "source": [
-    "response = query_engine.query(\n",
-    "    \"Which driver was faster on average on sector 1?\"\n",
-    ")"
+    "from llama_index.core.query_engine import NLSQLTableQueryEngine\n",
+    "\n",
+    "hamilton_query_engine = NLSQLTableQueryEngine(\n",
+    "    sql_database=sql_database, tables=[\"hamilton\"])\n",
+    "query_str = \"This database contains the data for the driver HAM in the first free practice, based on this data provide a short overview on how was his speed and time taking into account his compound\"\n",
+    "response = hamilton_query_engine.query(query_str)\n",
+    "\n",
+    "display(Markdown(f\"{response}\"))"
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": 9,
+   "cell_type": "markdown",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      ">>>  RUS was faster on average on sector 1.\n"
-     ]
-    }
-   ],
    "source": [
-    "print(\">>> \", str(response))"
+    "This query engine should be used in any case where you can specify the tables you want to query over beforehand, or the total size of all the table schema plus the rest of the prompt fits your context window.\n",
+    "\n",
+    "# Part 2: Query-Time Retrieval of Tables for Text-to-SQL\n",
+    "\n",
+    "If we don't know ahead of time which table we would like to use, and the total size of the table schema overflows your context window size, we should store the table schema in an index so that during query time we can retrieve the right schema.\n",
+    "\n",
+    "The way we can do this is using the `SQLTableNodeMapping` object, which takes in a `SQLDatabase` and produces a Node object for each `SQLTableSchema` object passed into the `ObjectIndex` constructor.\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      ">>>  HAM\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "response2 = query_engine.query(\n",
-    "    \"Which driver was faster on average on sector 2?\"\n",
+    "from llama_index.core.indices.struct_store.sql_query import SQLTableRetrieverQueryEngine\n",
+    "from llama_index.core.objects import (\n",
+    "    SQLTableNodeMapping,\n",
+    "    ObjectIndex,\n",
+    "    SQLTableSchema,\n",
     ")\n",
+    "from llama_index.core import VectorStoreIndex\n",
     "\n",
-    "print(\">>> \", str(response2))"
+    "table_node_mapping = SQLTableNodeMapping(sql_database)\n",
+    "table_schema_objs = [\n",
+    "    (SQLTableSchema(table_name=\"hamilton\", context_str=\"This table contains the lap time for the first Hamilton's free practice at Spain.\"))]\n",
+    "\n",
+    "obj_index = ObjectIndex.from_objects(objects=table_schema_objs,\n",
+    "                                     object_mapping=table_node_mapping,\n",
+    "                                     index_cls=VectorStoreIndex)\n",
+    "\n",
+    "table_retriever = obj_index.as_retriever(similarity_top_k=1)\n",
+    "\n",
+    "query_engine = SQLTableRetrieverQueryEngine(\n",
+    "    sql_database=sql_database, table_retriever=table_retriever)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Define the Retriever and the Query Engine\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 7,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      ">>>  The driver HAM is 28 years old and has a lap time of 98.0705652173913 seconds with a top speed of 289.0 km/h.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "response3 = query_engine.query(\n",
-    "    \"Write a summary about the driver HAM\"\n",
+    "from llama_index.core.retrievers import NLSQLRetriever\n",
+    "from llama_index.core.query_engine import RetrieverQueryEngine\n",
+    "\n",
+    "nl_sql_retriever = NLSQLRetriever(\n",
+    "    sql_database, tables=[\"mercedes\"], return_raw=True\n",
     ")\n",
     "\n",
-    "print(\">>> \", str(response3))"
+    "query_engine = RetrieverQueryEngine.from_args(nl_sql_retriever)"
    ]
   }
  ],