langtrace-python-sdk 3.3.21__py3-none-any.whl → 3.3.22__py3-none-any.whl

examples/awsbedrock_examples/__init__.py

@@ -1,10 +1,19 @@
-from examples.awsbedrock_examples.converse import use_converse
+from examples.awsbedrock_examples.converse import (
+    use_converse_stream,
+    use_converse,
+    use_invoke_model_anthropic,
+    use_invoke_model_cohere,
+    use_invoke_model_amazon,
+)
 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_stream()
         use_converse()
+        use_invoke_model_anthropic()
+        use_invoke_model_cohere()
+        use_invoke_model_amazon()

examples/awsbedrock_examples/converse.py

@@ -1,34 +1,174 @@
-import os
 import boto3
+import json
 from langtrace_python_sdk import langtrace
+from dotenv import load_dotenv
+import botocore
+
+load_dotenv()
+langtrace.init(write_spans_to_console=False)
+
+brt = boto3.client("bedrock-runtime", region_name="us-east-1")
+brc = boto3.client("bedrock", region_name="us-east-1")
+
+
+def use_converse_stream():
+    model_id = "anthropic.claude-3-haiku-20240307-v1:0"
+    conversation = [
+        {
+            "role": "user",
+            "content": [{"text": "what is the capital of France?"}],
+        }
+    ]
+
+    try:
+        response = brt.converse_stream(
+            modelId=model_id,
+            messages=conversation,
+            inferenceConfig={"maxTokens": 4096, "temperature": 0},
+            additionalModelRequestFields={"top_k": 250},
+        )
+        # response_text = response["output"]["message"]["content"][0]["text"]
+        print(response)
+
+    except Exception as e:
+        print(f"ERROR: Can't invoke '{model_id}'. Reason: {e}")
+        exit(1)
 
-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."}],
+            "content": [{"text": "what is the capital of France?"}],
         }
     ]
 
     try:
-        response = client.converse(
+        response = brt.converse(
             modelId=model_id,
             messages=conversation,
-            inferenceConfig={"maxTokens":4096,"temperature":0},
-            additionalModelRequestFields={"top_k":250}
+            inferenceConfig={"maxTokens": 4096, "temperature": 0},
+            additionalModelRequestFields={"top_k": 250},
         )
         response_text = response["output"]["message"]["content"][0]["text"]
         print(response_text)
 
-    except (Exception) as e:
+    except Exception as e:
         print(f"ERROR: Can't invoke '{model_id}'. Reason: {e}")
-        exit(1)
+        exit(1)
+
+
+def get_foundation_models():
+    for model in brc.list_foundation_models()["modelSummaries"]:
+        print(model["modelId"])
+
+
+# Invoke Model API
+# Amazon Titan Models
+def use_invoke_model_titan(stream=False):
+    try:
+        prompt_data = "what's the capital of France?"
+        body = json.dumps(
+            {
+                "inputText": prompt_data,
+                "textGenerationConfig": {
+                    "maxTokenCount": 1024,
+                    "topP": 0.95,
+                    "temperature": 0.2,
+                },
+            }
+        )
+        modelId = "amazon.titan-text-express-v1"  # "amazon.titan-tg1-large"
+        accept = "application/json"
+        contentType = "application/json"
+
+        if stream:
+
+            response = brt.invoke_model_with_response_stream(
+                body=body, modelId=modelId, accept=accept, contentType=contentType
+            )
+        else:
+            response = brt.invoke_model(
+                body=body, modelId=modelId, accept=accept, contentType=contentType
+            )
+            response_body = json.loads(response.get("body").read())
+
+    except botocore.exceptions.ClientError as error:
+
+        if error.response["Error"]["Code"] == "AccessDeniedException":
+            print(
+                f"\x1b[41m{error.response['Error']['Message']}\
+                    \nTo troubeshoot this issue please refer to the following resources.\
+                    \nhttps://docs.aws.amazon.com/IAM/latest/UserGuide/troubleshoot_access-denied.html\
+                    \nhttps://docs.aws.amazon.com/bedrock/latest/userguide/security-iam.html\x1b[0m\n"
+            )
+
+        else:
+            raise error
+
+
+# Anthropic Models
+def use_invoke_model_anthropic(stream=False):
+    body = json.dumps(
+        {
+            "anthropic_version": "bedrock-2023-05-31",
+            "max_tokens": 1024,
+            "temperature": 0.1,
+            "top_p": 0.9,
+            "messages": [{"role": "user", "content": "Hello, Claude"}],
+        }
+    )
+    modelId = "anthropic.claude-v2"
+    accept = "application/json"
+    contentType = "application/json"
+
+    if stream:
+        response = brt.invoke_model_with_response_stream(body=body, modelId=modelId)
+        stream_response = response.get("body")
+        if stream_response:
+            for event in stream_response:
+                chunk = event.get("chunk")
+                if chunk:
+                    print(json.loads(chunk.get("bytes").decode()))
+
+    else:
+        response = brt.invoke_model(
+            body=body, modelId=modelId, accept=accept, contentType=contentType
+        )
+        response_body = json.loads(response.get("body").read())
+        # text
+        print(response_body.get("completion"))
+
+
+def use_invoke_model_llama():
+    model_id = "meta.llama3-8b-instruct-v1:0"
+    prompt = "What is the capital of France?"
+    max_gen_len = 128
+    temperature = 0.1
+    top_p = 0.9
+
+    # Create request body.
+    body = json.dumps(
+        {
+            "prompt": prompt,
+            "max_gen_len": max_gen_len,
+            "temperature": temperature,
+            "top_p": top_p,
+        }
+    )
+    response = brt.invoke_model(body=body, modelId=model_id)
+
+    response_body = json.loads(response.get("body").read())
+
+    return response_body
+
+
+# print(get_foundation_models())
+def use_invoke_model_cohere():
+    model_id = "cohere.command-r-plus-v1"
+    prompt = "What is the capital of France?"
+    body = json.dumps({"prompt": prompt, "max_tokens": 1024, "temperature": 0.1})
+    response = brt.invoke_model(body=body, modelId=model_id)
+    response_body = json.loads(response.get("body").read())
+    print(response_body)

langtrace_python_sdk/constants/instrumentation/aws_bedrock.py

@@ -1,6 +1,10 @@
 from langtrace.trace_attributes import AWSBedrockMethods
 
 APIS = {
+    "INVOKE_MODEL": {
+        "METHOD": "aws_bedrock.invoke_model",
+        "ENDPOINT": "/invoke-model",
+    },
     "CONVERSE": {
         "METHOD": AWSBedrockMethods.CONVERSE.value,
         "ENDPOINT": "/converse",

langtrace_python_sdk/instrumentation/aws_bedrock/bedrock_streaming_wrapper.py

@@ -0,0 +1,43 @@
+import json
+from wrapt import ObjectProxy
+
+
+class StreamingWrapper(ObjectProxy):
+    def __init__(
+        self,
+        response,
+        stream_done_callback=None,
+    ):
+        super().__init__(response)
+
+        self._stream_done_callback = stream_done_callback
+        self._accumulating_body = {}
+
+    def __iter__(self):
+        for event in self.__wrapped__:
+            self._process_event(event)
+            yield event
+
+    def _process_event(self, event):
+        chunk = event.get("chunk")
+        if not chunk:
+            return
+
+        decoded_chunk = json.loads(chunk.get("bytes").decode())
+        type = decoded_chunk.get("type")
+
+        if type == "message_start":
+            self._accumulating_body = decoded_chunk.get("message")
+        elif type == "content_block_start":
+            self._accumulating_body["content"].append(
+                decoded_chunk.get("content_block")
+            )
+        elif type == "content_block_delta":
+            self._accumulating_body["content"][-1]["text"] += decoded_chunk.get(
+                "delta"
+            ).get("text")
+        elif type == "message_stop" and self._stream_done_callback:
+            self._accumulating_body["invocation_metrics"] = decoded_chunk.get(
+                "amazon-bedrock-invocationMetrics"
+            )
+            self._stream_done_callback(self._accumulating_body)

langtrace_python_sdk/instrumentation/aws_bedrock/instrumentation.py

@@ -22,22 +22,13 @@ 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
-)
+from langtrace_python_sdk.instrumentation.aws_bedrock.patch import patch_aws_bedrock
 
 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"]
 
@@ -46,13 +37,11 @@ class AWSBedrockInstrumentation(BaseInstrumentor):
         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)
+        _W(
+            module="boto3",
+            name="client",
+            wrapper=patch_aws_bedrock(tracer, version),
+        )
 
     def _uninstrument(self, **kwargs):
-        pass
+        pass

langtrace_python_sdk/instrumentation/aws_bedrock/patch.py

@@ -15,8 +15,12 @@ limitations under the License.
 """
 
 import json
-from functools import wraps
 
+from langtrace_python_sdk.instrumentation.aws_bedrock.bedrock_streaming_wrapper import (
+    StreamingWrapper,
+)
+from .stream_body_wrapper import BufferedStreamBody
+from functools import wraps
 from langtrace.trace_attributes import (
     LLMSpanAttributes,
     SpanAttributes,
@@ -39,73 +43,16 @@ from langtrace_python_sdk.utils.llm import (
     get_span_name,
     set_event_completion,
     set_span_attributes,
+    set_usage_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_langtrace_attributes(version, service_provider, vendor_type="llm"),
             **get_llm_request_attributes(kwargs),
             **get_llm_url(instance),
             SpanAttributes.LLM_PATH: APIS["CONVERSE_STREAM"]["ENDPOINT"],
@@ -129,29 +76,321 @@ def converse_stream(original_method, version, tracer):
                 span.record_exception(err)
                 span.set_status(Status(StatusCode.ERROR, str(err)))
                 raise err
-            
+
+    return traced_method
+
+
+def patch_aws_bedrock(tracer, version):
+    def traced_method(wrapped, instance, args, kwargs):
+        if args and args[0] != "bedrock-runtime":
+            return wrapped(*args, **kwargs)
+
+        client = wrapped(*args, **kwargs)
+        client.invoke_model = patch_invoke_model(client.invoke_model, tracer, version)
+
+        client.converse = patch_converse(client.converse, tracer, version)
+        client.converse_stream = patch_converse_stream(
+            client.converse_stream, tracer, version
+        )
+
+        return client
+
+    return traced_method
+
+
+def patch_converse_stream(original_method, tracer, version):
+    def traced_method(*args, **kwargs):
+        modelId = kwargs.get("modelId")
+        (vendor, _) = modelId.split(".")
+        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, vendor, vendor_type="framework"),
+            **get_llm_request_attributes(kwargs, model=modelId, prompts=input_content),
+            **get_llm_url(args[0] if args else None),
+            **get_extra_attributes(),
+        }
+        with tracer.start_as_current_span(
+            name=get_span_name("aws_bedrock.converse"),
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
+            set_span_attributes(span, span_attributes)
+            response = original_method(*args, **kwargs)
+
+            if span.is_recording():
+                set_span_streaming_response(span, response)
+            return response
+
+    return traced_method
+
+
+def patch_converse(original_method, tracer, version):
+    def traced_method(*args, **kwargs):
+        modelId = kwargs.get("modelId")
+        (vendor, _) = modelId.split(".")
+        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, vendor, vendor_type="framework"),
+            **get_llm_request_attributes(kwargs, model=modelId, prompts=input_content),
+            **get_llm_url(args[0] if args else None),
+            **get_extra_attributes(),
+        }
+        with tracer.start_as_current_span(
+            name=get_span_name("aws_bedrock.converse"),
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
+            set_span_attributes(span, span_attributes)
+            response = original_method(*args, **kwargs)
+
+            if span.is_recording():
+                _set_response_attributes(span, kwargs, response)
+            return response
+
+    return traced_method
+
+
+def patch_invoke_model(original_method, tracer, version):
+    def traced_method(*args, **kwargs):
+        modelId = kwargs.get("modelId")
+        (vendor, _) = modelId.split(".")
+        span_attributes = {
+            **get_langtrace_attributes(version, vendor, vendor_type="framework"),
+            **get_extra_attributes(),
+        }
+        with tracer.start_as_current_span(
+            name=get_span_name("aws_bedrock.invoke_model"),
+            kind=SpanKind.CLIENT,
+            context=set_span_in_context(trace.get_current_span()),
+        ) as span:
+            set_span_attributes(span, span_attributes)
+            response = original_method(*args, **kwargs)
+            if span.is_recording():
+                handle_call(span, kwargs, response)
+            return response
+
     return traced_method
 
 
+def handle_call(span, kwargs, response):
+    modelId = kwargs.get("modelId")
+    (vendor, model_name) = modelId.split(".")
+    response["body"] = BufferedStreamBody(
+        response["body"]._raw_stream, response["body"]._content_length
+    )
+    request_body = json.loads(kwargs.get("body"))
+    response_body = json.loads(response.get("body").read())
+
+    set_span_attribute(span, SpanAttributes.LLM_SYSTEM, vendor)
+    set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, modelId)
+    set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, modelId)
+
+    if vendor == "amazon":
+        set_amazon_attributes(span, request_body, response_body)
+
+    if vendor == "anthropic":
+        if "prompt" in request_body:
+            set_anthropic_completions_attributes(span, request_body, response_body)
+        elif "messages" in request_body:
+            set_anthropic_messages_attributes(span, request_body, response_body)
+
+    if vendor == "meta":
+        set_llama_meta_attributes(span, request_body, response_body)
+
+
+def set_llama_meta_attributes(span, request_body, response_body):
+    set_span_attribute(
+        span, SpanAttributes.LLM_REQUEST_TOP_P, request_body.get("top_p")
+    )
+    set_span_attribute(
+        span, SpanAttributes.LLM_REQUEST_TEMPERATURE, request_body.get("temperature")
+    )
+    set_span_attribute(
+        span, SpanAttributes.LLM_REQUEST_MAX_TOKENS, request_body.get("max_gen_len")
+    )
+
+    set_usage_attributes(
+        span,
+        {
+            "input_tokens": response_body.get("prompt_token_count"),
+            "output_tokens": response_body.get("generation_token_count"),
+        },
+    )
+
+    prompts = [
+        {
+            "role": "user",
+            "content": request_body.get("prompt"),
+        }
+    ]
+
+    completions = [
+        {
+            "role": "assistant",
+            "content": response_body.get("generation"),
+        }
+    ]
+    set_span_attribute(span, SpanAttributes.LLM_PROMPTS, json.dumps(prompts))
+    print(completions)
+    set_event_completion(span, completions)
+
+
+def set_amazon_attributes(span, request_body, response_body):
+    config = request_body.get("textGenerationConfig", {})
+    prompts = [
+        {
+            "role": "user",
+            "content": request_body.get("inputText"),
+        }
+    ]
+    completions = [
+        {
+            "role": "assistant",
+            "content": result.get("outputText"),
+        }
+        for result in response_body.get("results")
+    ]
+    set_span_attribute(
+        span, SpanAttributes.LLM_REQUEST_MAX_TOKENS, config.get("maxTokenCount")
+    )
+    set_span_attribute(
+        span, SpanAttributes.LLM_REQUEST_TEMPERATURE, config.get("temperature")
+    )
+    set_span_attribute(span, SpanAttributes.LLM_REQUEST_TOP_P, config.get("topP"))
+    set_span_attribute(span, SpanAttributes.LLM_PROMPTS, json.dumps(prompts))
+    set_usage_attributes(
+        span,
+        {
+            "input_tokens": response_body.get("inputTextTokenCount"),
+            "output_tokens": sum(
+                int(result.get("tokenCount")) for result in response_body.get("results")
+            ),
+        },
+    )
+    set_event_completion(span, completions)
+
+
+def set_anthropic_completions_attributes(span, request_body, response_body):
+    set_span_attribute(
+        span,
+        SpanAttributes.LLM_REQUEST_MAX_TOKENS,
+        request_body.get("max_tokens_to_sample"),
+    )
+    set_span_attribute(
+        span,
+        SpanAttributes.LLM_REQUEST_TEMPERATURE,
+        str(request_body.get("temperature")),
+    )
+    set_span_attribute(
+        span,
+        SpanAttributes.LLM_REQUEST_TOP_P,
+        str(request_body.get("top_p")),
+    )
+    prompts = [
+        {
+            "role": "user",
+            "content": request_body.get("prompt"),
+        }
+    ]
+    completions = [
+        {
+            "role": "assistant",
+            "content": response_body.get("completion"),
+        }
+    ]
+    set_span_attribute(span, SpanAttributes.LLM_PROMPTS, json.dumps(prompts))
+    set_event_completion(span, completions)
+
+
+def set_anthropic_messages_attributes(span, request_body, response_body):
+    set_span_attribute(
+        span,
+        SpanAttributes.LLM_REQUEST_MAX_TOKENS,
+        request_body.get("max_tokens_to_sample"),
+    )
+    set_span_attribute(
+        span,
+        SpanAttributes.LLM_REQUEST_TEMPERATURE,
+        str(request_body.get("temperature")),
+    )
+    set_span_attribute(
+        span,
+        SpanAttributes.LLM_REQUEST_TOP_P,
+        str(request_body.get("top_p")),
+    )
+    set_span_attribute(
+        span, SpanAttributes.LLM_PROMPTS, json.dumps(request_body.get("messages"))
+    )
+    set_event_completion(span, response_body.get("content"))
+    set_usage_attributes(span, response_body.get("usage"))
+
+
 @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', [])
+    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
-        ]
+        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:
+    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'),
-            }
+                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"
+                ),
+            },
+        )
+
+
+def set_span_streaming_response(span, response):
+    streaming_response = ""
+    role = None
+    for event in response["stream"]:
+        if "messageStart" in event:
+            role = event["messageStart"]["role"]
+        elif "contentBlockDelta" in event:
+            delta = event["contentBlockDelta"]["delta"]
+            if "text" in delta:
+                streaming_response += delta["text"]
+        elif "metadata" in event and "usage" in event["metadata"]:
+            usage = event["metadata"]["usage"]
+            set_usage_attributes(
+                span,
+                {
+                    "input_tokens": usage.get("inputTokens"),
+                    "output_tokens": usage.get("outputTokens"),
+                },
+            )
+
+    if streaming_response:
+        set_event_completion(
+            span, [{"role": role or "assistant", "content": streaming_response}]
         )

langtrace_python_sdk/instrumentation/aws_bedrock/stream_body_wrapper.py

@@ -0,0 +1,41 @@
+from botocore.response import StreamingBody
+from botocore.exceptions import (
+    ReadTimeoutError,
+    ResponseStreamingError,
+)
+from urllib3.exceptions import ProtocolError as URLLib3ProtocolError
+from urllib3.exceptions import ReadTimeoutError as URLLib3ReadTimeoutError
+
+
+class BufferedStreamBody(StreamingBody):
+    def __init__(self, raw_stream, content_length):
+        super().__init__(raw_stream, content_length)
+        self._buffer = None
+        self._buffer_cursor = 0
+
+    def read(self, amt=None):
+        """Read at most amt bytes from the stream.
+
+        If the amt argument is omitted, read all data.
+        """
+        if self._buffer is None:
+            try:
+                self._buffer = self._raw_stream.read()
+            except URLLib3ReadTimeoutError as e:
+                # TODO: the url will be None as urllib3 isn't setting it yet
+                raise ReadTimeoutError(endpoint_url=e.url, error=e)
+            except URLLib3ProtocolError as e:
+                raise ResponseStreamingError(error=e)
+
+            self._amount_read += len(self._buffer)
+            if amt is None or (not self._buffer and amt > 0):
+                # If the server sends empty contents or
+                # we ask to read all of the contents, then we know
+                # we need to verify the content length.
+                self._verify_content_length()
+
+        if amt is None:
+            return self._buffer[self._buffer_cursor :]
+        else:
+            self._buffer_cursor += amt
+            return self._buffer[self._buffer_cursor - amt : self._buffer_cursor]

langtrace_python_sdk/version.py

@@ -1 +1 @@
-__version__ = "3.3.21"
+__version__ = "3.3.22"

{langtrace_python_sdk-3.3.21.dist-info → langtrace_python_sdk-3.3.22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langtrace-python-sdk
-Version: 3.3.21
+Version: 3.3.22
 Summary: Python SDK for LangTrace
 Project-URL: Homepage, https://github.com/Scale3-Labs/langtrace-python-sdk
 Author-email: Scale3 Labs <engineering@scale3labs.com>

{langtrace_python_sdk-3.3.21.dist-info → langtrace_python_sdk-3.3.22.dist-info}/RECORD

@@ -3,8 +3,8 @@ examples/anthropic_example/__init__.py,sha256=03us1YuvAJR6fqXX8NH2kROBfTmyz7KzFV
 examples/anthropic_example/completion.py,sha256=3_YEZrt0BLVNJT_RbLXg6JGP2bweuc_HPC2MWR73tOM,713
 examples/autogen_example/__init__.py,sha256=UJgpzL2yOmzir-DAiGFR1PB1Zz3YcQvYcq5bCN8nl0A,158
 examples/autogen_example/main.py,sha256=6OJ73VCdHgVrqnekF1S1nK8mXCUABLbUUkQtr7wOCdw,2312
-examples/awsbedrock_examples/__init__.py,sha256=MMaW1756Hqv8rRX6do_O_-SIfauLzoYxRgBemR9KL6g,263
-examples/awsbedrock_examples/converse.py,sha256=vra4yfXYynWyFenoO8wdUnksPx_o481BQlpuWkddLZY,1024
+examples/awsbedrock_examples/__init__.py,sha256=lVm2bgPNb4h83xRW2JbTsYf0DrZjrUr29g7-QleO4iQ,505
+examples/awsbedrock_examples/converse.py,sha256=-J0Hz6OSnIpO4idHol9nxL_Ri9GWwq95aqqi3wlCTRI,5317
 examples/azureopenai_example/__init__.py,sha256=PaZM90r6VN4eSOXxb6wGsyhf9-RJCNqBypzk1Xa2GJI,271
 examples/azureopenai_example/completion.py,sha256=K_GeU0TfJ9lLDfW5VI0Lmm8_I0JXf1x9Qi83ImJ350c,668
 examples/cerebras_example/__init__.py,sha256=ydfNi0DjFMGVcfo79XVG3VEbzIrHo5wYBgSJzl_asNA,295
@@ -109,12 +109,12 @@ examples/weaviate_example/__init__.py,sha256=8JMDBsRSEV10HfTd-YC7xb4txBjD3la56sn
 examples/weaviate_example/query_text.py,sha256=wPHQTc_58kPoKTZMygVjTj-2ZcdrIuaausJfMxNQnQc,127162
 langtrace_python_sdk/__init__.py,sha256=VZM6i71NR7pBQK6XvJWRelknuTYUhqwqE7PlicKa5Wg,1166
 langtrace_python_sdk/langtrace.py,sha256=nvPaJc426Iww3ildrhsSacXtLdzsZIa94_rlK2giyVM,13153
-langtrace_python_sdk/version.py,sha256=FgIp-K9CuZgCkI8xMNisdW4atuWh6YalE8E2kUmjXx4,23
+langtrace_python_sdk/version.py,sha256=RJrshlvNont5yxrTElCH9PRO5xojg3nwHJKaU-kz-pc,23
 langtrace_python_sdk/constants/__init__.py,sha256=3CNYkWMdd1DrkGqzLUgNZXjdAlM6UFMlf_F-odAToyc,146
 langtrace_python_sdk/constants/exporter/langtrace_exporter.py,sha256=EVCrouYCpY98f0KSaKr4PzNxPULTZZO6dSA_crEOyJU,106
 langtrace_python_sdk/constants/instrumentation/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langtrace_python_sdk/constants/instrumentation/anthropic.py,sha256=YX3llt3zwDY6XrYk3CB8WEVqgrzRXEw_ffyk56JoF3k,126
-langtrace_python_sdk/constants/instrumentation/aws_bedrock.py,sha256=f9eukqoxrPgPeaBJX2gpBUz1uu0dZIPahOpvoudfbH8,310
+langtrace_python_sdk/constants/instrumentation/aws_bedrock.py,sha256=QwKtO4NBarOZoGkt5cFCcpxAw3zvZxcMMWBbzPPGv-g,422
 langtrace_python_sdk/constants/instrumentation/chroma.py,sha256=hiPGYdHS0Yj4Kh3eaYBbuCAl_swqIygu80yFqkOgdak,955
 langtrace_python_sdk/constants/instrumentation/cohere.py,sha256=9yD133VdrYZ5BoJR4nJHlj67gHEImB9-KsD-NkzHW1I,1159
 langtrace_python_sdk/constants/instrumentation/common.py,sha256=DPDX8icb0Tj3OrgpbL9WeiIaMG7Si2IKiSL8YRwwor4,1203
@@ -143,8 +143,10 @@ langtrace_python_sdk/instrumentation/autogen/__init__.py,sha256=unDhpqWQIdHFw24l
 langtrace_python_sdk/instrumentation/autogen/instrumentation.py,sha256=MVDUCBi6XzLQYmZd6myAounI0HeM8QWX5leuul5Hj0Q,1262
 langtrace_python_sdk/instrumentation/autogen/patch.py,sha256=7Sq3C8Q5tT27UkWXd1SZgnLC1pbQf_tpTrxoBIYsDw4,5273
 langtrace_python_sdk/instrumentation/aws_bedrock/__init__.py,sha256=IHqPgR1kdDvcoV1nUb-B21PaJ_qbQB0jc011Udi1ioU,96
-langtrace_python_sdk/instrumentation/aws_bedrock/instrumentation.py,sha256=2l-WiyWYUEoGre92rmylq2jPZ5w4jcxTXmCTuQNC1RU,1911
-langtrace_python_sdk/instrumentation/aws_bedrock/patch.py,sha256=VAroMezSGKT2jQ5tggbdiMRIPr9mtLItGJJgZ-xoGls,6296
+langtrace_python_sdk/instrumentation/aws_bedrock/bedrock_streaming_wrapper.py,sha256=_EMlxfBhjNsB-2TrP0ek-vx1CK5bpgjXPmQWNT9eBcE,1370
+langtrace_python_sdk/instrumentation/aws_bedrock/instrumentation.py,sha256=M7Dyw1tG18ptD0ctCT9yHzO4UyKCcS8JkHg9WVaz9Ck,1473
+langtrace_python_sdk/instrumentation/aws_bedrock/patch.py,sha256=cZ20LS0lNCKoEeeJLVEG5aduuGrapu_u4v8Cl-z6gLg,13406
+langtrace_python_sdk/instrumentation/aws_bedrock/stream_body_wrapper.py,sha256=ENdhRVHBhdkIlJIc_tkf8ASijUzZdVZM-oonNLdNM48,1584
 langtrace_python_sdk/instrumentation/cerebras/__init__.py,sha256=9rHNg7PWcZ7a9jExQZlqwWPkvLGcPT-DGWot0_6Bx9k,92
 langtrace_python_sdk/instrumentation/cerebras/instrumentation.py,sha256=WPsaYxHanYnoxGjDk7fILGJSnSRUs_zoQ30JCyPBMII,1927
 langtrace_python_sdk/instrumentation/cerebras/patch.py,sha256=HR4slOrE3pMp0ABafnlYeTK61G-EnGhOgq3pd9A_G88,4697
@@ -270,8 +272,8 @@ tests/pinecone/cassettes/test_query.yaml,sha256=b5v9G3ssUy00oG63PlFUR3JErF2Js-5A
 tests/pinecone/cassettes/test_upsert.yaml,sha256=neWmQ1v3d03V8WoLl8FoFeeCYImb8pxlJBWnFd_lITU,38607
 tests/qdrant/conftest.py,sha256=9n0uHxxIjWk9fbYc4bx-uP8lSAgLBVx-cV9UjnsyCHM,381
 tests/qdrant/test_qdrant.py,sha256=pzjAjVY2kmsmGfrI2Gs2xrolfuaNHz7l1fqGQCjp5_o,3353
-langtrace_python_sdk-3.3.21.dist-info/METADATA,sha256=nXKNrJYvlgb66g57jWoUtOqOldTmw7X2BQ1nQtHp3NA,15676
-langtrace_python_sdk-3.3.21.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-langtrace_python_sdk-3.3.21.dist-info/entry_points.txt,sha256=1_b9-qvf2fE7uQNZcbUei9vLpFZBbbh9LrtGw95ssAo,70
-langtrace_python_sdk-3.3.21.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-langtrace_python_sdk-3.3.21.dist-info/RECORD,,
+langtrace_python_sdk-3.3.22.dist-info/METADATA,sha256=U0Uj5LhWApEzItJxDzJOIawCkV5_Xx_1gJU4kS05QjU,15676
+langtrace_python_sdk-3.3.22.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+langtrace_python_sdk-3.3.22.dist-info/entry_points.txt,sha256=1_b9-qvf2fE7uQNZcbUei9vLpFZBbbh9LrtGw95ssAo,70
+langtrace_python_sdk-3.3.22.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+langtrace_python_sdk-3.3.22.dist-info/RECORD,,