Update botocore.py sink to add support for InvokeModel

Author: bitterpanda63

aikido_zen/sinks/botocore.py

@@ -4,6 +4,20 @@
 from aikido_zen.sinks import after, on_import, patch_function, before
 
 
+def get_tokens_from_converse(api_response):
+    usage = api_response.get("usage", {})
+    input_tokens = usage.get("inputTokens", 0)
+    output_tokens = usage.get("outputTokens", 0)
+    return int(input_tokens), int(output_tokens)
+
+
+def get_tokens_from_invoke_model(api_response):
+    headers = api_response.get("ResponseMetadata", {}).get("HTTPHeaders", {})
+    input_tokens_str = headers.get("x-amzn-bedrock-input-token-count", "0")
+    output_tokens_str = headers.get("x-amzn-bedrock-output-token-count", "0")
+    return int(input_tokens_str), int(output_tokens_str)
+
+
 @after
 def make_api_call_after(func, instance, args, kwargs, return_value):
     # Extract arguments to validate later
@@ -21,13 +35,13 @@ def make_api_call_after(func, instance, args, kwargs, return_value):
     if not model_id:
         return None
 
-    usage = return_value.get("usage", {})
-    on_ai_call(
-        provider="bedrock",
-        model=model_id,
-        input_tokens=usage.get("inputTokens", 0),
-        output_tokens=usage.get("outputTokens", 0),
-    )
+    input_tokens, output_tokens = (0, 0)
+    if operation_name == "Converse":
+        input_tokens, output_tokens = get_tokens_from_converse(return_value)
+    elif operation_name == "InvokeModel":
+        input_tokens, output_tokens = get_tokens_from_invoke_model(return_value)
+
+    on_ai_call("bedrock", model_id, input_tokens, output_tokens)
 
 
 @on_import("botocore.client")

aikido_zen/sinks/tests/aws_bedrock_test.py

@@ -21,7 +21,7 @@ def setup():
 
 @pytest.fixture
 def client():
-    client = boto3.client(service_name="bedrock-runtime", region_name="eu-west-1")
+    client = boto3.client(service_name="bedrock-runtime", region_name="us-east-1")
     return client
 
 
@@ -48,51 +48,33 @@ def test_boto3_converse(client):
     )
     output = response["output"]["message"]["content"][0]["text"]
 
-    assert get_ai_stats()[0]["model"] == "claude-3-sonnet-20240229"
+    assert get_ai_stats()[0]["model"] == "anthropic.claude-3-sonnet-20240229-v1:0"
     assert get_ai_stats()[0]["calls"] == 1
-    assert get_ai_stats()[0]["provider"] == "anthropic"
+    assert get_ai_stats()[0]["provider"] == "bedrock"
     assert get_ai_stats()[0]["tokens"]["input"] == 13
     assert get_ai_stats()[0]["tokens"]["output"] == 20
     assert get_ai_stats()[0]["tokens"]["total"] == 33
 
 
 @skip_no_api_key
 def test_boto3_invoke_model_claude_3_sonnet(client):
-    model_id = (
-        "anthropic.claude-3-sonnet-20240229-v1:0"  # Example model ID for Amazon Bedrock
-    )
+    model_id = "us.anthropic.claude-3-5-sonnet-20241022-v2:0"  # Example model ID for Amazon Bedrock
     input_payload = {
-        "inputText": "Hello, how are you?",
-        "textGenerationConfig": {"maxTokenCount": 100, "temperature": 0.7},
-    }
-    response = client.invoke_model(modelId=model_id, body=json.dumps(input_payload))
-    stats = get_ai_stats()[0]
-    assert stats["model"] == "llama2-70b-chat"
-    assert stats["calls"] == 1
-    assert stats["provider"] == "meta"
-    assert stats["tokens"]["input"] == 0
-    assert stats["tokens"]["output"] == 30
-
-
-@skip_no_api_key
-def test_boto3_invoke_model_meta_llama3_8b(client):
-    metadata = {
-        "model": "meta.llama3-8b-instruct-v1:0",
-        "prompt": "Who painted the Mona Lisa?",
+        "messages": [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": "Are tomatoes a vegetable?"}],
+            }
+        ],
         "max_tokens": 20,
+        "anthropic_version": "bedrock-2023-05-31",
     }
-    body = {
-        "prompt": metadata["prompt"],
-        "max_gen_len": metadata["max_tokens"],
-        "temperature": 0.5,
-    }
-    response = client.invoke_model(
-        modelId=metadata["model"],
-        body=json.dumps(body).encode("utf-8"),
-    )
+    response = client.invoke_model(modelId=model_id, body=json.dumps(input_payload))
+    print(response)
     stats = get_ai_stats()[0]
-    assert stats["model"] == "llama3-8b-instruct"
+    assert stats["model"] == "us.anthropic.claude-3-5-sonnet-20241022-v2:0"
     assert stats["calls"] == 1
-    assert stats["provider"] == "meta"
-    assert stats["tokens"]["input"] == 3
+    assert stats["provider"] == "bedrock"
+    assert stats["tokens"]["input"] == 14
     assert stats["tokens"]["output"] == 20
+    assert stats["tokens"]["total"] == 34