Obinna/s3 en 2723 add bedrock instrumentation

pyproject.toml

@@ -18,7 +18,7 @@ classifiers = [
   "Operating System :: OS Independent",
 ]
 dependencies = [
-  'trace-attributes==7.0.4',
+  'trace-attributes==7.1.0',
   'opentelemetry-api>=1.25.0',
   'opentelemetry-sdk>=1.25.0',
   'opentelemetry-instrumentation>=0.47b0',
@@ -57,7 +57,8 @@ dev = [
   "google-generativeai",
   "google-cloud-aiplatform",
   "mistralai",
-  "embedchain"
+  "boto3",
+  "embedchain",
 ]
 
 test = ["pytest", "pytest-vcr", "pytest-asyncio"]

src/examples/awsbedrock_examples/__init__.py

@@ -0,0 +1,10 @@
+from examples.awsbedrock_examples.converse import use_converse
+from langtrace_python_sdk import langtrace, with_langtrace_root_span
+
+langtrace.init()
+
+
+class AWSBedrockRunner:
+    @with_langtrace_root_span("AWS_Bedrock")
+    def run(self):
+        use_converse()

src/examples/awsbedrock_examples/converse.py

@@ -0,0 +1,34 @@
+import os
+import boto3
+from langtrace_python_sdk import langtrace
+
+langtrace.init(api_key=os.environ["LANGTRACE_API_KEY"])
+
+def use_converse():
+    model_id = "anthropic.claude-3-haiku-20240307-v1:0"
+    client = boto3.client(
+        "bedrock-runtime",
+        region_name="us-east-1",
+        aws_access_key_id=os.environ["AWS_ACCESS_KEY_ID"],
+        aws_secret_access_key=os.environ["AWS_SECRET_ACCESS_KEY"],
+    )
+    conversation = [
+        {
+            "role": "user",
+            "content": [{"text": "Write a story about a magic backpack."}],
+        }
+    ]
+
+    try:
+        response = client.converse(
+            modelId=model_id,
+            messages=conversation,
+            inferenceConfig={"maxTokens":4096,"temperature":0},
+            additionalModelRequestFields={"top_k":250}
+        )
+        response_text = response["output"]["message"]["content"][0]["text"]
+        print(response_text)
+
+    except (Exception) as e:
+        print(f"ERROR: Can't invoke '{model_id}'. Reason: {e}")
+        exit(1)

src/langtrace_python_sdk/constants/instrumentation/aws_bedrock.py

@@ -0,0 +1,12 @@
+from langtrace.trace_attributes import AWSBedrockMethods
+
+APIS = {
+    "CONVERSE": {
+        "METHOD": AWSBedrockMethods.CONVERSE.value,
+        "ENDPOINT": "/converse",
+    },
+    "CONVERSE_STREAM": {
+        "METHOD": AWSBedrockMethods.CONVERSE_STREAM.value,
+        "ENDPOINT": "/converse-stream",
+    },
+}

src/langtrace_python_sdk/constants/instrumentation/common.py

@@ -34,6 +34,7 @@
     "EMBEDCHAIN": "Embedchain",
     "AUTOGEN": "Autogen",
     "XAI": "XAI",
+    "AWS_BEDROCK": "AWS Bedrock",
 }
 
 LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY = "langtrace_additional_attributes"

src/langtrace_python_sdk/instrumentation/__init__.py

@@ -18,6 +18,7 @@
 from .vertexai import VertexAIInstrumentation
 from .gemini import GeminiInstrumentation
 from .mistral import MistralInstrumentation
+from .aws_bedrock import AWSBedrockInstrumentation
 from .embedchain import EmbedchainInstrumentation
 from .litellm import LiteLLMInstrumentation
 
@@ -44,4 +45,5 @@
     "VertexAIInstrumentation",
     "GeminiInstrumentation",
     "MistralInstrumentation",
+    "AWSBedrockInstrumentation",
 ]

src/langtrace_python_sdk/instrumentation/aws_bedrock/__init__.py

@@ -0,0 +1,3 @@
+from .instrumentation import AWSBedrockInstrumentation
+
+__all__ = ["AWSBedrockInstrumentation"]

src/langtrace_python_sdk/instrumentation/aws_bedrock/instrumentation.py

@@ -0,0 +1,58 @@
+"""
+Copyright (c) 2024 Scale3 Labs
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import importlib.metadata
+import logging
+from typing import Collection
+
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from wrapt import wrap_function_wrapper as _W
+
+from langtrace_python_sdk.instrumentation.aws_bedrock.patch import (
+    converse, converse_stream
+)
+
+logging.basicConfig(level=logging.FATAL)
+
+def _patch_client(client, version: str, tracer) -> None:
+
+    # Store original methods
+    original_converse = client.converse
+
+    # Replace with wrapped versions
+    client.converse = converse("aws_bedrock.converse", version, tracer)(original_converse)
+
+class AWSBedrockInstrumentation(BaseInstrumentor):
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["boto3 >= 1.35.31"]
+
+    def _instrument(self, **kwargs):
+        tracer_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", tracer_provider)
+        version = importlib.metadata.version("boto3")
+
+        def wrap_create_client(wrapped, instance, args, kwargs):
+            result = wrapped(*args, **kwargs)
+            if args and args[0] == 'bedrock-runtime':
+                _patch_client(result, version, tracer)
+            return result
+
+        _W("boto3", "client", wrap_create_client)
+
+    def _uninstrument(self, **kwargs):
+        pass

src/langtrace_python_sdk/instrumentation/aws_bedrock/patch.py

@@ -0,0 +1,157 @@
+"""
+Copyright (c) 2024 Scale3 Labs
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import json
+from functools import wraps
+
+from langtrace.trace_attributes import (
+    LLMSpanAttributes,
+    SpanAttributes,
+)
+from langtrace_python_sdk.utils import set_span_attribute
+from langtrace_python_sdk.utils.silently_fail import silently_fail
+from opentelemetry import trace
+from opentelemetry.trace import SpanKind
+from opentelemetry.trace.status import Status, StatusCode
+from opentelemetry.trace.propagation import set_span_in_context
+from langtrace_python_sdk.constants.instrumentation.common import (
+    SERVICE_PROVIDERS,
+)
+from langtrace_python_sdk.constants.instrumentation.aws_bedrock import APIS
+from langtrace_python_sdk.utils.llm import (
+    get_extra_attributes,
+    get_langtrace_attributes,
+    get_llm_request_attributes,
+    get_llm_url,
+    get_span_name,
+    set_event_completion,
+    set_span_attributes,
+)
+
+
+def traced_aws_bedrock_call(api_name: str, operation_name: str):
+    def decorator(method_name: str, version: str, tracer):
+        def wrapper(original_method):
+            @wraps(original_method)
+            def wrapped_method(*args, **kwargs):
+                service_provider = SERVICE_PROVIDERS["AWS_BEDROCK"]
+
+                input_content = [
+                    {
+                        'role': message.get('role', 'user'),
+                        'content': message.get('content', [])[0].get('text', "")
+                    }
+                    for message in kwargs.get('messages', [])
+                ]
+
+                span_attributes = {
+                    **get_langtrace_attributes(version, service_provider, vendor_type="framework"),
+                    **get_llm_request_attributes(kwargs, operation_name=operation_name, prompts=input_content),
+                    **get_llm_url(args[0] if args else None),
+                    SpanAttributes.LLM_PATH: APIS[api_name]["ENDPOINT"],
+                    **get_extra_attributes(),
+                }
+
+                if api_name == "CONVERSE":
+                    span_attributes.update({
+                        SpanAttributes.LLM_REQUEST_MODEL: kwargs.get('modelId'),
+                        SpanAttributes.LLM_REQUEST_MAX_TOKENS: kwargs.get('inferenceConfig', {}).get('maxTokens'),
+                        SpanAttributes.LLM_REQUEST_TEMPERATURE: kwargs.get('inferenceConfig', {}).get('temperature'),
+                        SpanAttributes.LLM_REQUEST_TOP_P: kwargs.get('inferenceConfig', {}).get('top_p'),
+                    })
+
+                attributes = LLMSpanAttributes(**span_attributes)
+
+                with tracer.start_as_current_span(
+                    name=get_span_name(APIS[api_name]["METHOD"]),
+                    kind=SpanKind.CLIENT,
+                    context=set_span_in_context(trace.get_current_span()),
+                ) as span:
+                    set_span_attributes(span, attributes)
+                    try:
+                        result = original_method(*args, **kwargs)
+                        _set_response_attributes(span, kwargs, result)
+                        span.set_status(StatusCode.OK)
+                        return result
+                    except Exception as err:
+                        span.record_exception(err)
+                        span.set_status(Status(StatusCode.ERROR, str(err)))
+                        raise err
+
+            return wrapped_method
+        return wrapper
+    return decorator
+
+
+converse = traced_aws_bedrock_call("CONVERSE", "converse")
+
+
+def converse_stream(original_method, version, tracer):
+    def traced_method(wrapped, instance, args, kwargs):
+        service_provider = SERVICE_PROVIDERS["AWS_BEDROCK"]
+
+        span_attributes = {
+            **get_langtrace_attributes
+            (version, service_provider, vendor_type="llm"),
+            **get_llm_request_attributes(kwargs),
+            **get_llm_url(instance),
+            SpanAttributes.LLM_PATH: APIS["CONVERSE_STREAM"]["ENDPOINT"],
+            **get_extra_attributes(),
+        }
+
+        attributes = LLMSpanAttributes(**span_attributes)
+
+        with tracer.start_as_current_span(
+            name=get_span_name(APIS["CONVERSE_STREAM"]["METHOD"]),
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
+            set_span_attributes(span, attributes)
+            try:
+                result = wrapped(*args, **kwargs)
+                _set_response_attributes(span, kwargs, result)
+                span.set_status(StatusCode.OK)
+                return result
+            except Exception as err:
+                span.record_exception(err)
+                span.set_status(Status(StatusCode.ERROR, str(err)))
+                raise err
+
+    return traced_method
+
+
+@silently_fail
+def _set_response_attributes(span, kwargs, result):
+    set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, kwargs.get('modelId'))
+    set_span_attribute(span, SpanAttributes.LLM_TOP_K, kwargs.get('additionalModelRequestFields', {}).get('top_k'))
+    content = result.get('output', {}).get('message', {}).get('content', [])
+    if len(content) > 0:
+        role = result.get('output', {}).get('message', {}).get('role', "assistant")
+        responses = [
+            {"role": role, "content": c.get('text', "")}
+            for c in content
+        ]
+        set_event_completion(span, responses)
+
+    if 'usage' in result:
+        set_span_attributes(
+            span,
+            {
+                SpanAttributes.LLM_USAGE_COMPLETION_TOKENS: result['usage'].get('outputTokens'),
+                SpanAttributes.LLM_USAGE_PROMPT_TOKENS: result['usage'].get('inputTokens'),
+                SpanAttributes.LLM_USAGE_TOTAL_TOKENS: result['usage'].get('totalTokens'),
+            }
+        )

src/langtrace_python_sdk/langtrace.py

@@ -56,6 +56,7 @@
     LiteLLMInstrumentation,
     LlamaindexInstrumentation,
     MistralInstrumentation,
+    AWSBedrockInstrumentation,
     OllamaInstrumentor,
     OpenAIInstrumentation,
     PineconeInstrumentation,
@@ -277,6 +278,7 @@ def init(
         "google-cloud-aiplatform": VertexAIInstrumentation(),
         "google-generativeai": GeminiInstrumentation(),
         "mistralai": MistralInstrumentation(),
+        "boto3": AWSBedrockInstrumentation(),
         "autogen": AutogenInstrumentation(),
     }
 

src/langtrace_python_sdk/types/__init__.py

@@ -30,6 +30,7 @@ class InstrumentationType(Enum):
     SQLALCHEMY = "sqlalchemy"
     VERTEXAI = "vertexai"
     WEAVIATE = "weaviate"
+    AWS_BEDROCK = "boto3"
 
     @staticmethod
     def from_string(value: str):
@@ -62,6 +63,11 @@ class VendorMethods(TypedDict):
         "mistral.embeddings.create",
     ]
 
+    AwsBedrockMethods = Literal[
+        "aws_bedrock.converse",
+        "aws_bedrock.converse_stream",
+    ]
+
     ChromadbMethods = Literal[
         "chromadb.collection.add",
         "chromadb.collection.query",
@@ -112,6 +118,7 @@ class InstrumentationMethods(TypedDict):
     open_ai: List[VendorMethods.OpenaiMethods]
     groq: List[VendorMethods.GroqMethods]
     mistral: List[VendorMethods.MistralMethods]
+    aws_bedrock: List[VendorMethods.AwsBedrockMethods]
     pinecone: List[VendorMethods.PineconeMethods]
     llamaindex: List[VendorMethods.LlamaIndexMethods]
     chromadb: List[VendorMethods.ChromadbMethods]

src/langtrace_python_sdk/version.py

@@ -1 +1 @@
-__version__ = "3.1.3"
+__version__ = "3.2.0"

src/run_example.py

@@ -16,10 +16,11 @@
     "weaviate": False,
     "ollama": False,
     "groq": False,
-    "autogen": True,
+    "autogen": False,
     "vertexai": False,
     "gemini": False,
     "mistral": False,
+    "awsbedrock": True,
 }
 
 if ENABLED_EXAMPLES["anthropic"]:
@@ -123,3 +124,9 @@
 
     print(Fore.BLUE + "Running Azure OpenAI example" + Fore.RESET)
     AzureOpenAIRunner().run()
+
+if ENABLED_EXAMPLES["awsbedrock"]:
+    from examples.awsbedrock_examples import AWSBedrockRunner
+
+    print(Fore.BLUE + "Running AWS Bedrock example" + Fore.RESET)
+    AWSBedrockRunner().run()