Merge pull request #459 from Scale3-Labs/ali/bedrock-session

add support for `Session.client` bedrock to handle langchain

Author: alizenhom

src/examples/awsbedrock_examples/__init__.py

@@ -1,5 +1,6 @@
 from examples.awsbedrock_examples.converse import (
     use_invoke_model_anthropic,
+    init_bedrock_langchain,
     use_invoke_model_titan,
     use_invoke_model_llama,
 )
@@ -15,4 +16,5 @@ def run(self):
         # use_invoke_model_anthropic(stream=True)
         # use_invoke_model_cohere()
         # use_invoke_model_llama(stream=False)
-        use_invoke_model_titan(stream=False)
+        # use_invoke_model_titan(stream=False)
+        init_bedrock_langchain()

src/examples/awsbedrock_examples/converse.py

@@ -3,6 +3,7 @@
 from langtrace_python_sdk import langtrace
 from dotenv import load_dotenv
 import botocore
+from langchain_aws import ChatBedrock
 
 load_dotenv()
 langtrace.init(write_spans_to_console=False)
@@ -186,3 +187,13 @@ def use_invoke_model_cohere():
     response = brt.invoke_model(body=body, modelId=model_id)
     response_body = json.loads(response.get("body").read())
     print(response_body)
+
+
+def init_bedrock_langchain(temperature=0.1):
+    chat = ChatBedrock(
+        model_id="anthropic.claude-v2",
+        streaming=True,
+        model_kwargs={"temperature": temperature},
+        region_name="us-east-1",
+    )
+    return chat.invoke("What is the capital of France?")

src/langtrace_python_sdk/instrumentation/aws_bedrock/instrumentation.py

@@ -42,6 +42,11 @@ def _instrument(self, **kwargs):
             name="client",
             wrapper=patch_aws_bedrock(tracer, version),
         )
+        _W(
+            module="boto3.session",
+            name="Session.client",
+            wrapper=patch_aws_bedrock(tracer, version),
+        )
 
     def _uninstrument(self, **kwargs):
         pass

src/langtrace_python_sdk/version.py

@@ -1 +1 @@
-__version__ = "3.3.25"
+__version__ = "3.3.26"