Merge pull request #2415 from anik120/otel-for-streamingq

LCORE-2981: Add OpenTelemetry instrumentation for POST /v1/streaming_query endpoint

Author: tisnik

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()

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,14 +161,15 @@ 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,
     turn_summary: TurnSummary,
     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")