Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions src/langtrace_python_sdk/instrumentation/aws_bedrock/patch.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import json

from wrapt import ObjectProxy
from itertools import tee
from .stream_body_wrapper import BufferedStreamBody
from functools import wraps
from langtrace.trace_attributes import (
Expand Down Expand Up @@ -128,7 +129,9 @@ def traced_method(*args, **kwargs):
response = original_method(*args, **kwargs)

if span.is_recording():
set_span_streaming_response(span, response)
stream1, stream2 = tee(response["stream"])
set_span_streaming_response(span, stream1)
response["stream"] = stream2
return response

return traced_method
Expand Down Expand Up @@ -442,10 +445,10 @@ def _set_response_attributes(span, kwargs, result):
)


def set_span_streaming_response(span, response):
def set_span_streaming_response(span, response_stream):
streaming_response = ""
role = None
for event in response["stream"]:
for event in response_stream:
if "messageStart" in event:
role = event["messageStart"]["role"]
elif "contentBlockDelta" in event:
Expand Down
2 changes: 1 addition & 1 deletion src/langtrace_python_sdk/version.py
Original file line number Diff line number Diff line change
@@ -1 +1 @@
__version__ = "3.8.18"
__version__ = "3.8.19"
Loading