Fix region_name and is_recording problem

traceloop · Sep 19, 2024 · 6f5a11b · 6f5a11b
1 parent 934e99e
commit 6f5a11b
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 22 deletions.
diff --git a/...s/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/__init__.py b/...s/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/__init__.py
@@ -240,7 +240,7 @@ def _handle_call(span, kwargs, response, metricParams):
 
     metricParams.vendor = vendor
     metricParams.model = model
-    metricParams.isStream = True
+    metricParams.isStream = False
 
     _set_span_attribute(span, SpanAttributes.LLM_SYSTEM, vendor)
     _set_span_attribute(span, SpanAttributes.LLM_REQUEST_MODEL, model)

diff --git a/packages/opentelemetry-instrumentation-bedrock/tests/conftest.py b/packages/opentelemetry-instrumentation-bedrock/tests/conftest.py
@@ -24,6 +24,16 @@ def brt():
     )
 
 
+@pytest.fixture
+def brt2():
+    return boto3.client(
+        service_name="bedrock-runtime",
+        aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID"),
+        aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"),
+        region_name="us-west-2",
+    )
+
+
 @pytest.fixture(scope="module")
 def vcr_config():
     return {"filter_headers": ["authorization"]}
diff --git a/...ation-bedrock/tests/metrics/cassettes/test_bedrock_metrics/test_invoke_model_metrics.yaml b/...ation-bedrock/tests/metrics/cassettes/test_bedrock_metrics/test_invoke_model_metrics.yaml
@@ -13,38 +13,42 @@ interactions:
         YXBwbGljYXRpb24vanNvbg==
       User-Agent:
       - !!binary |
-        Qm90bzMvMS4zNS4yMSBtZC9Cb3RvY29yZSMxLjM1LjIxIHVhLzIuMCBvcy9tYWNvcyMyMy42LjAg
-        bWQvYXJjaCNhcm02NCBsYW5nL3B5dGhvbiMzLjkuNiBtZC9weWltcGwjQ1B5dGhvbiBjZmcvcmV0
-        cnktbW9kZSNsZWdhY3kgQm90b2NvcmUvMS4zNS4yMQ==
+        Qm90bzMvMS4zNC4xNjIgbWQvQm90b2NvcmUjMS4zNC4xNjIgdWEvMi4wIG9zL21hY29zIzIzLjYu
+        MCBtZC9hcmNoI2FybTY0IGxhbmcvcHl0aG9uIzMuMTEuNSBtZC9weWltcGwjQ1B5dGhvbiBjZmcv
+        cmV0cnktbW9kZSNsZWdhY3kgQm90b2NvcmUvMS4zNC4xNjI=
       X-Amz-Date:
       - !!binary |
-        MjAyNDA5MTdUMjMyOTQzWg==
+        MjAyNDA5MTlUMDgxNTMyWg==
       amz-sdk-invocation-id:
       - !!binary |
-        NTE5ZWJlYzUtNDMzYy00YzIzLWJmNmItYWZhMGMyNjBmMzQx
+        YmNjNTcxNjYtYjRmOC00MzQzLTg4YmYtNWE4YzBhODBmZTM5
       amz-sdk-request:
       - !!binary |
         YXR0ZW1wdD0x
     method: POST
-    uri: https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke
+    uri: https://bedrock-runtime.us-west-2.amazonaws.com/model/amazon.titan-text-express-v1/invoke
   response:
     body:
-      string: '{"message":"You don''t have access to the model with the specified
-        model ID."}'
+      string: '{"inputTextTokenCount":9,"results":[{"tokenCount":17,"outputText":"\nWhat
+        do you call a bear with no teeth?\nA gummy bear.","completionReason":"FINISH"}]}'
     headers:
       Connection:
       - keep-alive
       Content-Length:
-      - '77'
+      - '154'
       Content-Type:
       - application/json
       Date:
-      - Tue, 17 Sep 2024 23:29:43 GMT
-      x-amzn-ErrorType:
-      - AccessDeniedException:http://internal.amazon.com/coral/com.amazon.bedrock/
+      - Thu, 19 Sep 2024 08:15:34 GMT
+      X-Amzn-Bedrock-Input-Token-Count:
+      - '9'
+      X-Amzn-Bedrock-Invocation-Latency:
+      - '1229'
+      X-Amzn-Bedrock-Output-Token-Count:
+      - '17'
       x-amzn-RequestId:
-      - 4caffe59-bb3e-4e98-abe9-f761237b02db
+      - 6278b8bf-c1a7-46d7-822b-9f85ee7805b6
     status:
-      code: 403
-      message: Forbidden
+      code: 200
+      message: OK
 version: 1
diff --git a/packages/opentelemetry-instrumentation-bedrock/tests/metrics/conftest.py b/packages/opentelemetry-instrumentation-bedrock/tests/metrics/conftest.py
@@ -7,6 +7,9 @@
 from opentelemetry.sdk.metrics.export import InMemoryMetricReader
 from opentelemetry.instrumentation.bedrock import BedrockInstrumentor
 
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+
 
 @pytest.fixture(scope="session")
 def metrics_test_context():
@@ -16,7 +19,11 @@ def metrics_test_context():
 
     metrics.set_meter_provider(provider)
 
-    BedrockInstrumentor().instrument()
+    # Without the following line, span.is_recording() is False
+    # so that _handle_call and _handle_stream_call will be skipped
+    trace.set_tracer_provider(TracerProvider())
+
+    BedrockInstrumentor(enrich_token_usage=True).instrument()
 
     return provider, reader
 

diff --git a/packages/opentelemetry-instrumentation-bedrock/tests/metrics/test_bedrock_metrics.py b/packages/opentelemetry-instrumentation-bedrock/tests/metrics/test_bedrock_metrics.py
@@ -5,12 +5,12 @@
 
 
 @pytest.mark.vcr
-def test_invoke_model_metrics(metrics_test_context, brt):
-    if brt is None:
+def test_invoke_model_metrics(metrics_test_context, brt2):
+    if brt2 is None:
         print("test_invoke_model_metrics test skipped.")
         return
 
-    provider, reader = metrics_test_context
+    _, reader = metrics_test_context
 
     body = json.dumps(
         {
@@ -23,7 +23,7 @@ def test_invoke_model_metrics(metrics_test_context, brt):
         }
     )
 
-    brt.invoke_model(
+    brt2.invoke_model(
         body=body,
         modelId='amazon.titan-text-express-v1',
         accept='application/json',
@@ -48,7 +48,7 @@ def test_invoke_model_metrics(metrics_test_context, brt):
                             "output",
                             "input",
                         ]
-                        assert data_point.value > 0
+                        assert data_point.sum > 0
 
                 if metric.name == Meters.LLM_OPERATION_DURATION:
                     found_duration_metric = True