|
| 1 | +{ |
| 2 | + "cells": [ |
| 3 | + { |
| 4 | + "cell_type": "markdown", |
| 5 | + "id": "83c16ef6-98e7-48d0-b82f-4029a730ff00", |
| 6 | + "metadata": {}, |
| 7 | + "source": [ |
| 8 | + "[](https://colab.research.google.com/github/openlayer-ai/examples-gallery/blob/main/monitoring/llms/rag-tracing/rag_tracer.ipynb)\n", |
| 9 | + "\n", |
| 10 | + "\n", |
| 11 | + "# <a id=\"top\">Tracing a RAG system</a>" |
| 12 | + ] |
| 13 | + }, |
| 14 | + { |
| 15 | + "cell_type": "code", |
| 16 | + "execution_count": null, |
| 17 | + "id": "21137554-ad8e-444b-bf2e-49393f072956", |
| 18 | + "metadata": {}, |
| 19 | + "outputs": [], |
| 20 | + "source": [ |
| 21 | + "import os\n", |
| 22 | + "import openai\n", |
| 23 | + "\n", |
| 24 | + "# OpenAI env variable\n", |
| 25 | + "os.environ[\"OPENAI_API_KEY\"] = \"YOUR_OPENAI_KEY_HERE\"\n", |
| 26 | + "\n", |
| 27 | + "# Openlayer env variables\n", |
| 28 | + "os.environ[\"OPENLAYER_API_KEY\"] = \"YOUR_OPENLAYER_API_KEY_HERE\"\n", |
| 29 | + "os.environ[\"OPENLAYER_PROJECT_NAME\"] = \"YOUR_OPENLAYER_PROJECT_NAME_HERE\" # Where the traces will be uploaded to" |
| 30 | + ] |
| 31 | + }, |
| 32 | + { |
| 33 | + "cell_type": "markdown", |
| 34 | + "id": "20b25a1f-529e-45c5-90e5-26485914f511", |
| 35 | + "metadata": {}, |
| 36 | + "source": [ |
| 37 | + "## Defining and decorating our RAG system" |
| 38 | + ] |
| 39 | + }, |
| 40 | + { |
| 41 | + "cell_type": "code", |
| 42 | + "execution_count": null, |
| 43 | + "id": "9e2f8d80-d49a-48f0-8c12-350045dff985", |
| 44 | + "metadata": {}, |
| 45 | + "outputs": [], |
| 46 | + "source": [ |
| 47 | + "%%bash\n", |
| 48 | + "\n", |
| 49 | + "if [ ! -e \"context.txt\" ]; then\n", |
| 50 | + " curl \"https://raw.githubusercontent.com/openlayer-ai/examples-gallery/main/monitoring/llms/rag-tracing/context.txt\" --output \"context.txt\"\n", |
| 51 | + "fi" |
| 52 | + ] |
| 53 | + }, |
| 54 | + { |
| 55 | + "cell_type": "code", |
| 56 | + "execution_count": null, |
| 57 | + "id": "60d470d7-3aa0-4703-a9e7-cab24325a4a5", |
| 58 | + "metadata": {}, |
| 59 | + "outputs": [], |
| 60 | + "source": [ |
| 61 | + "import random\n", |
| 62 | + "import time\n", |
| 63 | + "\n", |
| 64 | + "import numpy as np\n", |
| 65 | + "from openai import OpenAI\n", |
| 66 | + "from sklearn.feature_extraction.text import TfidfVectorizer\n", |
| 67 | + "from sklearn.metrics.pairwise import cosine_similarity\n", |
| 68 | + "\n", |
| 69 | + "from openlayer import llm_monitors\n", |
| 70 | + "from openlayer.tracing import tracer" |
| 71 | + ] |
| 72 | + }, |
| 73 | + { |
| 74 | + "cell_type": "code", |
| 75 | + "execution_count": null, |
| 76 | + "id": "c8070d3f-ebec-4faf-8959-23e6ac22737d", |
| 77 | + "metadata": {}, |
| 78 | + "outputs": [], |
| 79 | + "source": [ |
| 80 | + "class RagPipeline:\n", |
| 81 | + " def __init__(self, context_path: str):\n", |
| 82 | + " # Wrap OpenAI client with Openlayer's OpenAIMonitor to trace it \n", |
| 83 | + " self.openai_client = OpenAI()\n", |
| 84 | + " llm_monitors.OpenAIMonitor(client=self.openai_client)\n", |
| 85 | + " \n", |
| 86 | + " self.vectorizer = TfidfVectorizer()\n", |
| 87 | + " with open(context_path, 'r', encoding='utf-8') as file:\n", |
| 88 | + " self.context_sections = file.read().split('\\n\\n') \n", |
| 89 | + " self.tfidf_matrix = self.vectorizer.fit_transform(self.context_sections)\n", |
| 90 | + "\n", |
| 91 | + " # Decorate the functions you'd like to trace with @tracer.trace()\n", |
| 92 | + " @tracer.trace()\n", |
| 93 | + " def query(self, user_query: str) -> str:\n", |
| 94 | + " \"\"\"Main method.\n", |
| 95 | + "\n", |
| 96 | + " Answers to a user query with the LLM.\n", |
| 97 | + " \"\"\"\n", |
| 98 | + " context = self.retrieve_context(user_query)\n", |
| 99 | + " prompt = self.inject_prompt(user_query, context)\n", |
| 100 | + " answer = self.generate_answer_with_gpt(prompt)\n", |
| 101 | + " return answer\n", |
| 102 | + "\n", |
| 103 | + " @tracer.trace()\n", |
| 104 | + " def retrieve_context(self, query: str) -> str:\n", |
| 105 | + " \"\"\"Context retriever. \n", |
| 106 | + " \n", |
| 107 | + " Given the query, returns the most similar context (using TFIDF).\n", |
| 108 | + " \"\"\"\n", |
| 109 | + " query_vector = self.vectorizer.transform([query])\n", |
| 110 | + " cosine_similarities = cosine_similarity(query_vector, self.tfidf_matrix).flatten()\n", |
| 111 | + " most_relevant_idx = np.argmax(cosine_similarities)\n", |
| 112 | + " return self.context_sections[most_relevant_idx]\n", |
| 113 | + "\n", |
| 114 | + " @tracer.trace()\n", |
| 115 | + " def inject_prompt(self, query: str, context: str):\n", |
| 116 | + " \"\"\"Combines the query with the context and returns\n", |
| 117 | + " the prompt (formatted to conform with OpenAI models).\"\"\"\n", |
| 118 | + " return [\n", |
| 119 | + " {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n", |
| 120 | + " {\"role\": \"user\", \"content\": f\"Answer the user query using only the following context: {context}. \\nUser query: {query}\"}\n", |
| 121 | + " ]\n", |
| 122 | + "\n", |
| 123 | + " @tracer.trace()\n", |
| 124 | + " def generate_answer_with_gpt(self, prompt):\n", |
| 125 | + " \"\"\"Forwards the prompt to GPT and returns the answer.\"\"\"\n", |
| 126 | + " response = self.openai_client.chat.completions.create(\n", |
| 127 | + " messages=prompt,\n", |
| 128 | + " model=\"gpt-3.5-turbo\",\n", |
| 129 | + " )\n", |
| 130 | + " return response.choices[0].message.content.strip()" |
| 131 | + ] |
| 132 | + }, |
| 133 | + { |
| 134 | + "cell_type": "code", |
| 135 | + "execution_count": null, |
| 136 | + "id": "f96f7073-7be4-4254-a6c9-eb808312beb8", |
| 137 | + "metadata": {}, |
| 138 | + "outputs": [], |
| 139 | + "source": [ |
| 140 | + "rag = RagPipeline(\"context.txt\")" |
| 141 | + ] |
| 142 | + }, |
| 143 | + { |
| 144 | + "cell_type": "code", |
| 145 | + "execution_count": null, |
| 146 | + "id": "50e046fd-68f1-4f66-b2a1-03aa95b9b367", |
| 147 | + "metadata": {}, |
| 148 | + "outputs": [], |
| 149 | + "source": [ |
| 150 | + "rag.query(\"Who were the founders of Apple?\")" |
| 151 | + ] |
| 152 | + }, |
| 153 | + { |
| 154 | + "cell_type": "code", |
| 155 | + "execution_count": null, |
| 156 | + "id": "afc7f963-fc13-4e93-b3ef-98aa183770a3", |
| 157 | + "metadata": {}, |
| 158 | + "outputs": [], |
| 159 | + "source": [ |
| 160 | + "rag.query(\"When did Apple IPO?\")" |
| 161 | + ] |
| 162 | + }, |
| 163 | + { |
| 164 | + "cell_type": "markdown", |
| 165 | + "id": "42f1e832-4c3f-4a6a-8013-8607ff141f67", |
| 166 | + "metadata": {}, |
| 167 | + "source": [ |
| 168 | + "That's it! After each inference, the traces are uploaded to Openlayer. If you navigate to your project, you should see the traces for these two inferences with our RAG system." |
| 169 | + ] |
| 170 | + }, |
| 171 | + { |
| 172 | + "cell_type": "code", |
| 173 | + "execution_count": null, |
| 174 | + "id": "f960a36f-3438-4c81-8cdb-ca078aa509cd", |
| 175 | + "metadata": {}, |
| 176 | + "outputs": [], |
| 177 | + "source": [] |
| 178 | + } |
| 179 | + ], |
| 180 | + "metadata": { |
| 181 | + "kernelspec": { |
| 182 | + "display_name": "Python 3 (ipykernel)", |
| 183 | + "language": "python", |
| 184 | + "name": "python3" |
| 185 | + }, |
| 186 | + "language_info": { |
| 187 | + "codemirror_mode": { |
| 188 | + "name": "ipython", |
| 189 | + "version": 3 |
| 190 | + }, |
| 191 | + "file_extension": ".py", |
| 192 | + "mimetype": "text/x-python", |
| 193 | + "name": "python", |
| 194 | + "nbconvert_exporter": "python", |
| 195 | + "pygments_lexer": "ipython3", |
| 196 | + "version": "3.9.18" |
| 197 | + } |
| 198 | + }, |
| 199 | + "nbformat": 4, |
| 200 | + "nbformat_minor": 5 |
| 201 | +} |
0 commit comments