{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generating Session Summary with LLMs\n",
    "\n",
    "This project utilizes LlamaIndex and AI technology to analyze Formula 1 car data and generate session summaries. It is particularly beneficial for race engineers seeking detailed insights and performance analysis during races.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "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()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optional: Setup Observability\n",
    "\n",
    "Here we will use our Arize Phoenix integration to view traces through the query engine. It will be available at http://localhost:6006\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "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"
     ]
    }
   ],
   "source": [
    "from openinference.instrumentation.llama_index import LlamaIndexInstrumentor\n",
    "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",
    "tracer_provider = trace_sdk.TracerProvider()\n",
    "tracer_provider.add_span_processor(\n",
    "    SimpleSpanProcessor(OTLPSpanExporter(endpoint)))\n",
    "\n",
    "LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the LLM and the Embedding\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core import Settings\n",
    "\n",
    "Settings.llm = OpenAI(model=\"gpt-3.5-turbo\", temperature=0)\n",
    "Settings.embed_model = OpenAIEmbedding(model=\"text-embedding-ada-002\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start SQL Database\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sqlalchemy import create_engine\n",
    "from llama_index.core import SQLDatabase\n",
    "\n",
    "engine = create_engine('sqlite:///Spain_2024_FP1.db')\n",
    "\n",
    "sql_database = SQLDatabase(engine)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 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": 9,
   "metadata": {},
   "outputs": [
    {
     "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": [
    "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": "markdown",
   "metadata": {},
   "source": [
    "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "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": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "query_engine = RetrieverQueryEngine.from_args(nl_sql_retriever)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama",
   "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
}