Add support for bedrock

Author: alizenhom

src/examples/awsbedrock_examples/__init__.py

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

src/examples/awsbedrock_examples/converse.py

@@ -1,23 +1,46 @@
 import boto3
-import botocore
 import json
 from langtrace_python_sdk import langtrace
 from dotenv import load_dotenv
-
+import botocore
 
 load_dotenv()
-langtrace.init()
+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)
+
+
 def use_converse():
     model_id = "anthropic.claude-3-haiku-20240307-v1:0"
     conversation = [
         {
             "role": "user",
-            "content": [{"text": "Write a story about a magic backpack."}],
+            "content": [{"text": "what is the capital of France?"}],
         }
     ]
 
@@ -37,17 +60,15 @@ def use_converse():
 
 
 def get_foundation_models():
-    models = []
     for model in brc.list_foundation_models()["modelSummaries"]:
-        models.append(model["modelId"])
-    return models
+        print(model["modelId"])
 
 
 # Invoke Model API
 # Amazon Titan Models
-def use_invoke_model_titan():
+def use_invoke_model_titan(stream=False):
     try:
-        prompt_data = "what's 1+1?"
+        prompt_data = "what's the capital of France?"
         body = json.dumps(
             {
                 "inputText": prompt_data,
@@ -62,12 +83,16 @@ def use_invoke_model_titan():
         accept = "application/json"
         contentType = "application/json"
 
-        response = brt.invoke_model(
-            body=body, modelId=modelId, accept=accept, contentType=contentType
-        )
-        response_body = json.loads(response.get("body").read())
+        if stream:
 
-        # print(response_body.get("results"))
+            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:
 
@@ -84,8 +109,66 @@ def use_invoke_model_titan():
 
 
 # Anthropic Models
-def use_invoke_model_anthropic():
-    pass
+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)

src/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)

src/langtrace_python_sdk/instrumentation/aws_bedrock/instrumentation.py

@@ -22,12 +22,7 @@
 from opentelemetry.trace import get_tracer
 from wrapt import wrap_function_wrapper as _W
 
-from langtrace_python_sdk.instrumentation.aws_bedrock.patch import (
-    converse,
-    invoke_model,
-    converse_stream,
-    patch_aws_bedrock,
-)
+from langtrace_python_sdk.instrumentation.aws_bedrock.patch import patch_aws_bedrock
 
 logging.basicConfig(level=logging.FATAL)