Address review comments

Author: stanek-michal

instrumentation/opentelemetry-instrumentation-botocore/examples/bedrock-runtime/zero-code/invoke_agent.py

@@ -1,23 +1,30 @@
 #!/usr/bin/env python3
 
-import boto3
 import json
-import uuid
 import logging
+import os
+import uuid
+
+import boto3
 from botocore.exceptions import ClientError
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+
 class WeatherAgentDemo:
     def __init__(self):
         # Initialize Bedrock Agents Runtime client
-        self.client = boto3.client('bedrock-agent-runtime')
+        self.client = boto3.client("bedrock-agent-runtime")
 
-        # Replace these with your actual agent IDs
-        self.agent_id = "TFMZVIWXR7"
-        self.agent_alias_id = "ZT7YYHO8TO"
+        self.agent_id = os.environ.get("BEDROCK_AGENT_ID")
+        self.agent_alias_id = os.environ.get("BEDROCK_AGENT_ALIAS_ID")
+
+        if not self.agent_id or not self.agent_alias_id:
+            raise ValueError(
+                "Environment variables BEDROCK_AGENT_ID and BEDROCK_AGENT_ALIAS_ID must be set."
+            )
 
         # Generate a unique session ID for this conversation
         self.session_id = str(uuid.uuid4())
@@ -54,12 +61,15 @@ def invoke_agent(self, input_text, session_state=None, end_session=False):
                     trace_data.append(event["trace"])
                 elif "returnControl" in event:
                     # Handle the case where the agent returns control
-                    return {"type": "return_control", "data": event["returnControl"]}
+                    return {
+                        "type": "return_control",
+                        "data": event["returnControl"],
+                    }
 
             return {
                 "type": "completion",
                 "completion": completion,
-                "trace": trace_data
+                "trace": trace_data,
             }
 
         except ClientError as e:
@@ -73,21 +83,23 @@ def handle_return_control(self, return_control_data):
 
         # Simulate weather API response (using function result format)
         weather_response = {
-            "actionGroup": invocation_inputs["functionInvocationInput"]["actionGroup"],
-            "function": invocation_inputs["functionInvocationInput"]["function"],
+            "actionGroup": invocation_inputs["functionInvocationInput"][
+                "actionGroup"
+            ],
+            "function": invocation_inputs["functionInvocationInput"][
+                "function"
+            ],
             "responseBody": {
-                "TEXT": {
-                    "body": "It's 55F in Seattle today, rainy."
-                }
-            }
+                "TEXT": {"body": "It's 55F in Seattle today, rainy."}
+            },
         }
 
         # Prepare session state with the returned control results
         session_state = {
             "invocationId": invocation_id,
             "returnControlInvocationResults": [
                 {"functionResult": weather_response}
-            ]
+            ],
         }
 
         return session_state
@@ -110,7 +122,7 @@ def run_demo(self):
             final_response = self.invoke_agent(
                 "Process the weather data",
                 session_state=session_state,
-                end_session=False # Keeping session open for potential follow-up
+                end_session=False,  # Keeping session open for potential follow-up
             )
 
             if final_response["type"] == "completion":
@@ -122,11 +134,13 @@ def run_demo(self):
                         print(json.dumps(trace, indent=2))
             else:
                 # Should not happen if session_state is provided correctly
-                logger.warning(f"Unexpected response type after providing session state: {final_response.get('type')}")
+                logger.warning(
+                    f"Unexpected response type after providing session state: {final_response.get('type')}"
+                )
                 print("Received unexpected response after providing data.")
 
         elif response["type"] == "completion":
-             # Agent answered directly without needing function call/return control
+            # Agent answered directly without needing function call/return control
             print("\nAgent responded directly:")
             print(response["completion"])
             if response["trace"]:
@@ -135,12 +149,17 @@ def run_demo(self):
                     print(json.dumps(trace, indent=2))
         else:
             # Handle other unexpected response types
-            logger.error(f"Unexpected initial response type: {response.get('type')}")
+            logger.error(
+                f"Unexpected initial response type: {response.get('type')}"
+            )
             print(f"Received unexpected response type: {response.get('type')}")
 
+
 if __name__ == "__main__":
     try:
         demo = WeatherAgentDemo()
         demo.run_demo()
     except Exception as e:
-        logger.error(f"Demo failed unexpectedly: {e}", exc_info=True) # Add traceback info
+        logger.error(
+            f"Demo failed unexpectedly: {e}", exc_info=True
+        )  # Add traceback info

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/__init__.py

@@ -32,8 +32,10 @@ def loader():
 
 
 _KNOWN_EXTENSIONS = {
+    "bedrock-agent-runtime": _lazy_load(
+        ".bedrock", "_BedrockRuntimeExtension"
+    ),
     "bedrock-runtime": _lazy_load(".bedrock", "_BedrockRuntimeExtension"),
-    "bedrock-agent-runtime": _lazy_load(".bedrock", "_BedrockRuntimeExtension"),
     "dynamodb": _lazy_load(".dynamodb", "_DynamoDbExtension"),
     "lambda": _lazy_load(".lmbd", "_LambdaExtension"),
     "sns": _lazy_load(".sns", "_SnsExtension"),

instrumentation/opentelemetry-instrumentation-botocore/src/opentelemetry/instrumentation/botocore/extensions/bedrock.py

@@ -45,6 +45,8 @@
     ERROR_TYPE,
 )
 from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import (
+    GEN_AI_AGENT_ID,
+    GEN_AI_AGENT_NAME,
     GEN_AI_OPERATION_NAME,
     GEN_AI_REQUEST_MAX_TOKENS,
     GEN_AI_REQUEST_MODEL,
@@ -54,18 +56,15 @@
     GEN_AI_RESPONSE_FINISH_REASONS,
     GEN_AI_SYSTEM,
     GEN_AI_TOKEN_TYPE,
-    GEN_AI_USAGE_INPUT_TOKENS,
-    GEN_AI_USAGE_OUTPUT_TOKENS,
-    GEN_AI_AGENT_ID,
-    GEN_AI_AGENT_NAME,
     GEN_AI_TOOL_CALL_ID,
     GEN_AI_TOOL_NAME,
     GEN_AI_TOOL_TYPE,
+    GEN_AI_USAGE_INPUT_TOKENS,
+    GEN_AI_USAGE_OUTPUT_TOKENS,
     GenAiOperationNameValues,
     GenAiSystemValues,
     GenAiTokenTypeValues,
 )
-
 from opentelemetry.semconv._incubating.metrics.gen_ai_metrics import (
     GEN_AI_CLIENT_OPERATION_DURATION,
     GEN_AI_CLIENT_TOKEN_USAGE,
@@ -189,7 +188,9 @@ def extract_attributes(self, attributes: _AttributeMapT):
             agent_alias_id = self._call_context.params.get("agentAliasId")
 
             self._set_if_not_none(attributes, GEN_AI_AGENT_ID, agent_id)
-            self._set_if_not_none(attributes, GEN_AI_AGENT_NAME, agent_alias_id)
+            self._set_if_not_none(
+                attributes, GEN_AI_AGENT_NAME, agent_alias_id
+            )
             return
 
         # Handle non-agent chat completions
@@ -506,7 +507,9 @@ def _invoke_agent_on_success(
         instrumentor_context: _BotocoreInstrumentorContext,
     ):
         try:
-            if "completion" in result and isinstance(result["completion"], EventStream):
+            if "completion" in result and isinstance(
+                result["completion"], EventStream
+            ):
                 event_stream = result["completion"]
 
                 # Drain the stream so we can instrument AND keep events
@@ -522,15 +525,15 @@ def _invoke_agent_on_success(
             # Record metrics
             metrics = instrumentor_context.metrics
             metrics_attributes = self._extract_metrics_attributes()
-            if operation_duration_histogram := metrics.get(GEN_AI_CLIENT_OPERATION_DURATION):
+            if operation_duration_histogram := metrics.get(
+                GEN_AI_CLIENT_OPERATION_DURATION
+            ):
                 duration = max((default_timer() - self._operation_start), 0)
                 operation_duration_histogram.record(
                     duration,
                     attributes=metrics_attributes,
                 )
 
-        except json.JSONDecodeError:
-            _logger.debug("Error: Unable to parse the response body as JSON")
         except Exception as exc:  # pylint: disable=broad-exception-caught
             _logger.debug("Error processing response: %s", exc)
 
@@ -548,7 +551,9 @@ def _handle_return_control(self, span: Span, event: dict):
                     func_input = input_item["functionInvocationInput"]
                     action_group = func_input.get("actionGroup")
                     function = func_input.get("function")
-                    span.set_attribute(GEN_AI_TOOL_NAME, action_group)
+                    span.set_attribute(
+                        GEN_AI_TOOL_NAME, action_group + "." + function
+                    )
                     span.set_attribute(GEN_AI_TOOL_TYPE, "function")
 
                 # Handle API invocation
@@ -847,6 +852,7 @@ def on_error(
                 attributes=metrics_attributes,
             )
 
+
 def _replay_events(events):
     """
     Helper so that user can still iterate EventStream

instrumentation/opentelemetry-instrumentation-botocore/test-requirements-0.txt

@@ -37,5 +37,3 @@ zipp==3.19.2
 -e opentelemetry-instrumentation
 -e propagator/opentelemetry-propagator-aws-xray
 -e instrumentation/opentelemetry-instrumentation-botocore
-git+https://github.com/open-telemetry/opentelemetry-python.git@main#subdirectory=opentelemetry-semantic-conventions
-

instrumentation/opentelemetry-instrumentation-botocore/test-requirements-1.txt

@@ -37,5 +37,3 @@ zipp==3.19.2
 -e opentelemetry-instrumentation
 -e propagator/opentelemetry-propagator-aws-xray
 -e instrumentation/opentelemetry-instrumentation-botocore
-git+https://github.com/open-telemetry/opentelemetry-python.git@main#subdirectory=opentelemetry-semantic-conventions
-

instrumentation/opentelemetry-instrumentation-botocore/tests/bedrock_utils.py

@@ -273,31 +273,52 @@ def assert_message_in_logs(log, event_name, expected_content, parent_span):
     assert_log_parent(log, parent_span)
 
 
-def assert_invoke_agent_attributes(span, agent_id, agent_alias_id, session_id, has_tool_call=False, is_result_call=False):
+def assert_invoke_agent_attributes(
+    span,
+    agent_id,
+    agent_alias_id,
+    session_id,
+    has_tool_call=False,
+    is_result_call=False,
+):
     # Check system and operation name
-    assert span.attributes.get(GenAIAttributes.GEN_AI_SYSTEM) == GenAIAttributes.GenAiSystemValues.AWS_BEDROCK.value
-    assert span.attributes.get(GenAIAttributes.GEN_AI_OPERATION_NAME) == "invoke_agent"
+    assert (
+        span.attributes.get(GenAIAttributes.GEN_AI_SYSTEM)
+        == GenAIAttributes.GenAiSystemValues.AWS_BEDROCK.value
+    )
+    assert (
+        span.attributes.get(GenAIAttributes.GEN_AI_OPERATION_NAME)
+        == "invoke_agent"
+    )
 
     # Check agent attributes
     assert span.attributes.get(GenAIAttributes.GEN_AI_AGENT_ID) == agent_id
-    assert span.attributes.get(GenAIAttributes.GEN_AI_AGENT_NAME) == agent_alias_id
+    assert (
+        span.attributes.get(GenAIAttributes.GEN_AI_AGENT_NAME)
+        == agent_alias_id
+    )
 
     # If tool call exists, check tool attributes
     if has_tool_call:
         assert GenAIAttributes.GEN_AI_TOOL_CALL_ID in span.attributes
         assert GenAIAttributes.GEN_AI_TOOL_NAME in span.attributes
         assert GenAIAttributes.GEN_AI_TOOL_TYPE in span.attributes
         allowed_tool_types = {"extension", "function", "datastore"}
-        assert span.attributes.get(GenAIAttributes.GEN_AI_TOOL_TYPE) in allowed_tool_types, \
-            f"Unexpected tool type in span: {span.attributes.get(GenAIAttributes.GEN_AI_TOOL_TYPE)}"
+        assert (
+            span.attributes.get(GenAIAttributes.GEN_AI_TOOL_TYPE)
+            in allowed_tool_types
+        ), f"Unexpected tool type in span: {span.attributes.get(GenAIAttributes.GEN_AI_TOOL_TYPE)}"
     elif is_result_call:
         assert GenAIAttributes.GEN_AI_TOOL_CALL_ID not in span.attributes
         assert GenAIAttributes.GEN_AI_TOOL_NAME not in span.attributes
         assert GenAIAttributes.GEN_AI_TOOL_TYPE not in span.attributes
 
 
 def assert_all_metric_attributes(
-    data_point, operation_name: str, model: str, error_type: str | None = None
+    data_point,
+    operation_name: str,
+    model: str | None,
+    error_type: str | None = None,
 ):
     assert GenAIAttributes.GEN_AI_OPERATION_NAME in data_point.attributes
     assert (
@@ -311,7 +332,10 @@ def assert_all_metric_attributes(
     )
     if model is not None:
         assert GenAIAttributes.GEN_AI_REQUEST_MODEL in data_point.attributes
-        assert data_point.attributes[GenAIAttributes.GEN_AI_REQUEST_MODEL] == model
+        assert (
+            data_point.attributes[GenAIAttributes.GEN_AI_REQUEST_MODEL]
+            == model
+        )
 
     if error_type is not None:
         assert ERROR_TYPE in data_point.attributes

instrumentation/opentelemetry-instrumentation-botocore/tests/conftest.py

@@ -77,9 +77,10 @@ def fixture_meter_provider(metric_reader):
 def bedrock_runtime_client():
     return boto3.client("bedrock-runtime")
 
+
 @pytest.fixture
 def bedrock_agent_client():
-    return boto3.client('bedrock-agent-runtime')
+    return boto3.client("bedrock-agent-runtime")
 
 
 @pytest.fixture(autouse=True)