Merge pull request #1918 from asimurka/relocate_streaming_utils

LCORE-2311: Relocated stream interruption utilities

Author: tisnik

src/app/endpoints/streaming_query.py

@@ -1,7 +1,5 @@
 """Streaming query handler using Responses API."""
 
-# pylint: disable=too-many-lines
-
 import asyncio
 import datetime
 from collections.abc import AsyncIterator
@@ -10,7 +8,6 @@
 from fastapi import APIRouter, Depends, HTTPException, Request
 from fastapi.responses import StreamingResponse
 from llama_stack_api import (
-    OpenAIResponseMessage,
     OpenAIResponseObject,
     OpenAIResponseObjectStream,
 )
@@ -56,7 +53,6 @@
     MEDIA_TYPE_EVENT_STREAM,
     MEDIA_TYPE_JSON,
     MEDIA_TYPE_TEXT,
-    TOPIC_SUMMARY_INTERRUPT_TIMEOUT_SECONDS,
 )
 from log import get_logger
 from metrics import recording
@@ -100,7 +96,6 @@
     is_context_length_error,
     prepare_input,
     store_query_results,
-    update_conversation_topic_summary,
     validate_attachments_metadata,
     validate_model_provider_override,
 )
@@ -118,11 +113,14 @@
     prepare_responses_params,
 )
 from utils.shields import (
-    append_turn_to_conversation,
     run_shield_moderation,
     validate_shield_ids_override,
 )
-from utils.stream_interrupts import get_stream_interrupt_registry
+from utils.stream_interrupts import (
+    deregister_stream,
+    persist_interrupted_turn,
+    register_interrupt_callback,
+)
 from utils.streaming_sse import (
     http_exception_stream_event,
     shield_violation_generator,
@@ -430,43 +428,6 @@ async def retrieve_response_generator(
         raise HTTPException(**error_response.model_dump()) from e
 
 
-async def _background_update_topic_summary(
-    context: ResponseGeneratorContext,
-    model: str,
-) -> None:
-    """Generate topic summary and update DB/cache in the background.
-
-    Runs as a fire-and-forget task after an interrupted turn is persisted.
-    All errors are caught and logged.
-    """
-    try:
-        topic_summary = await asyncio.wait_for(
-            get_topic_summary(
-                context.query_request.query,
-                context.client,
-                model,
-            ),
-            timeout=TOPIC_SUMMARY_INTERRUPT_TIMEOUT_SECONDS,
-        )
-        if topic_summary:
-            update_conversation_topic_summary(
-                context.conversation_id,
-                topic_summary,
-                user_id=context.user_id,
-                skip_userid_check=context.skip_userid_check,
-            )
-    except asyncio.TimeoutError:
-        logger.warning(
-            "Topic summary timed out for interrupted turn, request %s",
-            context.request_id,
-        )
-    except Exception:  # pylint: disable=broad-except
-        logger.exception(
-            "Failed to generate topic summary for interrupted turn, request %s",
-            context.request_id,
-        )
-
-
 async def shutdown_background_topic_summary_tasks() -> None:
     """Cancel and await outstanding background topic summary tasks on shutdown.
 
@@ -485,148 +446,6 @@ async def shutdown_background_topic_summary_tasks() -> None:
     await asyncio.gather(*tasks, return_exceptions=True)
 
 
-async def _persist_interrupted_turn(
-    context: ResponseGeneratorContext,
-    responses_params: ResponsesApiParams,
-    turn_summary: TurnSummary,
-    original_input: Optional[ResponseInput] = None,
-) -> None:
-    """Persist the user query and an interrupted response into the conversation.
-
-    Called when a streaming request is cancelled so the exchange is not lost.
-    Persists immediately with topic_summary=None so the conversation exists
-    when the client fetches. Topic summary is generated in a background task
-    and updated when ready.
-
-    Parameters:
-    ----------
-        context: The response generator context.
-        responses_params: The Responses API parameters.
-        turn_summary: TurnSummary with llm_response already set to the
-            interrupted message.
-        original_input: In compacted mode, the original user input before the
-            explicit-input rewrite. When set, the turn is persisted against it
-            (the ``conversation`` parameter was dropped, and
-            ``responses_params.input`` is the explicit rewrite); ``None``
-            otherwise (LCORE-1572).
-    """
-    try:
-        if original_input is not None:
-            await append_turn_items_to_conversation(
-                context.client,
-                responses_params.conversation,
-                original_input,
-                [
-                    OpenAIResponseMessage(
-                        role="assistant", content=INTERRUPTED_RESPONSE_MESSAGE
-                    )
-                ],
-            )
-        else:
-            await append_turn_to_conversation(
-                context.client,
-                responses_params.conversation,
-                cast(str, responses_params.input),
-                INTERRUPTED_RESPONSE_MESSAGE,
-            )
-    except Exception:  # pylint: disable=broad-except
-        logger.exception(
-            "Failed to append interrupted turn to conversation for request %s",
-            context.request_id,
-        )
-
-    try:
-        completed_at = datetime.datetime.now(datetime.UTC).strftime(
-            "%Y-%m-%dT%H:%M:%SZ"
-        )
-        store_query_results(
-            user_id=context.user_id,
-            conversation_id=context.conversation_id,
-            model=responses_params.model,
-            completed_at=completed_at,
-            started_at=context.started_at,
-            summary=turn_summary,
-            query=context.query_request.query,
-            skip_userid_check=context.skip_userid_check,
-            topic_summary=None,
-        )
-
-        if (
-            not context.query_request.conversation_id
-            and context.query_request.generate_topic_summary
-        ):
-            task = asyncio.create_task(
-                _background_update_topic_summary(
-                    context=context,
-                    model=responses_params.model,
-                )
-            )
-            _background_topic_summary_tasks.append(task)
-            task.add_done_callback(_background_topic_summary_tasks.remove)
-    except Exception:  # pylint: disable=broad-except
-        logger.exception(
-            "Failed to store interrupted query results for request %s",
-            context.request_id,
-        )
-
-
-def _register_interrupt_callback(
-    context: ResponseGeneratorContext,
-    responses_params: ResponsesApiParams,
-    turn_summary: TurnSummary,
-    original_input: Optional[ResponseInput] = None,
-) -> list[bool]:
-    """Build an interrupt callback and register the stream for cancellation.
-
-    The callback is invoked by ``cancel_stream`` when the client
-    interrupts, so persistence runs regardless of where the
-    ``CancelledError`` is raised in the ASGI stack.
-
-    A mutable one-element list is used as a shared guard so the
-    callback and the in-generator ``CancelledError`` handler never
-    both persist the same turn.
-
-    Parameters:
-    ----------
-        context: The response generator context.
-        responses_params: The Responses API parameters.
-        turn_summary: TurnSummary populated during streaming.
-
-    Returns:
-    -------
-        A mutable list ``[False]`` used as a persist-done guard; the
-        caller should check ``guard[0]`` before persisting and set
-        it to ``True`` afterwards.
-    """
-    guard: list[bool] = [False]
-
-    async def _on_interrupt() -> None:
-        if guard[0]:
-            return
-        guard[0] = True
-        turn_summary.llm_response = INTERRUPTED_RESPONSE_MESSAGE
-        await _persist_interrupted_turn(
-            context, responses_params, turn_summary, original_input
-        )
-
-    current_task = asyncio.current_task()
-    if current_task is not None:
-        get_stream_interrupt_registry().register_stream(
-            request_id=context.request_id,
-            user_id=context.user_id,
-            task=current_task,
-            on_interrupt=_on_interrupt,
-        )
-    else:
-        logger.warning(
-            "No current asyncio task for request %s; "
-            "stream interruption will not be available",
-            context.request_id,
-        )
-
-    return guard
-
-
 async def generate_response_with_compaction(
     context: ResponseGeneratorContext,
     responses_params: ResponsesApiParams,
@@ -759,8 +578,12 @@ async def generate_response(  # pylint: disable=too-many-arguments,too-many-posi
     Yields:
         SSE-formatted strings from the wrapped generator
     """
-    persist_guard = _register_interrupt_callback(
-        context, responses_params, turn_summary, original_input
+    persist_guard = register_interrupt_callback(
+        context,
+        responses_params,
+        turn_summary,
+        _background_topic_summary_tasks,
+        original_input,
     )
 
     stream_completed = False
@@ -802,12 +625,16 @@ async def generate_response(  # pylint: disable=too-many-arguments,too-many-posi
         if not persist_guard[0]:
             persist_guard[0] = True
             turn_summary.llm_response = INTERRUPTED_RESPONSE_MESSAGE
-            await _persist_interrupted_turn(
-                context, responses_params, turn_summary, original_input
+            await persist_interrupted_turn(
+                context,
+                responses_params,
+                turn_summary,
+                _background_topic_summary_tasks,
+                original_input,
             )
         yield stream_interrupted_event(context.request_id)
     finally:
-        get_stream_interrupt_registry().deregister_stream(context.request_id)
+        deregister_stream(context.request_id)
 
     if not stream_completed:
         return