feat(integrations): add AWS Bedrock tracer

examples/tracing/bedrock/bedrock_tracing.ipynb

@@ -0,0 +1,164 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "091d7544",
+   "metadata": {},
+   "source": [
+    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/openlayer-ai/openlayer-python/blob/main/examples/tracing/bedrock/bedrock_tracing.ipynb)\n",
+    "\n",
+    "\n",
+    "# <a id=\"top\">Tracing a AWS Bedrock model invocation</a>\n",
+    "\n",
+    "## 1. Set the environment variables"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5c1adbce",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "# Openlayer env variables\n",
+    "os.environ[\"OPENLAYER_API_KEY\"] = \"YOUR_OPENLAYER_API_KEY_HERE\"\n",
+    "os.environ[\"OPENLAYER_INFERENCE_PIPELINE_ID\"] = \"YOUR_OPENLAYER_INFERENCE_PIPELINE_ID_HERE\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "13c44cbd",
+   "metadata": {},
+   "source": [
+    "## 2. Initialize the session"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9c82b04f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "\n",
+    "import boto3"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "21659c33",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Initialize a session using Amazon Bedrock\n",
+    "session = boto3.Session(\n",
+    "    aws_access_key_id='YOUR_AWS_ACCESS_KEY_ID_HERE',\n",
+    "    aws_secret_access_key='YOUR_AWS_SECRET_ACCESS_KEY_HERE',\n",
+    "    region_name='us-east-1'  # Change to your desired region\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "017c53be",
+   "metadata": {},
+   "source": [
+    "## 3. Wrap the Bedrock client in Openlayer's `trace_bedrock` function"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "24ddd361",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from openlayer.lib import trace_bedrock"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fecb56cb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "bedrock_client = trace_bedrock(session.client(service_name='bedrock-runtime'))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4eb11465",
+   "metadata": {},
+   "source": [
+    "## 4. Invoke the model normally\n",
+    "\n",
+    "That's it! Now you can continue using the traced Bedrock client normally. The data is automatically published to Openlayer and you can start creating tests around it!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c3f97c28",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Define the model ID and the input prompt\n",
+    "model_id = 'anthropic.claude-3-5-sonnet-20240620-v1:0'  # Replace with your model ID\n",
+    "input_data = {\n",
+    "  \"max_tokens\": 256,\n",
+    "  \"messages\": [{\"role\": \"user\", \"content\": \"Hello, world\"}],\n",
+    "  \"anthropic_version\": \"bedrock-2023-05-31\"\n",
+    "}\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1dcd33b8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Invoke the model\n",
+    "response = bedrock_client.invoke_model(\n",
+    "    body=json.dumps(input_data),\n",
+    "    contentType='application/json',\n",
+    "    accept='application/json',\n",
+    "    modelId=model_id\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3a647127",
+   "metadata": {},
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "bedrock-test",
+   "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.12.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/openlayer/lib/__init__.py

@@ -9,6 +9,7 @@
     "trace_groq",
     "trace_async_openai",
     "trace_async",
+    "trace_bedrock",
 ]
 
 # ---------------------------------- Tracing --------------------------------- #
@@ -84,3 +85,26 @@ def trace_groq(client):
     if not isinstance(client, groq.Groq):
         raise ValueError("Invalid client. Please provide a Groq client.")
     return groq_tracer.trace_groq(client)
+
+
+def trace_bedrock(client):
+    """Trace AWS Bedrock model invocations."""
+    # pylint: disable=import-outside-toplevel
+    try:
+        import boto3
+    except ImportError:
+        raise ImportError(
+            "boto3 is required for Bedrock tracing. Install with: pip install boto3"
+        )
+
+    from .integrations import bedrock_tracer
+
+    # Check if it's a boto3 client for bedrock-runtime service
+    if (
+        not hasattr(client, "_service_model")
+        or client._service_model.service_name != "bedrock-runtime"
+    ):
+        raise ValueError(
+            "Invalid client. Please provide a boto3 bedrock-runtime client."
+        )
+    return bedrock_tracer.trace_bedrock(client)