kickstart refactoring

alizenhom · alizenhom · commit 545859bf92ef · 2025-01-08T14:45:04.000+02:00
diff --git a/src/examples/awsbedrock_examples/__init__.py b/src/examples/awsbedrock_examples/__init__.py
@@ -1,4 +1,4 @@
-from examples.awsbedrock_examples.converse import use_converse
+from examples.awsbedrock_examples.converse import use_converse, use_invoke_model_titan
 from langtrace_python_sdk import langtrace, with_langtrace_root_span
 
 langtrace.init()
@@ -7,4 +7,5 @@
 class AWSBedrockRunner:
     @with_langtrace_root_span("AWS_Bedrock")
     def run(self):
-        use_converse()
+        # use_converse()
+        use_invoke_model_titan()
diff --git a/src/examples/awsbedrock_examples/converse.py b/src/examples/awsbedrock_examples/converse.py
@@ -1,17 +1,19 @@
-import os
 import boto3
+import botocore
+import json
 from langtrace_python_sdk import langtrace
+from dotenv import load_dotenv
+
+
+load_dotenv()
+langtrace.init()
+
+brt = boto3.client("bedrock-runtime", region_name="us-east-1")
+brc = boto3.client("bedrock", region_name="us-east-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",
@@ -20,15 +22,70 @@ def use_converse():
     ]
 
     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():
+    models = []
+    for model in brc.list_foundation_models()["modelSummaries"]:
+        models.append(model["modelId"])
+    return models
+
+
+# Invoke Model API
+# Amazon Titan Models
+def use_invoke_model_titan():
+    try:
+        prompt_data = "what's 1+1?"
+        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"
+
+        response = brt.invoke_model(
+            body=body, modelId=modelId, accept=accept, contentType=contentType
+        )
+        response_body = json.loads(response.get("body").read())
+
+        # print(response_body.get("results"))
+
+    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():
+    pass
+
+
+# print(get_foundation_models())
diff --git a/src/langtrace_python_sdk/constants/instrumentation/aws_bedrock.py b/src/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",
diff --git a/src/langtrace_python_sdk/instrumentation/aws_bedrock/instrumentation.py b/src/langtrace_python_sdk/instrumentation/aws_bedrock/instrumentation.py
@@ -23,21 +23,17 @@
 from wrapt import wrap_function_wrapper as _W
 
 from langtrace_python_sdk.instrumentation.aws_bedrock.patch import (
-    converse, converse_stream
+    converse,
+    invoke_model,
+    converse_stream,
+    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 +42,11 @@ def _instrument(self, **kwargs):
         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
diff --git a/src/langtrace_python_sdk/instrumentation/aws_bedrock/patch.py b/src/langtrace_python_sdk/instrumentation/aws_bedrock/patch.py
@@ -48,30 +48,42 @@ def wrapper(original_method):
             @wraps(original_method)
             def wrapped_method(*args, **kwargs):
                 service_provider = SERVICE_PROVIDERS["AWS_BEDROCK"]
-
+                print("Here's the kwargs: ", kwargs)
                 input_content = [
                     {
-                        'role': message.get('role', 'user'),
-                        'content': message.get('content', [])[0].get('text', "")
+                        "role": message.get("role", "user"),
+                        "content": message.get("content", [])[0].get("text", ""),
                     }
-                    for message in kwargs.get('messages', [])
+                    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_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'),
-                    })
+                    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)
 
@@ -92,20 +104,22 @@ def wrapped_method(*args, **kwargs):
                         raise err
 
             return wrapped_method
+
         return wrapper
+
     return decorator
 
 
 converse = traced_aws_bedrock_call("CONVERSE", "converse")
+invoke_model = traced_aws_bedrock_call("INVOKE_MODEL", "invoke_model")
 
 
 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 +143,87 @@ def traced_method(wrapped, instance, args, kwargs):
                 span.record_exception(err)
                 span.set_status(Status(StatusCode.ERROR, str(err)))
                 raise err
-            
+
     return traced_method
 
 
 @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 patch_aws_bedrock(tracer, version):
+    def traced_method(wrapped, instance, args, kwargs):
+        if args and args[0] != "bedrock-runtime":
+            return
+
+        client = wrapped(*args, **kwargs)
+        print("Here's the client: ", client)
+        client.invoke_model = patch_invoke_model(client.invoke_model, tracer, version)
+        client.invoke_model_with_response_stream = patch_invoke_model(
+            client.invoke_model_with_response_stream, tracer, version
+        )
+        client.converse = patch_invoke_model(client.converse, tracer, version)
+        client.converse_stream = patch_invoke_model(
+            client.converse_stream, tracer, version
         )
+        return client
+
+    return traced_method
+
+
+def patch_invoke_model(original_method, tracer, version):
+    def traced_method(*args, **kwargs):
+        service_provider = SERVICE_PROVIDERS["AWS_BEDROCK"]
+        span_attributes = {
+            **get_langtrace_attributes(
+                version, service_provider, 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)
+            set_invoke_model_attributes(span, kwargs)
+            response = original_method(*args, **kwargs)
+            return response
+
+    return traced_method
+
+
+def set_invoke_model_attributes(span, kwargs):
+    modelId = kwargs.get("modelId")
+    (vendor, model_name) = modelId.split(".")
+    
+    print("Here's the vendor: ", vendor)
+    print("Here's the model_name: ", model_name)
+    print("Here's the kwargs: ", kwargs)
diff --git a/src/run_example.py b/src/run_example.py
@@ -4,7 +4,7 @@
     "anthropic": False,
     "azureopenai": False,
     "chroma": False,
-    "cohere": True,
+    "cohere": False,
     "fastapi": False,
     "langchain": False,
     "llamaindex": False,
@@ -20,7 +20,7 @@
     "vertexai": False,
     "gemini": False,
     "mistral": False,
-    "awsbedrock": False,
+    "awsbedrock": True,
     "cerebras": False,
 }
 
