[OPIK-271] Add integration with aws bedrock (#529)

* Init bedrock namespace * Initial implementation for bedrock converse integration (non-streaming mode) * Draft implementation for converse_stream support * Update implementation for bedrock integration. Stream inside the response is now wrapped and response is returned. Metadata from the start of the span/trace is now merged with the metadata from the end of the span/trace. * Fix lint errors * Update _generators_handler * Fix lint errors * Add comments to SpanData and TraceData classes * Add comment to SpanData and TraceData class * Add comment * created_by -> created_from for unification with other integrations * add docstrings * Add first draft of Bedrock cookbook * Slight update to OpenAI cookbook to avoid recreating the client every time * Add streaming example to BedRock example * Add streaming example to BedRock example * Apply suggestions from code review Co-authored-by: Jacques Verré <[email protected]> * Update bedrock cookbook Pass project_name directly when possible Add screenshots * Introduce Opik.search_spans (#545) * Added Opik.search_spans method * [OPIK-309] Create prompt endpoint (#531) * [OPIK-309] Create prompt endpoint * Add logic to create first version when specified * Address PR review comments * Move piece of docstring to comment * Add protection from None in _generators_handler * Fix track-stacking when using track_bedrock or track_openai multiple times * Fix output structure in bedrock chunks aggregator --------- Co-authored-by: Boris Feld <[email protected]> Co-authored-by: Jacques Verré <[email protected]> Co-authored-by: Thiago dos Santos Hora <[email protected]>
comet-ml · Nov 5, 2024 · 9eca8dd · 9eca8dd
1 parent c1b4f7a
commit 9eca8dd
Show file tree

Hide file tree

Showing 16 changed files with 684 additions and 35 deletions.
diff --git a/apps/opik-documentation/documentation/docs/cookbook/bedrock.ipynb b/apps/opik-documentation/documentation/docs/cookbook/bedrock.ipynb
@@ -0,0 +1,309 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Using Opik with AWS Bedrock\n",
+    "\n",
+    "Opik integrates with AWS Bedrock to provide a simple way to log traces for all Bedrock LLM calls. This works for all supported models, including if you are using the streaming API.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Creating an account on Comet.com\n",
+    "\n",
+    "[Comet](https://www.comet.com/site?from=llm&utm_source=opik&utm_medium=colab&utm_content=bedrock&utm_campaign=opik) provides a hosted version of the Opik platform, [simply create an account](https://www.comet.com/signup?from=llm&utm_source=opik&utm_medium=colab&utm_content=bedrock&utm_campaign=opik) and grab you API Key.\n",
+    "\n",
+    "> You can also run the Opik platform locally, see the [installation guide](https://www.comet.com/docs/opik/self-host/overview/?from=llm&utm_source=opik&utm_medium=colab&utm_content=bedrock&utm_campaign=opik) for more information."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install --upgrade opik boto3"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import opik\n",
+    "\n",
+    "opik.configure(use_local=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Preparing our environment\n",
+    "\n",
+    "First, we will set up our bedrock client. Uncomment the following lines to pass AWS Credentials manually or [checkout other ways of passing credentials to Boto3](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html). You will also need to request access to the model in the UI before being able to generate text, here we are gonna use the Llama 3.2 model, you can request access to it in [this page for the us-east1](https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/providers?model=meta.llama3-2-3b-instruct-v1:0) region."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import boto3\n",
+    "\n",
+    "REGION = \"us-east-1\"\n",
+    "\n",
+    "MODEL_ID = \"us.meta.llama3-2-3b-instruct-v1:0\"\n",
+    "\n",
+    "bedrock = boto3.client(\n",
+    "    service_name=\"bedrock-runtime\",\n",
+    "    region_name=REGION,\n",
+    "    # aws_access_key_id=ACCESS_KEY,\n",
+    "    # aws_secret_access_key=SECRET_KEY,\n",
+    "    # aws_session_token=SESSION_TOKEN,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Logging traces\n",
+    "\n",
+    "In order to log traces to Opik, we need to wrap our Bedrock calls with the `track_bedrock` function:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "from opik.integrations.bedrock import track_bedrock\n",
+    "\n",
+    "bedrock_client = track_bedrock(bedrock, project_name=\"bedrock-integration-demo\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "PROMPT = \"Why is it important to use a LLM Monitoring like CometML Opik tool that allows you to log traces and spans when working with LLM Models hosted on AWS Bedrock?\"\n",
+    "\n",
+    "response = bedrock_client.converse(\n",
+    "    modelId=MODEL_ID,\n",
+    "    messages=[{\"role\": \"user\", \"content\": [{\"text\": PROMPT}]}],\n",
+    "    inferenceConfig={\"temperature\": 0.5, \"maxTokens\": 512, \"topP\": 0.9},\n",
+    ")\n",
+    "print(\"Response\", response[\"output\"][\"message\"][\"content\"][0][\"text\"])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The prompt and response messages are automatically logged to Opik and can be viewed in the UI.\n",
+    "\n",
+    "![Bedrock Integration](https://raw.githubusercontent.com/comet-ml/opik/main/apps/opik-documentation/documentation/static/img/cookbook/bedrock_trace_cookbook.png)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Logging traces with streaming"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def stream_conversation(\n",
+    "    bedrock_client,\n",
+    "    model_id,\n",
+    "    messages,\n",
+    "    system_prompts,\n",
+    "    inference_config,\n",
+    "):\n",
+    "    \"\"\"\n",
+    "    Sends messages to a model and streams the response.\n",
+    "    Args:\n",
+    "        bedrock_client: The Boto3 Bedrock runtime client.\n",
+    "        model_id (str): The model ID to use.\n",
+    "        messages (JSON) : The messages to send.\n",
+    "        system_prompts (JSON) : The system prompts to send.\n",
+    "        inference_config (JSON) : The inference configuration to use.\n",
+    "        additional_model_fields (JSON) : Additional model fields to use.\n",
+    "\n",
+    "    Returns:\n",
+    "        Nothing.\n",
+    "\n",
+    "    \"\"\"\n",
+    "\n",
+    "    response = bedrock_client.converse_stream(\n",
+    "        modelId=model_id,\n",
+    "        messages=messages,\n",
+    "        system=system_prompts,\n",
+    "        inferenceConfig=inference_config,\n",
+    "    )\n",
+    "\n",
+    "    stream = response.get(\"stream\")\n",
+    "    if stream:\n",
+    "        for event in stream:\n",
+    "            if \"messageStart\" in event:\n",
+    "                print(f\"\\nRole: {event['messageStart']['role']}\")\n",
+    "\n",
+    "            if \"contentBlockDelta\" in event:\n",
+    "                print(event[\"contentBlockDelta\"][\"delta\"][\"text\"], end=\"\")\n",
+    "\n",
+    "            if \"messageStop\" in event:\n",
+    "                print(f\"\\nStop reason: {event['messageStop']['stopReason']}\")\n",
+    "\n",
+    "            if \"metadata\" in event:\n",
+    "                metadata = event[\"metadata\"]\n",
+    "                if \"usage\" in metadata:\n",
+    "                    print(\"\\nToken usage\")\n",
+    "                    print(f\"Input tokens: {metadata['usage']['inputTokens']}\")\n",
+    "                    print(f\":Output tokens: {metadata['usage']['outputTokens']}\")\n",
+    "                    print(f\":Total tokens: {metadata['usage']['totalTokens']}\")\n",
+    "                if \"metrics\" in event[\"metadata\"]:\n",
+    "                    print(f\"Latency: {metadata['metrics']['latencyMs']} milliseconds\")\n",
+    "\n",
+    "\n",
+    "system_prompt = \"\"\"You are an app that creates playlists for a radio station\n",
+    "  that plays rock and pop music. Only return song names and the artist.\"\"\"\n",
+    "\n",
+    "# Message to send to the model.\n",
+    "input_text = \"Create a list of 3 pop songs.\"\n",
+    "\n",
+    "\n",
+    "message = {\"role\": \"user\", \"content\": [{\"text\": input_text}]}\n",
+    "messages = [message]\n",
+    "\n",
+    "# System prompts.\n",
+    "system_prompts = [{\"text\": system_prompt}]\n",
+    "\n",
+    "# inference parameters to use.\n",
+    "temperature = 0.5\n",
+    "top_p = 0.9\n",
+    "# Base inference parameters.\n",
+    "inference_config = {\"temperature\": temperature, \"topP\": 0.9}\n",
+    "\n",
+    "\n",
+    "stream_conversation(\n",
+    "    bedrock_client,\n",
+    "    MODEL_ID,\n",
+    "    messages,\n",
+    "    system_prompts,\n",
+    "    inference_config,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "![Bedrock Integration](https://raw.githubusercontent.com/comet-ml/opik/main/apps/opik-documentation/documentation/static/img/cookbook/bedrock_trace_streaming_cookbook.png)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Using it with the `track` decorator\n",
+    "\n",
+    "If you have multiple steps in your LLM pipeline, you can use the `track` decorator to log the traces for each step. If Bedrock is called within one of these steps, the LLM call with be associated with that corresponding step:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from opik import track\n",
+    "from opik.integrations.bedrock import track_bedrock\n",
+    "\n",
+    "bedrock = boto3.client(\n",
+    "    service_name=\"bedrock-runtime\",\n",
+    "    region_name=REGION,\n",
+    "    # aws_access_key_id=ACCESS_KEY,\n",
+    "    # aws_secret_access_key=SECRET_KEY,\n",
+    "    # aws_session_token=SESSION_TOKEN,\n",
+    ")\n",
+    "\n",
+    "os.environ[\"OPIK_PROJECT_NAME\"] = \"bedrock-integration-demo\"\n",
+    "bedrock_client = track_bedrock(bedrock)\n",
+    "\n",
+    "\n",
+    "@track\n",
+    "def generate_story(prompt):\n",
+    "    res = bedrock_client.converse(\n",
+    "        modelId=MODEL_ID, messages=[{\"role\": \"user\", \"content\": [{\"text\": prompt}]}]\n",
+    "    )\n",
+    "    return res[\"output\"][\"message\"][\"content\"][0][\"text\"]\n",
+    "\n",
+    "\n",
+    "@track\n",
+    "def generate_topic():\n",
+    "    prompt = \"Generate a topic for a story about Opik.\"\n",
+    "    res = bedrock_client.converse(\n",
+    "        modelId=MODEL_ID, messages=[{\"role\": \"user\", \"content\": [{\"text\": prompt}]}]\n",
+    "    )\n",
+    "    return res[\"output\"][\"message\"][\"content\"][0][\"text\"]\n",
+    "\n",
+    "\n",
+    "@track\n",
+    "def generate_opik_story():\n",
+    "    topic = generate_topic()\n",
+    "    story = generate_story(topic)\n",
+    "    return story\n",
+    "\n",
+    "\n",
+    "generate_opik_story()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The trace can now be viewed in the UI:\n",
+    "\n",
+    "![Bedrock Integration](https://raw.githubusercontent.com/comet-ml/opik/main/apps/opik-documentation/documentation/static/img/cookbook/bedrock_trace_decorator_cookbook.png)"
+   ]
+  }
+ ],
+ "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.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}
diff --git a/apps/opik-documentation/documentation/docs/cookbook/openai.ipynb b/apps/opik-documentation/documentation/docs/cookbook/openai.ipynb
@@ -26,7 +26,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "%pip install --upgrade --quiet opik openai"
+    "%pip install --upgrade opik openai"
    ]
   },
   {
@@ -81,10 +81,17 @@
     "from openai import OpenAI\n",
     "\n",
     "os.environ[\"OPIK_PROJECT_NAME\"] = \"openai-integration-demo\"\n",
-    "client = OpenAI()\n",
-    "\n",
-    "openai_client = track_openai(client)\n",
     "\n",
+    "client = OpenAI()\n",
+    "openai_client = track_openai(client)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "prompt = \"\"\"\n",
     "Write a short two sentence story about Opik.\n",
     "\"\"\"\n",

diff --git a/...opik-documentation/documentation/static/img/cookbook/bedrock_trace_cookbook.png b/...opik-documentation/documentation/static/img/cookbook/bedrock_trace_cookbook.png
diff --git a/...entation/documentation/static/img/cookbook/bedrock_trace_decorator_cookbook.png b/...entation/documentation/static/img/cookbook/bedrock_trace_decorator_cookbook.png
diff --git a/...entation/documentation/static/img/cookbook/bedrock_trace_streaming_cookbook.png b/...entation/documentation/static/img/cookbook/bedrock_trace_streaming_cookbook.png