Obinna/fix awsbedrock streaming response (#538)

* fix aws bedrock streaming bug

* bump version

* add deprecated dependency

* restrict current pinecone instrumentation to v6.0.2

* hard code boto3 dependency version

Author: obinnascale3

pyproject.toml

@@ -32,8 +32,9 @@ dependencies = [
   "transformers>=4.11.3",
   "sentry-sdk>=2.14.0",
   "ujson>=5.10.0",
-  "boto3",
+  "boto3==1.38.0",
   "setuptools",
+  "Deprecated==1.2.18",
 ]
 
 requires-python = ">=3.9"
@@ -47,7 +48,7 @@ dev = [
   "qdrant-client",
   "graphlit-client",
   "python-dotenv",
-  "pinecone",
+  "pinecone>=3.1.0,<=6.0.2",
   "langchain",
   "langchain-community",
   "langchain-openai",

src/langtrace_python_sdk/instrumentation/aws_bedrock/patch.py

@@ -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 (
@@ -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
@@ -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:

src/langtrace_python_sdk/instrumentation/pinecone/instrumentation.py

@@ -33,7 +33,7 @@ class PineconeInstrumentation(BaseInstrumentor):
     The PineconeInstrumentation class represents the Pinecone instrumentation"""
 
     def instrumentation_dependencies(self) -> Collection[str]:
-        return ["pinecone >= 3.1.0"]
+        return ["pinecone >= 3.1.0", "pinecone <= 6.0.2"]
 
     def _instrument(self, **kwargs):
         tracer_provider = kwargs.get("tracer_provider")

src/langtrace_python_sdk/version.py

@@ -1 +1 @@
-__version__ = "3.8.18"
+__version__ = "3.8.19"