{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8d10861f-a550-4443-bc63-4ce2ae13b841",
   "metadata": {},
   "source": [
    "# Infino\n",
    "\n",
    "This example shows how one can track the following while calling OpenAI and ChatOpenAI models via `LangChain` and [Infino](https://github.com/infinohq/infino):\n",
    "\n",
    "* prompt input,\n",
    "* response from `ChatGPT` or any other `LangChain` model,\n",
    "* latency,\n",
    "* errors,\n",
    "* number of tokens consumed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64d14c88-b71c-4524-ab1b-4250a7dbb62b",
   "metadata": {},
   "source": [
    "## Initializing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ed46c894-caa6-49b2-85d1-f275374fa308",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install necessary dependencies.\n",
    "!pip install -q infinopy\n",
    "!pip install -q matplotlib\n",
    "!pip install -q tiktoken"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3a5a0976-9953-41d8-880c-eb3f2992e936",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime as dt\n",
    "import json\n",
    "from langchain.llms import OpenAI\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.dates as md\n",
    "import os\n",
    "import time\n",
    "import sys\n",
    "\n",
    "from infinopy import InfinoClient\n",
    "from langchain.callbacks import InfinoCallbackHandler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f90210d-c805-4a0c-81e4-d5298942afc4",
   "metadata": {},
   "source": [
    "## Start Infino server, initialize the Infino client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "748b9858-5145-4351-976a-ca2d54e836a6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a1159e99c6bdb3101139157acee6aba7ae9319375e77ab6fbc79beff75abeca3\n"
     ]
    }
   ],
   "source": [
    "# Start server using the Infino docker image.\n",
    "!docker run --rm --detach --name infino-example -p 3000:3000 infinohq/infino:latest\n",
    "\n",
    "# Create Infino client.\n",
    "client = InfinoClient()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6b81cda-b841-43ee-8c5e-b1576555765f",
   "metadata": {},
   "source": [
    "## Read the questions dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b659fd0c-0d8c-470e-8b6c-867a117f2a27",
   "metadata": {},
   "outputs": [],
   "source": [
    "# These are a subset of questions from Stanford's QA dataset -\n",
    "# https://rajpurkar.github.io/SQuAD-explorer/\n",
    "data = \"\"\"In what country is Normandy located?\n",
    "When were the Normans in Normandy?\n",
    "From which countries did the Norse originate?\n",
    "Who was the Norse leader?\n",
    "What century did the Normans first gain their separate identity?\n",
    "Who gave their name to Normandy in the 1000's and 1100's\n",
    "What is France a region of?\n",
    "Who did King Charles III swear fealty to?\n",
    "When did the Frankish identity emerge?\n",
    "Who was the duke in the battle of Hastings?\n",
    "Who ruled the duchy of Normandy\n",
    "What religion were the Normans\n",
    "What type of major impact did the Norman dynasty have on modern Europe?\n",
    "Who was famed for their Christian spirit?\n",
    "Who assimilted the Roman language?\n",
    "Who ruled the country of Normandy?\n",
    "What principality did William the conquerer found?\n",
    "What is the original meaning of the word Norman?\n",
    "When was the Latin version of the word Norman first recorded?\n",
    "What name comes from the English words Normans/Normanz?\"\"\"\n",
    "\n",
    "questions = data.split(\"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dce1b820-3f1a-4b94-b848-4c6032cadc18",
   "metadata": {},
   "source": [
    "## Example 1: LangChain OpenAI Q&A; Publish metrics and logs to Infino"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d5cebf35-2d10-48b8-ab11-c4a574c595d2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "In what country is Normandy located?\n",
      "generations=[[Generation(text='\\n\\nNormandy is located in France.', generation_info={'finish_reason': 'stop', 'logprobs': None})]] llm_output={'token_usage': {'total_tokens': 16, 'prompt_tokens': 7, 'completion_tokens': 9}, 'model_name': 'text-davinci-003'} run=[RunInfo(run_id=UUID('67a516e3-d48a-4e83-92ba-a139079bd3b1'))]\n",
      "When were the Normans in Normandy?\n",
      "generations=[[Generation(text='\\n\\nThe Normans first settled in Normandy in the late 9th century.', generation_info={'finish_reason': 'stop', 'logprobs': None})]] llm_output={'token_usage': {'total_tokens': 24, 'prompt_tokens': 8, 'completion_tokens': 16}, 'model_name': 'text-davinci-003'} run=[RunInfo(run_id=UUID('6417a773-c863-4942-9607-c8a0c5d486e7'))]\n",
      "From which countries did the Norse originate?\n",
      "generations=[[Generation(text='\\n\\nThe Norse originated from Scandinavia, which includes the modern-day countries of Norway, Sweden, and Denmark.', generation_info={'finish_reason': 'stop', 'logprobs': None})]] llm_output={'token_usage': {'total_tokens': 32, 'prompt_tokens': 8, 'completion_tokens': 24}, 'model_name': 'text-davinci-003'} run=[RunInfo(run_id=UUID('70547d72-7925-454e-97fb-5539f8788c3f'))]\n",
      "Who was the Norse leader?\n",
      "generations=[[Generation(text='\\n\\nThe most famous Norse leader was the legendary Viking king Ragnar Lodbrok. He was a legendary Viking hero and ruler who is said to have lived in the 9th century. He is known for his legendary exploits, including leading a Viking raid on Paris in 845.', generation_info={'finish_reason': 'stop', 'logprobs': None})]] llm_output={'token_usage': {'total_tokens': 62, 'prompt_tokens': 6, 'completion_tokens': 56}, 'model_name': 'text-davinci-003'} run=[RunInfo(run_id=UUID('04500e37-44ab-4e56-9017-76fe8c19e2ca'))]\n",
      "What century did the Normans first gain their separate identity?\n",
      "generations=[[Generation(text='\\n\\nThe Normans first gained their separate identity in the 11th century.', generation_info={'finish_reason': 'stop', 'logprobs': None})]] llm_output={'token_usage': {'total_tokens': 28, 'prompt_tokens': 12, 'completion_tokens': 16}, 'model_name': 'text-davinci-003'} run=[RunInfo(run_id=UUID('adf319b7-1022-40df-9afe-1d65f869d83d'))]\n",
      "Who gave their name to Normandy in the 1000's and 1100's\n",
      "generations=[[Generation(text='\\n\\nThe Normans, a people from northern France, gave their name to Normandy in the 1000s and 1100s. The Normans were descendants of Vikings who had settled in the region in the late 800s.', generation_info={'finish_reason': 'stop', 'logprobs': None})]] llm_output={'token_usage': {'total_tokens': 57, 'prompt_tokens': 13, 'completion_tokens': 44}, 'model_name': 'text-davinci-003'} run=[RunInfo(run_id=UUID('1a0503bc-d033-4b69-a5fa-5e1796566133'))]\n",
      "What is France a region of?\n",
      "generations=[[Generation(text='\\n\\nFrance is a region of Europe.', generation_info={'finish_reason': 'stop', 'logprobs': None})]] llm_output={'token_usage': {'total_tokens': 16, 'prompt_tokens': 7, 'completion_tokens': 9}, 'model_name': 'text-davinci-003'} run=[RunInfo(run_id=UUID('7485d954-1c14-4dff-988a-25a0aa0871cc'))]\n",
      "Who did King Charles III swear fealty to?\n",
      "generations=[[Generation(text='\\n\\nKing Charles III swore fealty to King Philip II of Spain.', generation_info={'finish_reason': 'stop', 'logprobs': None})]] llm_output={'token_usage': {'total_tokens': 25, 'prompt_tokens': 10, 'completion_tokens': 15}, 'model_name': 'text-davinci-003'} run=[RunInfo(run_id=UUID('292c7143-4a08-43cd-a1e1-42cb1f594f33'))]\n",
      "When did the Frankish identity emerge?\n",
      "generations=[[Generation(text='\\n\\nThe Frankish identity began to emerge in the late 5th century, when the Franks began to expand their power and influence in the region. The Franks were a Germanic tribe that had settled in the area of modern-day France and Germany. They eventually established the Merovingian dynasty, which ruled much of Western Europe from the mid-6th century until 751.', generation_info={'finish_reason': 'stop', 'logprobs': None})]] llm_output={'token_usage': {'total_tokens': 85, 'prompt_tokens': 8, 'completion_tokens': 77}, 'model_name': 'text-davinci-003'} run=[RunInfo(run_id=UUID('3d9475c2-931e-4217-8bc3-b3e970e7597c'))]\n",
      "Who was the duke in the battle of Hastings?\n",
      "generations=[[Generation(text='\\n\\nThe Duke of Normandy, William the Conqueror, was the leader of the Norman forces at the Battle of Hastings in 1066.', generation_info={'finish_reason': 'stop', 'logprobs': None})]] llm_output={'token_usage': {'total_tokens': 39, 'prompt_tokens': 11, 'completion_tokens': 28}, 'model_name': 'text-davinci-003'} run=[RunInfo(run_id=UUID('b8f84619-ea5f-4c18-b411-b62194f36fe0'))]\n"
     ]
    }
   ],
   "source": [
    "# Set your key here.\n",
    "# os.environ[\"OPENAI_API_KEY\"] = \"YOUR_API_KEY\"\n",
    "\n",
    "# Create callback handler. This logs latency, errors, token usage, prompts as well as prompt responses to Infino.\n",
    "handler = InfinoCallbackHandler(\n",
    "    model_id=\"test_openai\", model_version=\"0.1\", verbose=False\n",
    ")\n",
    "\n",
    "# Create LLM.\n",
    "llm = OpenAI(temperature=0.1)\n",
    "\n",
    "# Number of questions to ask the OpenAI model. We limit to a short number here to save $$ while running this demo.\n",
    "num_questions = 10\n",
    "\n",
    "questions = questions[0:num_questions]\n",
    "for question in questions:\n",
    "    print(question)\n",
    "\n",
    "    # We send the question to OpenAI API, with Infino callback.\n",
    "    llm_result = llm.generate([question], callbacks=[handler])\n",
    "    print(llm_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b68ec697-c922-4fd9-aad1-f49c6ac24e8a",
   "metadata": {},
   "source": [
    "## Create Metric Charts\n",
    "\n",
    "We now use matplotlib to create graphs of latency, errors and tokens consumed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2eeeb6aa-e7b4-4a12-bbc6-6c4f1459810a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper function to create a graph using matplotlib.\n",
    "def plot(data, title):\n",
    "    data = json.loads(data)\n",
    "\n",
    "    # Extract x and y values from the data\n",
    "    timestamps = [item[\"time\"] for item in data]\n",
    "    dates = [dt.datetime.fromtimestamp(ts) for ts in timestamps]\n",
    "    y = [item[\"value\"] for item in data]\n",
    "\n",
    "    plt.rcParams[\"figure.figsize\"] = [6, 4]\n",
    "    plt.subplots_adjust(bottom=0.2)\n",
    "    plt.xticks(rotation=25)\n",
    "    ax = plt.gca()\n",
    "    xfmt = md.DateFormatter(\"%Y-%m-%d %H:%M:%S\")\n",
    "    ax.xaxis.set_major_formatter(xfmt)\n",
    "\n",
    "    # Create the plot\n",
    "    plt.plot(dates, y)\n",
    "\n",
    "    # Set labels and title\n",
    "    plt.xlabel(\"Time\")\n",
    "    plt.ylabel(\"Value\")\n",
    "    plt.title(title)\n",
    "\n",
    "    plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f078c612-89e0-4a1d-b1a8-bf36b664a10e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 39504 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 20548 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 33192 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 38140 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 38320 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "response = client.search_ts(\"__name__\", \"latency\", 0, int(time.time()))\n",
    "plot(response.text, \"Latency\")\n",
    "\n",
    "response = client.search_ts(\"__name__\", \"error\", 0, int(time.time()))\n",
    "plot(response.text, \"Errors\")\n",
    "\n",
    "response = client.search_ts(\"__name__\", \"prompt_tokens\", 0, int(time.time()))\n",
    "plot(response.text, \"Prompt Tokens\")\n",
    "\n",
    "response = client.search_ts(\"__name__\", \"completion_tokens\", 0, int(time.time()))\n",
    "plot(response.text, \"Completion Tokens\")\n",
    "\n",
    "response = client.search_ts(\"__name__\", \"total_tokens\", 0, int(time.time()))\n",
    "plot(response.text, \"Total Tokens\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3d61822-1781-4bc6-97a2-2abc5c2b2e75",
   "metadata": {},
   "source": [
    "## Full text query on prompt or prompt outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a0f051f0-e2bc-44e7-8dfb-bfd5bbd0fc9f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Results for normandy : [{\"time\":1696947743,\"fields\":{\"prompt_response\":\"\\n\\nThe Normans, a people from northern France, gave their name to Normandy in the 1000s and 1100s. The Normans were descendants of Vikings who had settled in the region in the late 800s.\"},\"text\":\"\\n\\nThe Normans, a people from northern France, gave their name to Normandy in the 1000s and 1100s. The Normans were descendants of Vikings who had settled in the region in the late 800s.\"},{\"time\":1696947740,\"fields\":{\"prompt\":\"Who gave their name to Normandy in the 1000's and 1100's\"},\"text\":\"Who gave their name to Normandy in the 1000's and 1100's\"},{\"time\":1696947733,\"fields\":{\"prompt_response\":\"\\n\\nThe Normans first settled in Normandy in the late 9th century.\"},\"text\":\"\\n\\nThe Normans first settled in Normandy in the late 9th century.\"},{\"time\":1696947732,\"fields\":{\"prompt_response\":\"\\n\\nNormandy is located in France.\"},\"text\":\"\\n\\nNormandy is located in France.\"},{\"time\":1696947731,\"fields\":{\"prompt\":\"In what country is Normandy located?\"},\"text\":\"In what country is Normandy located?\"}]\n",
      "===\n",
      "Results for king charles III : [{\"time\":1696947745,\"fields\":{\"prompt_response\":\"\\n\\nKing Charles III swore fealty to King Philip II of Spain.\"},\"text\":\"\\n\\nKing Charles III swore fealty to King Philip II of Spain.\"},{\"time\":1696947744,\"fields\":{\"prompt\":\"Who did King Charles III swear fealty to?\"},\"text\":\"Who did King Charles III swear fealty to?\"}]\n"
     ]
    }
   ],
   "source": [
    "# Search for a particular prompt text.\n",
    "query = \"normandy\"\n",
    "response = client.search_log(query, 0, int(time.time()))\n",
    "print(\"Results for\", query, \":\", response.text)\n",
    "\n",
    "print(\"===\")\n",
    "\n",
    "query = \"king charles III\"\n",
    "response = client.search_log(\"king charles III\", 0, int(time.time()))\n",
    "print(\"Results for\", query, \":\", response.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32dc6598-e405-47a1-ac8c-5e41c4e1ff57",
   "metadata": {},
   "source": [
    "# Example 2: Summarize a piece of text using ChatOpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2cdca23e-4247-4876-aa6b-042cb7d3b126",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set your key here.\n",
    "# os.environ[\"OPENAI_API_KEY\"] = \"YOUR_API_KEY\"\n",
    "\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.document_loaders import WebBaseLoader\n",
    "from langchain.chains.summarize import load_summarize_chain\n",
    "\n",
    "# Create callback handler. This logs latency, errors, token usage, prompts, as well as prompt responses to Infino.\n",
    "handler = InfinoCallbackHandler(\n",
    "    model_id=\"test_chatopenai\", model_version=\"0.1\", verbose=False\n",
    ")\n",
    "\n",
    "urls = [\"https://lilianweng.github.io/posts/2023-06-23-agent/\",\n",
    "        \"https://medium.com/lyft-engineering/lyftlearn-ml-model-training-infrastructure-built-on-kubernetes-aef8218842bb\",\n",
    "        \"https://blog.langchain.dev/week-of-10-2-langchain-release-notes/\"]\n",
    "\n",
    "for url in urls:\n",
    "    loader = WebBaseLoader(url)\n",
    "    docs = loader.load()\n",
    "\n",
    "    llm = ChatOpenAI(temperature=0, model_name=\"gpt-3.5-turbo-16k\", callbacks=[handler])\n",
    "    chain = load_summarize_chain(llm, chain_type=\"stuff\", verbose=False)\n",
    "\n",
    "    chain.run(docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cafd7df7-c1f9-42ea-bb0e-dff623eac52c",
   "metadata": {},
   "source": [
    "## Create Metric Charts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f1ef0b90-8d09-45c4-8b23-5fad5ba6d0d5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 35336 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 22192 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 37212 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 39244 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "response = client.search_ts(\"__name__\", \"latency\", 0, int(time.time()))\n",
    "plot(response.text, \"Latency\")\n",
    "\n",
    "response = client.search_ts(\"__name__\", \"error\", 0, int(time.time()))\n",
    "plot(response.text, \"Errors\")\n",
    "\n",
    "response = client.search_ts(\"__name__\", \"prompt_tokens\", 0, int(time.time()))\n",
    "plot(response.text, \"Prompt Tokens\")\n",
    "\n",
    "response = client.search_ts(\"__name__\", \"completion_tokens\", 0, int(time.time()))\n",
    "plot(response.text, \"Completion Tokens\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2ee18e1e-46ab-4080-92e6-f13affb384ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Full text query on prompt or prompt outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "edcc9d45-563e-42dd-9a0a-75caa36775e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "===\n"
     ]
    }
   ],
   "source": [
    "# Search for a particular prompt text.\n",
    "query = \"machine learning\"\n",
    "response = client.search_log(query, 0, int(time.time()))\n",
    "\n",
    "# The output can be verbose - uncomment below if it needs to be printed.\n",
    "# print(\"Results for\", query, \":\", response.text)\n",
    "\n",
    "print(\"===\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ada8d518-ff99-4796-835d-9234b898e9c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Stop Infino server"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1566f92f-29f5-499e-996a-86b0de4e4456",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "infino-example\n"
     ]
    }
   ],
   "source": [
    "!docker rm -f infino-example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b16dc8d-e60b-4195-a0bd-72ef4e0b6b83",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
