diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/server/base.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/server/base.py index 8915aadb172b..48bd03a0e88c 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/server/base.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/server/base.py @@ -60,21 +60,18 @@ def set_request_id_to_context_var(self, request): request_context.set(ctx) def set_run_context_to_context_var(self, run_context): - agent_id, agent_name = "", "" - agent_obj = run_context.get_agent_id_object() - if agent_obj: - agent_name = getattr(agent_obj, "name", "") - agent_version = getattr(agent_obj, "version", "") - agent_id = f"{agent_name}:{agent_version}" + agent_name, agent_id = get_agent_name(run_context) res = { "azure.ai.agentserver.response_id": run_context.response_id or "", "azure.ai.agentserver.conversation_id": run_context.conversation_id or "", "azure.ai.agentserver.streaming": str(run_context.stream or False), + "gen_ai.operation.name": "invoke_agent", "gen_ai.agent.id": agent_id, - "gen_ai.agent.name": agent_name, + "gen_ai.agent.name": agent_name or "", "gen_ai.provider.name": "AzureAI Hosted Agents", "gen_ai.response.id": run_context.response_id or "", + "gen_ai.conversation.id": run_context.conversation_id or "", } ctx = request_context.get() or {} ctx.update(res) @@ -87,10 +84,12 @@ async def runs_endpoint(request): # Set up tracing context and span context = request.state.agent_run_context ctx = request_context.get() + agent_name, _ = get_agent_name(context) + span_name = f"invoke_agent {agent_name}" if agent_name else "invoke_agent" with self.tracer.start_as_current_span( - name=f"HostedAgents-{context.response_id}", + name=span_name, attributes=ctx, - kind=trace.SpanKind.SERVER, + kind=trace.SpanKind.CLIENT, ): try: logger.info("Start processing CreateResponse request:") @@ -304,6 +303,30 @@ def setup_otlp_exporter(self, endpoint, provider): logger.info(f"Tracing setup with OTLP exporter: {endpoint}") +def get_agent_name(agent_run_context: AgentRunContext) -> tuple: + """ + Extract the agent name and agent id from an AgentRunContext. + + :param agent_run_context: The AgentRunContext instance to extract from. + :type agent_run_context: AgentRunContext + :return: Tuple of (agent_name, agent_id). Agent id is formatted as "name:version". + Returns (None, "") if agent information is not available. + :rtype: tuple + """ + agent_name = None + agent_id = "" + + if agent_run_context: + agent_obj = agent_run_context.get_agent_id_object() + if agent_obj: + agent_name = getattr(agent_obj, "name", None) + agent_version = getattr(agent_obj, "version", "") + if agent_name: + agent_id = f"{agent_name}:{agent_version}" + + return agent_name, agent_id + + def _event_to_sse_chunk(event: ResponseStreamEvent) -> str: event_data = json.dumps(event.as_dict()) if event.type: diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/test_get_agent_name.py b/sdk/agentserver/azure-ai-agentserver-core/tests/test_get_agent_name.py new file mode 100644 index 000000000000..3bf546e18e1e --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_get_agent_name.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Unit tests for get_agent_name function.""" + +import sys +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +# Add the project root to the path +project_root = Path(__file__).parent.parent +sys.path.insert(0, str(project_root)) + +from azure.ai.agentserver.core.server.base import get_agent_name +from azure.ai.agentserver.core.server.common.agent_run_context import AgentRunContext + + +class TestGetAgentName: + """Test suite for get_agent_name function.""" + + def test_get_agent_name_with_valid_agent_object(self): + """Test get_agent_name with valid agent object.""" + # Create a mock agent object + mock_agent = MagicMock() + mock_agent.name = "TestAgent" + mock_agent.version = "1.0.0" + + # Create a mock AgentRunContext + mock_context = MagicMock(spec=AgentRunContext) + mock_context.get_agent_id_object.return_value = mock_agent + + # Call get_agent_name + agent_name, agent_id = get_agent_name(mock_context) + + # Assert results + assert agent_name == "TestAgent" + assert agent_id == "TestAgent:1.0.0" + + def test_get_agent_name_with_no_version(self): + """Test get_agent_name when agent has no version.""" + # Create a mock agent object with no version + mock_agent = MagicMock() + mock_agent.name = "TestAgent" + mock_agent.version = "" + + # Create a mock AgentRunContext + mock_context = MagicMock(spec=AgentRunContext) + mock_context.get_agent_id_object.return_value = mock_agent + + # Call get_agent_name + agent_name, agent_id = get_agent_name(mock_context) + + # Assert results + assert agent_name == "TestAgent" + assert agent_id == "TestAgent:" + + def test_get_agent_name_with_none_agent_object(self): + """Test get_agent_name when agent object is None.""" + # Create a mock AgentRunContext that returns None + mock_context = MagicMock(spec=AgentRunContext) + mock_context.get_agent_id_object.return_value = None + + # Call get_agent_name + agent_name, agent_id = get_agent_name(mock_context) + + # Assert results + assert agent_name is None + assert agent_id == "" + + def test_get_agent_name_with_none_context(self): + """Test get_agent_name with None context.""" + # Call get_agent_name with None + agent_name, agent_id = get_agent_name(None) + + # Assert results + assert agent_name is None + assert agent_id == "" + + def test_get_agent_name_with_missing_name_attribute(self): + """Test get_agent_name when agent object has no name attribute.""" + # Create a mock agent object without name + mock_agent = MagicMock() + mock_agent.name = None + mock_agent.version = "1.0.0" + + # Create a mock AgentRunContext + mock_context = MagicMock(spec=AgentRunContext) + mock_context.get_agent_id_object.return_value = mock_agent + + # Call get_agent_name + agent_name, agent_id = get_agent_name(mock_context) + + # Assert results - agent_id should not be set if name is None + assert agent_name is None + assert agent_id == "" + + def test_get_agent_name_with_empty_name(self): + """Test get_agent_name when agent name is empty string.""" + # Create a mock agent object with empty name + mock_agent = MagicMock() + mock_agent.name = "" + mock_agent.version = "1.0.0" + + # Create a mock AgentRunContext + mock_context = MagicMock(spec=AgentRunContext) + mock_context.get_agent_id_object.return_value = mock_agent + + # Call get_agent_name + agent_name, agent_id = get_agent_name(mock_context) + + # Assert results - agent_id should not be set if name is empty + assert agent_name == "" + assert agent_id == "" + + def test_get_agent_name_return_type(self): + """Test that get_agent_name returns a tuple.""" + # Create a mock agent object + mock_agent = MagicMock() + mock_agent.name = "TestAgent" + mock_agent.version = "2.0.0" + + # Create a mock AgentRunContext + mock_context = MagicMock(spec=AgentRunContext) + mock_context.get_agent_id_object.return_value = mock_agent + + # Call get_agent_name + result = get_agent_name(mock_context) + + # Assert result is a tuple of length 2 + assert isinstance(result, tuple) + assert len(result) == 2 + + def test_get_agent_name_with_special_characters(self): + """Test get_agent_name with special characters in agent name and version.""" + # Create a mock agent object with special characters + mock_agent = MagicMock() + mock_agent.name = "Test-Agent_v2" + mock_agent.version = "2.0.0-beta.1" + + # Create a mock AgentRunContext + mock_context = MagicMock(spec=AgentRunContext) + mock_context.get_agent_id_object.return_value = mock_agent + + # Call get_agent_name + agent_name, agent_id = get_agent_name(mock_context) + + # Assert results + assert agent_name == "Test-Agent_v2" + assert agent_id == "Test-Agent_v2:2.0.0-beta.1" diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/test_otel_context_attributes.py b/sdk/agentserver/azure-ai-agentserver-core/tests/test_otel_context_attributes.py new file mode 100644 index 000000000000..168574e7746e --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_otel_context_attributes.py @@ -0,0 +1,255 @@ +#!/usr/bin/env python3 +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Unit tests for OTEL context attributes in AgentRunContextMiddleware.""" + +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +# Add the project root to the path +project_root = Path(__file__).parent.parent +sys.path.insert(0, str(project_root)) + +from azure.ai.agentserver.core.server.base import AgentRunContextMiddleware +from azure.ai.agentserver.core.server.common.agent_run_context import AgentRunContext + + +class TestAgentRunContextMiddlewareOtelAttributes: + """Test suite for OTEL attributes in AgentRunContextMiddleware.""" + + @pytest.fixture + def middleware(self): + """Create a middleware instance for testing.""" + mock_app = MagicMock() + return AgentRunContextMiddleware(mock_app) + + def test_set_run_context_with_valid_agent(self, middleware): + """Test that context is set with correct OTEL attributes for valid agent.""" + # Create a mock agent object + mock_agent = MagicMock() + mock_agent.name = "TestAgent" + mock_agent.version = "1.0.0" + + # Create a mock AgentRunContext + mock_context = MagicMock(spec=AgentRunContext) + mock_context.response_id = "resp-123" + mock_context.conversation_id = "conv-456" + mock_context.stream = False + mock_context.get_agent_id_object.return_value = mock_agent + + # Mock request_context + with patch("azure.ai.agentserver.core.server.base.request_context") as mock_req_context: + mock_req_context.get.return_value = {} + + # Call the method + middleware.set_run_context_to_context_var(mock_context) + + # Verify request_context.set was called + mock_req_context.set.assert_called_once() + + # Get the context dict that was set + call_args = mock_req_context.set.call_args + ctx = call_args[0][0] + + # Assert OTEL attributes are set correctly + assert ctx["gen_ai.operation.name"] == "invoke_agent" + assert ctx["gen_ai.agent.id"] == "TestAgent:1.0.0" + assert ctx["gen_ai.agent.name"] == "TestAgent" + assert ctx["gen_ai.provider.name"] == "AzureAI Hosted Agents" + + def test_set_run_context_with_no_agent(self, middleware): + """Test that context is set with empty values when agent is None.""" + # Create a mock AgentRunContext with no agent + mock_context = MagicMock(spec=AgentRunContext) + mock_context.response_id = "resp-123" + mock_context.conversation_id = "conv-456" + mock_context.stream = False + mock_context.get_agent_id_object.return_value = None + + # Mock request_context + with patch("azure.ai.agentserver.core.server.base.request_context") as mock_req_context: + mock_req_context.get.return_value = {} + + # Call the method + middleware.set_run_context_to_context_var(mock_context) + + # Get the context dict that was set + call_args = mock_req_context.set.call_args + ctx = call_args[0][0] + + # Assert OTEL attributes are set with empty values + assert ctx["gen_ai.operation.name"] == "invoke_agent" + assert ctx["gen_ai.agent.id"] == "" + assert ctx["gen_ai.agent.name"] == "" + assert ctx["gen_ai.provider.name"] == "AzureAI Hosted Agents" + + def test_set_run_context_preserves_existing_context(self, middleware): + """Test that existing context values are preserved.""" + # Create a mock agent object + mock_agent = MagicMock() + mock_agent.name = "TestAgent" + mock_agent.version = "1.0.0" + + # Create a mock AgentRunContext + mock_context = MagicMock(spec=AgentRunContext) + mock_context.response_id = "resp-123" + mock_context.conversation_id = "conv-456" + mock_context.stream = True + mock_context.get_agent_id_object.return_value = mock_agent + + # Mock request_context with existing values + existing_ctx = {"existing_key": "existing_value"} + with patch("azure.ai.agentserver.core.server.base.request_context") as mock_req_context: + mock_req_context.get.return_value = existing_ctx.copy() + + # Call the method + middleware.set_run_context_to_context_var(mock_context) + + # Get the context dict that was set + call_args = mock_req_context.set.call_args + ctx = call_args[0][0] + + # Assert existing key is preserved + assert ctx["existing_key"] == "existing_value" + # Assert new keys are added + assert ctx["gen_ai.operation.name"] == "invoke_agent" + assert ctx["gen_ai.agent.name"] == "TestAgent" + + def test_set_run_context_with_streaming(self, middleware): + """Test that streaming flag is correctly set in context.""" + # Create a mock agent object + mock_agent = MagicMock() + mock_agent.name = "TestAgent" + mock_agent.version = "1.0.0" + + # Create a mock AgentRunContext with streaming enabled + mock_context = MagicMock(spec=AgentRunContext) + mock_context.response_id = "resp-123" + mock_context.conversation_id = "conv-456" + mock_context.stream = True + mock_context.get_agent_id_object.return_value = mock_agent + + # Mock request_context + with patch("azure.ai.agentserver.core.server.base.request_context") as mock_req_context: + mock_req_context.get.return_value = {} + + # Call the method + middleware.set_run_context_to_context_var(mock_context) + + # Get the context dict that was set + call_args = mock_req_context.set.call_args + ctx = call_args[0][0] + + # Assert streaming flag is set + assert ctx["azure.ai.agentserver.streaming"] == "True" + + def test_set_run_context_with_none_values(self, middleware): + """Test that None response/conversation IDs are converted to empty strings.""" + # Create a mock agent object + mock_agent = MagicMock() + mock_agent.name = "TestAgent" + mock_agent.version = "1.0.0" + + # Create a mock AgentRunContext with None IDs + mock_context = MagicMock(spec=AgentRunContext) + mock_context.response_id = None + mock_context.conversation_id = None + mock_context.stream = False + mock_context.get_agent_id_object.return_value = mock_agent + + # Mock request_context + with patch("azure.ai.agentserver.core.server.base.request_context") as mock_req_context: + mock_req_context.get.return_value = {} + + # Call the method + middleware.set_run_context_to_context_var(mock_context) + + # Get the context dict that was set + call_args = mock_req_context.set.call_args + ctx = call_args[0][0] + + # Assert None values are converted to empty strings + assert ctx["azure.ai.agentserver.response_id"] == "" + assert ctx["azure.ai.agentserver.conversation_id"] == "" + assert ctx["gen_ai.response.id"] == "" + assert ctx["gen_ai.conversation.id"] == "" + + def test_set_run_context_otel_attributes_completeness(self, middleware): + """Test that all required OTEL attributes are present in context.""" + # Create a mock agent object + mock_agent = MagicMock() + mock_agent.name = "TestAgent" + mock_agent.version = "1.0.0" + + # Create a mock AgentRunContext + mock_context = MagicMock(spec=AgentRunContext) + mock_context.response_id = "resp-123" + mock_context.conversation_id = "conv-456" + mock_context.stream = False + mock_context.get_agent_id_object.return_value = mock_agent + + # Mock request_context + with patch("azure.ai.agentserver.core.server.base.request_context") as mock_req_context: + mock_req_context.get.return_value = {} + + # Call the method + middleware.set_run_context_to_context_var(mock_context) + + # Get the context dict that was set + call_args = mock_req_context.set.call_args + ctx = call_args[0][0] + + # Assert all required attributes are present + required_attributes = [ + "azure.ai.agentserver.response_id", + "azure.ai.agentserver.conversation_id", + "azure.ai.agentserver.streaming", + "gen_ai.operation.name", + "gen_ai.agent.id", + "gen_ai.agent.name", + "gen_ai.provider.name", + "gen_ai.response.id", + "gen_ai.conversation.id", + ] + for attr in required_attributes: + assert attr in ctx, f"Missing required attribute: {attr}" + + def test_set_run_context_agent_id_format(self, middleware): + """Test that agent ID is formatted correctly as 'name:version'.""" + # Test cases with different version values + test_cases = [ + ("Agent1", "1.0.0", "Agent1:1.0.0"), + ("Agent2", "2.1.0-beta", "Agent2:2.1.0-beta"), + ("Agent3", "", "Agent3:"), + ] + + for agent_name, agent_version, expected_id in test_cases: + # Create a mock agent object + mock_agent = MagicMock() + mock_agent.name = agent_name + mock_agent.version = agent_version + + # Create a mock AgentRunContext + mock_context = MagicMock(spec=AgentRunContext) + mock_context.response_id = "resp-123" + mock_context.conversation_id = "conv-456" + mock_context.stream = False + mock_context.get_agent_id_object.return_value = mock_agent + + # Mock request_context + with patch("azure.ai.agentserver.core.server.base.request_context") as mock_req_context: + mock_req_context.get.return_value = {} + + # Call the method + middleware.set_run_context_to_context_var(mock_context) + + # Get the context dict that was set + call_args = mock_req_context.set.call_args + ctx = call_args[0][0] + + # Assert agent ID format + assert ctx["gen_ai.agent.id"] == expected_id