diff --git a/src/app/endpoints/streaming_query.py b/src/app/endpoints/streaming_query.py index b83ae47af..02a6f5479 100644 --- a/src/app/endpoints/streaming_query.py +++ b/src/app/endpoints/streaming_query.py @@ -14,6 +14,7 @@ APIStatusError as LLSApiStatusError, ) from openai._exceptions import APIStatusError as OpenAIAPIStatusError +from opentelemetry import trace from authentication import get_auth_dependency from authentication.interface import AuthTuple @@ -65,6 +66,13 @@ ) from utils.mcp_headers import McpHeaders, mcp_headers_dependency from utils.mcp_oauth_probe import check_mcp_auth +from utils.otel_tracing import ( + SpanAttributes, + SpanEvents, + add_span_event, + anonymize_value, + set_span_attributes, +) from utils.query import ( extract_provider_and_model_from_model_id, handle_known_apistatus_errors, @@ -93,6 +101,7 @@ from utils.vector_search import build_rag_context logger = get_logger(__name__) +tracer = trace.get_tracer(__name__) router = APIRouter(tags=["streaming_query"]) # Tracks background topic summary tasks for graceful shutdown. @@ -158,11 +167,55 @@ async def streaming_query_endpoint_handler( # pylint: disable=too-many-locals - 500: Internal Server Error - Configuration not loaded or other server errors - 503: Service Unavailable - Unable to connect to OGX backend """ + root_span = tracer.start_span("streaming_query.handle_request") + try: + return await _handle_streaming_query_with_tracing( + request, query_request, auth, mcp_headers, root_span + ) + except Exception: + root_span.end() + raise + + +async def _handle_streaming_query_with_tracing( # pylint: disable=too-many-locals + request: Request, + query_request: QueryRequest, + auth: AuthTuple, + mcp_headers: McpHeaders, + root_span: trace.Span, +) -> StreamingResponse: + """Handle streaming query request with OTEL tracing instrumentation. + + Parameters: + request: The incoming HTTP request. + query_request: Request payload containing query and optional parameters. + auth: Authentication tuple (user_id, username, skip_check, token). + mcp_headers: Headers to be passed to MCP servers. + root_span: OpenTelemetry root span for this request. + + Returns: + StreamingResponse with SSE-formatted events. + + Raises: + HTTPException: On authentication, authorization, quota, or model errors. + """ check_configuration_loaded(configuration) user_id, _user_name, _skip_userid_check, token = auth started_at = datetime.datetime.now(datetime.UTC).strftime("%Y-%m-%dT%H:%M:%SZ") + # Set initial span attributes + set_span_attributes( + root_span, + { + SpanAttributes.USER_ID: anonymize_value(user_id), + SpanAttributes.INPUT: anonymize_value(query_request.query), + SpanAttributes.REQUEST_ATTACHMENTS_COUNT: ( + len(query_request.attachments) if query_request.attachments else 0 + ), + }, + ) + # Check MCP Auth await check_mcp_auth(configuration, mcp_headers, token, request.headers) @@ -181,6 +234,9 @@ async def streaming_query_endpoint_handler( # pylint: disable=too-many-locals if query_request.attachments: validate_attachments_metadata(query_request.attachments) + # Validation completed + add_span_event(root_span, SpanEvents.VALIDATION_COMPLETED) + # Retrieve conversation if conversation_id is provided user_conversation = None if query_request.conversation_id: @@ -291,6 +347,7 @@ async def streaming_query_endpoint_handler( # pylint: disable=too-many-locals responses_params=responses_params, endpoint_path=endpoint_path, image_attachments=image_attachments, + root_span=root_span, ), media_type=response_media_type, ) @@ -316,6 +373,7 @@ async def streaming_query_endpoint_handler( # pylint: disable=too-many-locals responses_params=responses_params, turn_summary=turn_summary, background_topic_summary_tasks=_background_topic_summary_tasks, + root_span=root_span, ), media_type=response_media_type, ) @@ -344,6 +402,7 @@ async def generate_response_with_compaction( responses_params: ResponsesApiParams, endpoint_path: str, image_attachments: Optional[list[Attachment]] = None, + root_span: Optional[trace.Span] = None, ) -> AsyncIterator[str]: """Stream a response for a conversation that requires compaction. @@ -359,79 +418,85 @@ async def generate_response_with_compaction( responses_params: The base Responses API parameters. endpoint_path: API endpoint path used for metric labeling. image_attachments: Image attachments for multimodal prompt construction. + root_span: OpenTelemetry root span for this request. Yields: SSE-formatted strings. """ - media_type = context.query_request.media_type or MEDIA_TYPE_JSON - yield stream_start_event( - conversation_id=context.conversation_id, - request_id=context.request_id, - ) - - compacted_original_input: Optional[ResponseInput] = None try: - async for item in apply_compaction( - context.client, - responses_params, - configuration.inference, - configuration.compaction, - emit_events=True, - cache=configured_conversation_cache(), - user_id=context.user_id, - skip_user_id_check=context.skip_userid_check, - ): - if isinstance(item, CompactionStartedEvent): - yield stream_compaction_event(context.conversation_id) - elif isinstance(item, CompactionResult): - responses_params = item.params - compacted_original_input = item.original_input - - generator, turn_summary = await retrieve_agent_response_generator( - responses_params=responses_params, - context=context, - endpoint_path=endpoint_path, - image_attachments=image_attachments, - ) - except HTTPException as e: - yield http_exception_stream_event(e) - return - except RuntimeError as e: # library mode wraps 413 into runtime error - error_response = ( - PromptTooLongResponse(model=responses_params.model) - if is_context_length_error(str(e)) - else InternalServerErrorResponse.generic() - ) - yield stream_http_error_event(error_response, media_type) - return - except APIConnectionError as e: - yield stream_http_error_event( - ServiceUnavailableResponse(backend_name="OGX", cause=str(e)), - media_type, - ) - return - except (LLSApiStatusError, OpenAIAPIStatusError) as e: - yield stream_http_error_event( - handle_known_apistatus_errors(e, responses_params.model), media_type - ) - return - - # Combine inline RAG results (BYOK + Solr) with tool-based results - if context.moderation_result.decision == "passed": - turn_summary.referenced_documents = deduplicate_referenced_documents( - context.inline_rag_context.referenced_documents - + turn_summary.referenced_documents + media_type = context.query_request.media_type or MEDIA_TYPE_JSON + yield stream_start_event( + conversation_id=context.conversation_id, + request_id=context.request_id, ) - # The start event was already emitted above; delegate the rest (re-yield, - # finalization, compacted-turn storage) to the shared generator. - async for event in generate_agent_response( - generator, - context, - responses_params, - turn_summary, - background_topic_summary_tasks=_background_topic_summary_tasks, - emit_start=False, - original_input=compacted_original_input, - ): - yield event + compacted_original_input: Optional[ResponseInput] = None + try: + async for item in apply_compaction( + context.client, + responses_params, + configuration.inference, + configuration.compaction, + emit_events=True, + cache=configured_conversation_cache(), + user_id=context.user_id, + skip_user_id_check=context.skip_userid_check, + ): + if isinstance(item, CompactionStartedEvent): + yield stream_compaction_event(context.conversation_id) + elif isinstance(item, CompactionResult): + responses_params = item.params + compacted_original_input = item.original_input + + generator, turn_summary = await retrieve_agent_response_generator( + responses_params=responses_params, + context=context, + endpoint_path=endpoint_path, + image_attachments=image_attachments, + ) + except HTTPException as e: + yield http_exception_stream_event(e) + return + except RuntimeError as e: # library mode wraps 413 into runtime error + error_response = ( + PromptTooLongResponse(model=responses_params.model) + if is_context_length_error(str(e)) + else InternalServerErrorResponse.generic() + ) + yield stream_http_error_event(error_response, media_type) + return + except APIConnectionError as e: + yield stream_http_error_event( + ServiceUnavailableResponse(backend_name="OGX", cause=str(e)), + media_type, + ) + return + except (LLSApiStatusError, OpenAIAPIStatusError) as e: + yield stream_http_error_event( + handle_known_apistatus_errors(e, responses_params.model), media_type + ) + return + + # Combine inline RAG results (BYOK + Solr) with tool-based results + if context.moderation_result.decision == "passed": + turn_summary.referenced_documents = deduplicate_referenced_documents( + context.inline_rag_context.referenced_documents + + turn_summary.referenced_documents + ) + + # The start event was already emitted above; delegate the rest (re-yield, + # finalization, compacted-turn storage) to the shared generator. + async for event in generate_agent_response( + generator, + context, + responses_params, + turn_summary, + background_topic_summary_tasks=_background_topic_summary_tasks, + emit_start=False, + original_input=compacted_original_input, + root_span=root_span, + ): + yield event + finally: + if root_span is not None: + root_span.end() diff --git a/src/utils/agents/streaming.py b/src/utils/agents/streaming.py index e03b2326b..4170e55f0 100644 --- a/src/utils/agents/streaming.py +++ b/src/utils/agents/streaming.py @@ -12,6 +12,7 @@ from fastapi import HTTPException from ogx_client import APIConnectionError, APIStatusError +from opentelemetry import trace from pydantic_ai import Agent, AgentRunError, AgentRunResultEvent, ToolReturnPart from pydantic_ai.messages import ( AgentStreamEvent, @@ -60,6 +61,13 @@ process_native_tool_result, ) from utils.conversations import append_turn_items_to_conversation +from utils.otel_tracing import ( + SpanAttributes, + SpanEvents, + add_span_event, + anonymize_value, + set_span_attributes, +) from utils.pydantic_ai_helpers import build_agent from utils.query import ( build_multimodal_input, @@ -153,7 +161,7 @@ async def retrieve_agent_response_generator( raise HTTPException(**response.model_dump()) from exc -async def generate_agent_response( +async def generate_agent_response( # pylint: disable=too-many-statements generator: AsyncIterator[str], context: ResponseGeneratorContext, responses_params: ResponsesApiParams, @@ -161,6 +169,7 @@ async def generate_agent_response( background_topic_summary_tasks: list[asyncio.Task[None]], emit_start: bool = True, original_input: Optional[ResponseInput] = None, + root_span: Optional[trace.Span] = None, ) -> AsyncIterator[str]: """Wrap an agent SSE generator with cleanup logic. @@ -179,6 +188,8 @@ async def generate_agent_response( original_input: In compacted mode, the original user input before the explicit-input rewrite. Used to persist the completed turn with its structured input (preserving attachments); ``None`` otherwise. + root_span: OpenTelemetry root span for this request. + Yields: SSE-formatted strings from the wrapped generator. """ @@ -241,6 +252,8 @@ async def generate_agent_response( deregister_stream(context.request_id) if not stream_completed: + if root_span is not None: + root_span.end() return should_generate_topic_summary = ( @@ -269,6 +282,8 @@ async def generate_agent_response( ), media_type, ) + if root_span is not None: + root_span.end() return logger.info("Consuming tokens") consume_query_tokens( @@ -302,6 +317,40 @@ async def generate_agent_response( skip_userid_check=context.skip_userid_check, topic_summary=topic_summary, ) + + # Set final OTEL span attributes + if root_span is not None: + add_span_event(root_span, SpanEvents.TURN_PERSISTED) + if turn_summary.tool_calls: + tool_names = [tc.name for tc in turn_summary.tool_calls] + set_span_attributes( + root_span, + { + SpanAttributes.TOOL_CALLS_COUNT: len(tool_names), + SpanAttributes.TOOL_CALLS_NAMES: tool_names, + }, + ) + add_span_event( + root_span, + SpanEvents.TOOL_EXECUTION_COMPLETED, + {"tool.calls": ", ".join(tool_names)}, + ) + set_span_attributes( + root_span, + { + SpanAttributes.SESSION_ID: context.conversation_id, + SpanAttributes.LLM_USAGE_INPUT_TOKENS: ( + turn_summary.token_usage.input_tokens + ), + SpanAttributes.LLM_USAGE_OUTPUT_TOKENS: ( + turn_summary.token_usage.output_tokens + ), + SpanAttributes.OUTPUT: anonymize_value(turn_summary.llm_response), + }, + ) + add_span_event(root_span, SpanEvents.LLM_RESPONSE_COMPLETED) + root_span.end() + logger.info("Agent streaming complete") diff --git a/tests/unit/app/endpoints/test_streaming_query.py b/tests/unit/app/endpoints/test_streaming_query.py index 3cdc4855b..b19d40884 100644 --- a/tests/unit/app/endpoints/test_streaming_query.py +++ b/tests/unit/app/endpoints/test_streaming_query.py @@ -4,9 +4,13 @@ from typing import Any import pytest -from fastapi import Request +from fastapi import HTTPException, Request from fastapi.responses import StreamingResponse from ogx_client import AsyncOgxClient +from opentelemetry import trace +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) from pytest_mock import MockerFixture from app.endpoints.streaming_query import ( @@ -26,6 +30,7 @@ TurnSummary, ) from models.config import Action +from utils.otel_tracing import SpanAttributes, SpanEvents INTERRUPTED_INDICATOR = f"\n\n*{INTERRUPTED_RESPONSE_MESSAGE}*" @@ -569,3 +574,336 @@ async def mock_generate_agent_response( ) mock_client_holder.update_azure_token.assert_called_once() + + +async def _drain_response(response: StreamingResponse) -> None: + """Consume a StreamingResponse body to trigger the generator.""" + async for _ in response.body_iterator: + pass + + +class TestStreamingQueryOtelInstrumentation: + """Tests for OpenTelemetry instrumentation in the streaming query endpoint.""" + + def _setup_common_mocks( + self, + mocker: MockerFixture, + setup_configuration: AppConfig, + tracer: Any, + ) -> None: + """Set up common mocks for OTEL tests.""" + mocker.patch("app.endpoints.streaming_query.configuration", setup_configuration) + mocker.patch("app.endpoints.streaming_query.check_configuration_loaded") + mocker.patch("app.endpoints.streaming_query.check_tokens_available") + mocker.patch("app.endpoints.streaming_query.validate_model_provider_override") + mocker.patch( + "app.endpoints.streaming_query.build_rag_context", + new=mocker.AsyncMock(return_value=RAGContext()), + ) + mocker.patch( + "app.endpoints.streaming_query.check_mcp_auth", + new=mocker.AsyncMock(), + ) + + mock_client = mocker.AsyncMock(spec=AsyncOgxClient) + mock_client_holder = mocker.Mock() + mock_client_holder.get_client.return_value = mock_client + mocker.patch( + "app.endpoints.streaming_query.AsyncOgxClientHolder", + return_value=mock_client_holder, + ) + + mock_responses_params = mocker.Mock(spec=ResponsesApiParams) + mock_responses_params.model = "provider1/model1" + mock_responses_params.conversation = "conv_123" + mock_responses_params.tools = None + mock_responses_params.model_dump.return_value = { + "input": "test", + "model": "provider1/model1", + } + mocker.patch( + "app.endpoints.streaming_query.prepare_responses_params", + new=mocker.AsyncMock(return_value=mock_responses_params), + ) + mocker.patch( + "app.endpoints.streaming_query.run_shield_moderation", + new=mocker.AsyncMock(return_value=ShieldModerationPassed()), + ) + + mocker.patch("app.endpoints.streaming_query.AzureEntraIDManager") + mocker.patch( + "app.endpoints.streaming_query.extract_provider_and_model_from_model_id", + return_value=("provider1", "model1"), + ) + mocker.patch("app.endpoints.streaming_query.recording.record_llm_call") + + async def mock_generator() -> AsyncIterator[str]: + yield "data: test\n\n" + + mock_turn_summary = TurnSummary() + mocker.patch( + "app.endpoints.streaming_query.retrieve_agent_response_generator", + new=mocker.AsyncMock(return_value=(mock_generator(), mock_turn_summary)), + ) + + async def mock_generate_agent_response( + *_args: Any, **_kwargs: Any + ) -> AsyncIterator[str]: + async for item in mock_generator(): + yield item + if span := _kwargs.get("root_span"): + span.end() + + mocker.patch( + "app.endpoints.streaming_query.generate_agent_response", + side_effect=mock_generate_agent_response, + ) + mocker.patch( + "app.endpoints.streaming_query.normalize_conversation_id", + return_value="123", + ) + + mocker.patch("app.endpoints.streaming_query.tracer", tracer) + mocker.patch( + "app.endpoints.streaming_query.anonymize_value", + side_effect=lambda v: f"[anon:{v}]", + ) + + @pytest.mark.asyncio + async def test_creates_root_span( + self, + dummy_request: Request, # pylint: disable=redefined-outer-name + setup_configuration: AppConfig, + mocker: MockerFixture, + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that the handler creates a root span with the correct name.""" + tracer, exporter = otel + self._setup_common_mocks(mocker, setup_configuration, tracer) + + response = await streaming_query_endpoint_handler( + request=dummy_request, + query_request=QueryRequest( + query="test" + ), # pyright: ignore[reportCallIssue] + auth=MOCK_AUTH_STREAMING, + mcp_headers={}, + ) + await _drain_response(response) + + spans = exporter.get_finished_spans() + root_spans = [s for s in spans if s.name == "streaming_query.handle_request"] + assert len(root_spans) == 1 + + @pytest.mark.asyncio + async def test_sets_initial_span_attributes( + self, + dummy_request: Request, # pylint: disable=redefined-outer-name + setup_configuration: AppConfig, + mocker: MockerFixture, + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that initial span attributes are set for user ID, input, and attachments.""" + tracer, exporter = otel + self._setup_common_mocks(mocker, setup_configuration, tracer) + + response = await streaming_query_endpoint_handler( + request=dummy_request, + query_request=QueryRequest( + query="What is Kubernetes?" + ), # pyright: ignore[reportCallIssue] + auth=MOCK_AUTH_STREAMING, + mcp_headers={}, + ) + await _drain_response(response) + + spans = exporter.get_finished_spans() + root = [s for s in spans if s.name == "streaming_query.handle_request"][0] + assert root.attributes is not None + assert root.attributes[SpanAttributes.USER_ID] == ( + "[anon:00000001-0001-0001-0001-000000000001]" + ) + assert root.attributes[SpanAttributes.INPUT] == "[anon:What is Kubernetes?]" + assert root.attributes[SpanAttributes.REQUEST_ATTACHMENTS_COUNT] == 0 + + @pytest.mark.asyncio + async def test_sets_attachments_count_when_present( + self, + dummy_request: Request, # pylint: disable=redefined-outer-name + setup_configuration: AppConfig, + mocker: MockerFixture, + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that attachment count reflects actual attachments.""" + tracer, exporter = otel + self._setup_common_mocks(mocker, setup_configuration, tracer) + mocker.patch("app.endpoints.streaming_query.validate_attachments_metadata") + + query_request = QueryRequest( + query="test", + attachments=[ + Attachment( + attachment_type="log", + content_type="text/plain", + content="log1", + ), + Attachment( + attachment_type="log", + content_type="text/plain", + content="log2", + ), + ], + ) # pyright: ignore[reportCallIssue] + + response = await streaming_query_endpoint_handler( + request=dummy_request, + query_request=query_request, + auth=MOCK_AUTH_STREAMING, + mcp_headers={}, + ) + await _drain_response(response) + + spans = exporter.get_finished_spans() + root = [s for s in spans if s.name == "streaming_query.handle_request"][0] + assert root.attributes is not None + assert root.attributes[SpanAttributes.REQUEST_ATTACHMENTS_COUNT] == 2 + + @pytest.mark.asyncio + async def test_emits_validation_completed_event( + self, + dummy_request: Request, # pylint: disable=redefined-outer-name + setup_configuration: AppConfig, + mocker: MockerFixture, + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that VALIDATION_COMPLETED event is emitted after validation.""" + tracer, exporter = otel + self._setup_common_mocks(mocker, setup_configuration, tracer) + + response = await streaming_query_endpoint_handler( + request=dummy_request, + query_request=QueryRequest( + query="test" + ), # pyright: ignore[reportCallIssue] + auth=MOCK_AUTH_STREAMING, + mcp_headers={}, + ) + await _drain_response(response) + + spans = exporter.get_finished_spans() + root = [s for s in spans if s.name == "streaming_query.handle_request"][0] + event_names = [e.name for e in root.events] + assert SpanEvents.VALIDATION_COMPLETED in event_names + + @pytest.mark.asyncio + async def test_passes_root_span_to_generate_agent_response( + self, + dummy_request: Request, # pylint: disable=redefined-outer-name + setup_configuration: AppConfig, + mocker: MockerFixture, + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that root_span is forwarded to generate_agent_response.""" + tracer, _exporter = otel + self._setup_common_mocks(mocker, setup_configuration, tracer) + + mock_gen = mocker.patch( + "app.endpoints.streaming_query.generate_agent_response", + ) + + async def gen_side_effect(*_a: Any, **_kw: Any) -> AsyncIterator[str]: + yield "data: test\n\n" + + mock_gen.side_effect = gen_side_effect + + response = await streaming_query_endpoint_handler( + request=dummy_request, + query_request=QueryRequest( + query="test" + ), # pyright: ignore[reportCallIssue] + auth=MOCK_AUTH_STREAMING, + mcp_headers={}, + ) + await _drain_response(response) + + mock_gen.assert_called_once() + assert mock_gen.call_args.kwargs["root_span"] is not None + + @pytest.mark.asyncio + async def test_span_ended_on_exception( + self, + dummy_request: Request, # pylint: disable=redefined-outer-name + setup_configuration: AppConfig, + mocker: MockerFixture, + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that root span is ended when an exception occurs.""" + tracer, exporter = otel + self._setup_common_mocks(mocker, setup_configuration, tracer) + + mocker.patch( + "app.endpoints.streaming_query.check_configuration_loaded", + side_effect=HTTPException(status_code=500, detail="not loaded"), + ) + + with pytest.raises(HTTPException): + await streaming_query_endpoint_handler( + request=dummy_request, + query_request=QueryRequest( + query="test" + ), # pyright: ignore[reportCallIssue] + auth=MOCK_AUTH_STREAMING, + mcp_headers={}, + ) + + spans = exporter.get_finished_spans() + root_spans = [s for s in spans if s.name == "streaming_query.handle_request"] + assert len(root_spans) == 1 + + @pytest.mark.asyncio + async def test_child_spans_nested_under_root( + self, + dummy_request: Request, # pylint: disable=redefined-outer-name + setup_configuration: AppConfig, + mocker: MockerFixture, + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that child spans created during the request nest under root.""" + tracer, exporter = otel + self._setup_common_mocks(mocker, setup_configuration, tracer) + + async def mock_generate_with_child( + *_args: Any, **_kwargs: Any + ) -> AsyncIterator[str]: + root_span = _kwargs.get("root_span") + if root_span is not None: + parent_ctx = trace.set_span_in_context(root_span) + child = tracer.start_span("child.operation", context=parent_ctx) + child.end() + root_span.end() + yield "data: test\n\n" + + mocker.patch( + "app.endpoints.streaming_query.generate_agent_response", + side_effect=mock_generate_with_child, + ) + + response = await streaming_query_endpoint_handler( + request=dummy_request, + query_request=QueryRequest( + query="test" + ), # pyright: ignore[reportCallIssue] + auth=MOCK_AUTH_STREAMING, + mcp_headers={}, + ) + await _drain_response(response) + + spans = exporter.get_finished_spans() + root_spans = [s for s in spans if s.name == "streaming_query.handle_request"] + assert len(root_spans) == 1 + child_spans = [s for s in spans if s.parent is not None] + assert len(child_spans) >= 1 + for child in child_spans: + assert child.parent is not None # pyright narrowing + assert root_spans[0].context is not None # pyright narrowing + assert child.parent.span_id == root_spans[0].context.span_id diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index a353f6c67..e5e1da337 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -6,11 +6,16 @@ import os from collections.abc import Callable, Generator from pathlib import Path -from typing import Optional +from typing import Any, Optional import httpx import pytest from ogx_client import AsyncOgxClient +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) from pytest_mock import AsyncMockType, MockerFixture from configuration import AppConfig @@ -49,6 +54,18 @@ def otel_anonymization_secret() -> Generator[None, None, None]: os.environ["OTEL_ANONYMIZATION_SECRET"] = original_value +@pytest.fixture(name="otel") +def otel_fixture() -> Generator[tuple[Any, InMemorySpanExporter], None, None]: + """Provide an isolated tracer and exporter for OTEL tests.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = provider.get_tracer("unit-test-tracer") + yield tracer, exporter + exporter.clear() + provider.shutdown() + + @pytest.fixture(autouse=True) def reset_logging_state() -> Generator[None, None, None]: """Reset logging state before and after each test. diff --git a/tests/unit/utils/agents/test_streaming.py b/tests/unit/utils/agents/test_streaming.py index f452f8106..16ca0a899 100644 --- a/tests/unit/utils/agents/test_streaming.py +++ b/tests/unit/utils/agents/test_streaming.py @@ -11,6 +11,9 @@ import pytest from fastapi import HTTPException from ogx_client import APIStatusError +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) from pydantic_ai import AgentRunResultEvent from pydantic_ai.exceptions import AgentRunError from pydantic_ai.messages import ( @@ -52,7 +55,7 @@ from models.common.query import Attachment as QueryAttachment from models.common.responses.contexts import ResponseGeneratorContext from models.common.responses.responses_api_params import ResponsesApiParams -from models.common.turn_summary import RAGContext, TurnSummary +from models.common.turn_summary import RAGContext, ToolCallSummary, TurnSummary from utils.agents.query import AgentFinishReason from utils.agents.streaming import ( DEFAULT_REFUSAL_RESPONSE, @@ -62,6 +65,7 @@ retrieve_agent_response_generator, serialize_event, ) +from utils.otel_tracing import SpanAttributes, SpanEvents from utils.token_counter import TokenCounter INTERRUPTED_INDICATOR = f"\n\n*{INTERRUPTED_RESPONSE_MESSAGE}*" @@ -811,6 +815,336 @@ async def inner() -> AsyncIterator[str]: persist_mock.assert_not_awaited() +class TestGenerateAgentResponseOtel: + """Tests for OTEL instrumentation in generate_agent_response.""" + + @pytest.mark.asyncio + async def test_sets_final_span_attributes_on_success( + self, + mocker: MockerFixture, + make_generator_context: Callable[..., ResponseGeneratorContext], + responses_params: ResponsesApiParams, + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that final OTEL attributes are set after successful stream.""" + tracer, exporter = otel + context = make_generator_context() + turn_summary = TurnSummary() + turn_summary.token_usage = TokenCounter(input_tokens=10, output_tokens=5) + turn_summary.llm_response = "The answer is 42" + background_tasks: list[asyncio.Task[None]] = [] + root_span = tracer.start_span("streaming_query.handle_request") + + async def inner() -> AsyncIterator[str]: + yield serialize_event( + TokenStreamPayload.create(chunk_id=0, token="Hi"), + MEDIA_TYPE_JSON, + ) + + mocker.patch("utils.agents.streaming.consume_query_tokens") + mocker.patch( + "utils.agents.streaming.get_available_quotas", + return_value={"daily": 100}, + ) + mocker.patch( + "utils.agents.streaming.maybe_get_topic_summary", + new=mocker.AsyncMock(return_value=None), + ) + mocker.patch("utils.agents.streaming.store_query_results") + mock_config = mocker.Mock() + mock_config.quota_limiters = [] + mocker.patch("utils.agents.streaming.configuration", mock_config) + + mocker.patch( + "utils.agents.streaming.anonymize_value", + side_effect=lambda v: f"[anon:{v}]", + ) + + [ + event + async for event in generate_agent_response( + inner(), + context, + responses_params, + turn_summary, + background_tasks, + root_span=root_span, + ) + ] + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.attributes is not None + assert span.attributes[SpanAttributes.SESSION_ID] == context.conversation_id + assert span.attributes[SpanAttributes.LLM_USAGE_INPUT_TOKENS] == 10 + assert span.attributes[SpanAttributes.LLM_USAGE_OUTPUT_TOKENS] == 5 + assert span.attributes[SpanAttributes.OUTPUT] == "[anon:The answer is 42]" + event_names = [e.name for e in span.events] + assert SpanEvents.TURN_PERSISTED in event_names + assert SpanEvents.LLM_RESPONSE_COMPLETED in event_names + + @pytest.mark.asyncio + async def test_sets_tool_call_span_attributes( + self, + mocker: MockerFixture, + make_generator_context: Callable[..., ResponseGeneratorContext], + responses_params: ResponsesApiParams, + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that tool call OTEL attributes are emitted on the root span.""" + tracer, exporter = otel + context = make_generator_context() + turn_summary = TurnSummary() + turn_summary.token_usage = TokenCounter(input_tokens=10, output_tokens=5) + turn_summary.llm_response = "Result" + turn_summary.tool_calls = [ + ToolCallSummary(id="tc-1", name="web_search", type="web_search_call"), + ToolCallSummary(id="tc-2", name="file_search", type="file_search_call"), + ] + background_tasks: list[asyncio.Task[None]] = [] + root_span = tracer.start_span("streaming_query.handle_request") + + async def inner() -> AsyncIterator[str]: + yield serialize_event( + TokenStreamPayload.create(chunk_id=0, token="Hi"), + MEDIA_TYPE_JSON, + ) + + mocker.patch("utils.agents.streaming.consume_query_tokens") + mocker.patch( + "utils.agents.streaming.get_available_quotas", + return_value={"daily": 100}, + ) + mocker.patch( + "utils.agents.streaming.maybe_get_topic_summary", + new=mocker.AsyncMock(return_value=None), + ) + mocker.patch("utils.agents.streaming.store_query_results") + mock_config = mocker.Mock() + mock_config.quota_limiters = [] + mocker.patch("utils.agents.streaming.configuration", mock_config) + mocker.patch( + "utils.agents.streaming.anonymize_value", + side_effect=lambda v: f"[anon:{v}]", + ) + + [ + event + async for event in generate_agent_response( + inner(), + context, + responses_params, + turn_summary, + background_tasks, + root_span=root_span, + ) + ] + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.attributes is not None + assert span.attributes[SpanAttributes.TOOL_CALLS_COUNT] == 2 + assert span.attributes[SpanAttributes.TOOL_CALLS_NAMES] == ( + "web_search", + "file_search", + ) + event_names = [e.name for e in span.events] + assert SpanEvents.TOOL_EXECUTION_COMPLETED in event_names + tool_event = next( + e for e in span.events if e.name == SpanEvents.TOOL_EXECUTION_COMPLETED + ) + assert tool_event.attributes is not None + assert tool_event.attributes["tool.calls"] == "web_search, file_search" + + @pytest.mark.asyncio + async def test_span_ended_on_stream_error( + self, + mocker: MockerFixture, + make_generator_context: Callable[..., ResponseGeneratorContext], + responses_params: ResponsesApiParams, + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that span is ended when streaming fails with an error.""" + tracer, exporter = otel + context = make_generator_context() + root_span = tracer.start_span("streaming_query.handle_request") + + async def inner() -> AsyncIterator[str]: + yield serialize_event( + TokenStreamPayload.create(chunk_id=0, token="partial"), + MEDIA_TYPE_JSON, + ) + raise AgentRunError("inference failure") + + mocker.patch( + "utils.agents.streaming.register_interrupt_callback", + return_value=[False], + ) + + [ + event + async for event in generate_agent_response( + inner(), + context, + responses_params, + TurnSummary(), + [], + root_span=root_span, + ) + ] + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == "streaming_query.handle_request" + + @pytest.mark.asyncio + async def test_span_ended_on_topic_summary_error( + self, + mocker: MockerFixture, + make_generator_context: Callable[..., ResponseGeneratorContext], + responses_params: ResponsesApiParams, + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that span is ended when topic summary generation fails.""" + tracer, exporter = otel + context = make_generator_context( + generate_topic_summary=True, conversation_id_in_request=None + ) + turn_summary = TurnSummary() + turn_summary.token_usage = TokenCounter(input_tokens=3, output_tokens=7) + root_span = tracer.start_span("streaming_query.handle_request") + + async def inner() -> AsyncIterator[str]: + yield serialize_event( + TokenStreamPayload.create(chunk_id=0, token="ok"), + MEDIA_TYPE_JSON, + ) + + mocker.patch("utils.agents.streaming.consume_query_tokens") + mocker.patch( + "utils.agents.streaming.get_available_quotas", + return_value={}, + ) + mocker.patch( + "utils.agents.streaming.maybe_get_topic_summary", + new=mocker.AsyncMock( + side_effect=HTTPException(status_code=500, detail="boom") + ), + ) + mock_config = mocker.Mock() + mock_config.quota_limiters = [] + mocker.patch("utils.agents.streaming.configuration", mock_config) + + [ + event + async for event in generate_agent_response( + inner(), + context, + responses_params, + turn_summary, + [], + root_span=root_span, + ) + ] + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + + @pytest.mark.asyncio + async def test_no_spans_finished_when_root_span_is_none( + self, + mocker: MockerFixture, + make_generator_context: Callable[..., ResponseGeneratorContext], + responses_params: ResponsesApiParams, + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that no spans are finished when root_span is None.""" + _tracer, exporter = otel + context = make_generator_context() + turn_summary = TurnSummary() + turn_summary.token_usage = TokenCounter(input_tokens=3, output_tokens=7) + + async def inner() -> AsyncIterator[str]: + yield serialize_event( + TokenStreamPayload.create(chunk_id=0, token="Hi"), + MEDIA_TYPE_JSON, + ) + + mocker.patch("utils.agents.streaming.consume_query_tokens") + mocker.patch( + "utils.agents.streaming.get_available_quotas", + return_value={"daily": 100}, + ) + mocker.patch( + "utils.agents.streaming.maybe_get_topic_summary", + new=mocker.AsyncMock(return_value=None), + ) + mocker.patch("utils.agents.streaming.store_query_results") + mock_config = mocker.Mock() + mock_config.quota_limiters = [] + mocker.patch("utils.agents.streaming.configuration", mock_config) + + [ + event + async for event in generate_agent_response( + inner(), + context, + responses_params, + turn_summary, + [], + root_span=None, + ) + ] + + assert len(exporter.get_finished_spans()) == 0 + + @pytest.mark.asyncio + async def test_span_ended_on_cancelled_error( + self, + mocker: MockerFixture, + make_generator_context: Callable[..., ResponseGeneratorContext], + responses_params: ResponsesApiParams, + otel: tuple[Any, InMemorySpanExporter], + ) -> None: + """Test that span is ended when stream is cancelled/interrupted.""" + tracer, exporter = otel + context = make_generator_context() + root_span = tracer.start_span("streaming_query.handle_request") + + async def inner() -> AsyncIterator[str]: + yield serialize_event( + TokenStreamPayload.create(chunk_id=0, token="partial"), + MEDIA_TYPE_JSON, + ) + raise asyncio.CancelledError() + + mocker.patch( + "utils.agents.streaming.persist_interrupted_turn", + new=mocker.AsyncMock(), + ) + mocker.patch( + "utils.agents.streaming.register_interrupt_callback", + return_value=[False], + ) + + [ + event + async for event in generate_agent_response( + inner(), + context, + responses_params, + TurnSummary(), + [], + root_span=root_span, + ) + ] + + spans = exporter.get_finished_spans() + assert len(spans) == 1 + + class TestAgentResponseGenerator: """Tests for agent_response_generator."""