✨(backend) allow use to stop conversation streaming

This introduces a new endpoint to allow user to stop the message
stream.

We currently use uWSGI workers, which will not automatically stop the streaming
when the request is cancelled by the client. Therefore, we need to
explicitly stop the streaming by calling the /stop-steaming endpoint.
When (if) we switch to Gunicorn with Uvicorn workers, this will not be necessary
as the Uvicorn workers will automatically stop the streaming when the request
is cancelled. BUT, we will then need to handle the streaming cancellation when
the user is simply offline and still waiting for a response and the conversation
to be updated with the result. So this endpoint will still be useful to be able
to detect the cancellation of the streaming versus the user being offline.

CHANGELOG.md

@@ -19,6 +19,7 @@ and this project adheres to
 - 🎉(conversations) bootstrap backend & frontend #1
 - ✨(web-search) add RAG capability to do web search #7
 - ✨(chat) add document RAG on document uploaded by user #8
+- ✨(backend) allow use to stop conversation streaming #14
 
 
 [unreleased]: https://github.com/numerique-gouv/conversations/compare/HEAD...main

src/backend/chat/clients/async_to_sync.py

@@ -9,6 +9,8 @@ import asyncio
 import queue
 import threading
 
+from chat.clients.exceptions import StreamCancelException
+
 
 def convert_async_generator_to_sync(async_gen):
     """Convert an async generator to a sync generator."""
@@ -20,6 +22,9 @@ def convert_async_generator_to_sync(async_gen):
         try:
             async for async_item in async_gen:
                 q.put(async_item)
+        except StreamCancelException:
+            # Handle cancellation gracefully, do not put anything in the queue
+            q.put(sentinel)
         except Exception as exc:  # pylint: disable=broad-except #noqa: BLE001
             q.put((exc_sentinel, exc))
         finally:

src/backend/chat/clients/exceptions.py

@@ -7,3 +7,11 @@ class WebSearchEmptyException(Exception):
     def __init__(self, message="Web search returned no results."):
         self.message = message
         super().__init__(self.message)
+
+
+class StreamCancelException(Exception):
+    """Exception raised when a streaming operation is cancelled."""
+
+    def __init__(self, message="Streaming operation was cancelled."):
+        self.message = message
+        super().__init__(self.message)

src/backend/chat/clients/pydantic_ai.py

@@ -9,11 +9,13 @@ changes are needed in views.py or tests.
 import dataclasses
 import json
 import logging
+import time
 import uuid
 from contextlib import AsyncExitStack
-from typing import Dict, List, Tuple
+from typing import Dict, List, Optional, Tuple
 
 from django.conf import settings
+from django.core.cache import cache
 from django.core.exceptions import ImproperlyConfigured
 from django.utils.module_loading import import_string
 from django.utils.translation import gettext_lazy as _
@@ -51,7 +53,7 @@ from chat.ai_sdk_types import (
     UIMessage,
 )
 from chat.clients.async_to_sync import convert_async_generator_to_sync
-from chat.clients.exceptions import WebSearchEmptyException
+from chat.clients.exceptions import StreamCancelException, WebSearchEmptyException
 from chat.clients.pydantic_ui_message_converter import (
     model_message_to_ui_message,
     ui_message_to_user_content,
@@ -99,6 +101,11 @@ class AIAgentService:
 
     def __init__(self, conversation):
         self.conversation = conversation
+        self._last_stop_check = 0
+
+    @property
+    def _stop_cache_key(self):
+        return f"streaming:stop:{self.conversation.pk}"
 
     # --------------------------------------------------------------------- #
     # Public streaming API (unchanged signatures)
@@ -112,21 +119,59 @@ class AIAgentService:
         """Return Vercel-AI-SDK formatted events."""
         return convert_async_generator_to_sync(self.stream_data_async(messages, force_web_search))
 
+    def stop_streaming(self):
+        """
+        Stop the current streaming operation.
+
+        This method is a placeholder for stopping the streaming operation.
+        """
+        logger.info("Stopping streaming for conversation %s", self.conversation.id)
+        cache.set(self._stop_cache_key, "1", timeout=30 * 60)  # 30 minutes timeout
+
     # --------------------------------------------------------------------- #
     # Async internals
     # --------------------------------------------------------------------- #
 
     async def stream_text_async(self, messages: List[UIMessage], force_web_search: bool = False):
         """Return only the assistant text deltas (legacy text mode)."""
+        await self._clean()
         async for delta in self._run_agent(messages, force_web_search):
             if delta["type"] == "0":
                 yield delta["payload"]
 
     async def stream_data_async(self, messages: List[UIMessage], force_web_search: bool = False):
         """Return Vercel-AI-SDK formatted events."""
+        await self._clean()
         async for delta in self._run_agent(messages, force_web_search):
             yield f"{delta['type']}:{json.dumps(delta['payload'])}\n"
 
+    async def _agent_stop_streaming(self, force_cache_check: Optional[bool] = False) -> None:
+        """Check if the agent should stop streaming."""
+        now = time.time()  # Current time in seconds since epoch
+
+        # Check if we should skip the cache check to avoid frequent checks
+        # This is useful to avoid unnecessary cache checks during streaming
+        # Check every 2 seconds
+        if not force_cache_check and now - self._last_stop_check < 2:
+            return
+        self._last_stop_check = now
+
+        if await cache.aget(self._stop_cache_key):
+            logger.info("Streaming stopped by cache key for conversation %s", self.conversation.id)
+            await cache.adelete(self._stop_cache_key)
+            raise StreamCancelException()
+        return
+
+    async def _clean(self):
+        """
+        Clean up the agent service.
+
+        This method is called when the agent service is no longer needed.
+        It can be used to release resources or perform any necessary cleanup.
+        """
+        self._last_stop_check = 0
+        await cache.adelete(self._stop_cache_key)
+
     # --------------------------------------------------------------------- #
     # Core agent runner
     # --------------------------------------------------------------------- #
@@ -312,7 +357,7 @@ class AIAgentService:
         self,
         messages: List[UIMessage],
         force_web_search: bool = False,
-    ):  # pylint: disable=too-many-branches,too-many-statements, too-many-locals
+    ):  # pylint: disable=too-many-branches,too-many-statements, too-many-locals, too-many-return-statements
         """Run the Pydantic AI agent and stream events."""
         if messages[-1].role != "user":
             return
@@ -341,6 +386,8 @@ class AIAgentService:
             yield {"type": "3", "payload": "attachment_summary_not_supported"}
             return
 
+        await self._agent_stop_streaming(force_cache_check=True)
+
         # If user uploaded documents and did not enforce a web search, we disable the intent
         if input_documents and not force_web_search:
             user_intent.web_search = False
@@ -377,6 +424,8 @@ class AIAgentService:
                 },
             }
 
+        await self._agent_stop_streaming(force_cache_check=True)
+
         # Prepare the prompt for the agent
         try:
             _new_prompt, _ui_sources = self.perform_rag(
@@ -415,6 +464,7 @@ class AIAgentService:
                 message_history=history,
             ) as run:
                 async for node in run:
+                    await self._agent_stop_streaming()
                     if Agent.is_user_prompt_node(node):
                         # A user prompt node => The user has provided input
                         pass
@@ -423,6 +473,7 @@ class AIAgentService:
                         # A model request node => We can stream tokens from the model's request
                         async with node.stream(run.ctx) as request_stream:
                             async for event in request_stream:
+                                await self._agent_stop_streaming()
                                 logger.debug("Received request_stream event: %s", type(event))
                                 if isinstance(event, PartStartEvent):
                                     logger.debug("PartStartEvent: %s", dataclasses.asdict(event))
@@ -464,6 +515,7 @@ class AIAgentService:
                         # potentially calls a tool
                         async with node.stream(run.ctx) as handle_stream:
                             async for event in handle_stream:
+                                await self._agent_stop_streaming()
                                 logger.debug(
                                     "Received request_stream event: %s, %s",
                                     type(event),
@@ -518,6 +570,8 @@ class AIAgentService:
                 usage["promptTokens"] = final_usage.request_tokens
                 usage["completionTokens"] = final_usage.response_tokens
 
+        await self._agent_stop_streaming(force_cache_check=True)
+
         # Persist conversation
         await sync_to_async(self._update_conversation)(
             final_output=run.result.new_messages(),

src/backend/chat/tests/clients/pydantic_ai/test_stream_methods.py

@@ -57,7 +57,7 @@ def test_stream_data_delegates_to_async(mock_convert, ui_messages):
 @pytest.mark.asyncio
 async def test_stream_text_async_filters_text_deltas(ui_messages):
     """Test stream_text_async only yields text deltas."""
-    conversation = sync_to_async(ChatConversationFactory)()
+    conversation = await sync_to_async(ChatConversationFactory)()
     service = AIAgentService(conversation)
 
     # Mock _run_agent to return various delta types
@@ -78,7 +78,7 @@ async def test_stream_text_async_filters_text_deltas(ui_messages):
 @pytest.mark.asyncio
 async def test_stream_data_async_formats_as_sdk_events(ui_messages):
     """Test stream_data_async formats events correctly."""
-    conversation = sync_to_async(ChatConversationFactory)()
+    conversation = await sync_to_async(ChatConversationFactory)()
     service = AIAgentService(conversation)
 
     async def mock_run_agent(*args, **kwargs):

src/backend/chat/tests/views/chat/conversations/test_stop_steaming.py

@@ -0,0 +1,50 @@
+"""Test the post_stop_steaming view."""
+
+from unittest.mock import patch
+
+import pytest
+from rest_framework import status
+
+from core.factories import UserFactory
+
+from chat import factories
+from chat.factories import ChatConversationFactory
+
+pytestmark = pytest.mark.django_db()
+
+
+def test_stop_streaming(api_client):
+    """Test that the stop_streaming method is called when the endpoint is called."""
+    chat_conversation = factories.ChatConversationFactory()
+    api_client.force_login(chat_conversation.owner)
+    url = f"/api/v1.0/chats/{chat_conversation.pk}/stop-streaming/"
+
+    with patch("chat.clients.pydantic_ai.AIAgentService.stop_streaming") as mock_stop_streaming:
+        response = api_client.post(url)
+        mock_stop_streaming.assert_called_once_with()
+
+    assert response.status_code == status.HTTP_200_OK
+    assert response.json() == {"status": "OK"}
+
+
+def test_stop_streaming_unauthenticated(api_client):
+    """Test that unauthenticated users cannot call the endpoint."""
+    chat_conversation = ChatConversationFactory()
+    url = f"/api/v1.0/chats/{chat_conversation.pk}/stop-streaming/"
+
+    response = api_client.post(url)
+
+    assert response.status_code == status.HTTP_401_UNAUTHORIZED
+
+
+def test_stop_streaming_for_conversation_of_another_user(
+    api_client,
+):
+    """Test that a user cannot stop streaming for a conversation of another user."""
+    chat_conversation = factories.ChatConversationFactory()
+    api_client.force_login(UserFactory())
+    url = f"/api/v1.0/chats/{chat_conversation.pk}/stop-streaming/"
+
+    response = api_client.post(url)
+
+    assert response.status_code == status.HTTP_404_NOT_FOUND

src/backend/chat/views.py

@@ -126,3 +126,39 @@ class ChatViewSet(  # pylint: disable=too-many-ancestors
             },
         )
         return response
+
+    @decorators.action(
+        methods=["post"],
+        detail=True,
+        url_path="stop-streaming",
+        url_name="stop-streaming",
+    )
+    def post_stop_steaming(self, request, pk):  # pylint: disable=unused-argument
+        """Handle POST requests to stop streaming the chat conversation.
+
+        This action will put a poison pill in the redis cache to stop any ongoing streaming.
+        It is used to stop the streaming when the user decides to cancel the chat.
+
+        Note:
+            We currently use uWSGI workers, which will not automatically stop the streaming
+            when the request is cancelled by the client. Therefore, we need to
+            explicitly stop the streaming by calling this endpoint.
+            When (if) we switch to Gunicorn with Uvicorn workers, this will not be necessary
+            as the Uvicorn workers will automatically stop the streaming when the request
+            is cancelled. BUT, we will then need to handle the streaming cancellation when the
+            user is simply offline and still waiting for a response and the conversation to
+            be updated with the result. So this endpoint will still be useful to be able to
+            detect the cancellation of the streaming versus the user being offline.
+
+        Args:
+            request: The HTTP request object.
+            pk: The primary key of the chat conversation.
+
+        Returns:
+            Response: A response indicating that the streaming has been stopped.
+        """
+        conversation = self.get_object()
+
+        AIAgentService(conversation=conversation).stop_streaming()
+
+        return Response({"status": "OK"}, status=status.HTTP_200_OK)

src/backend/conversations/settings.py

@@ -806,12 +806,8 @@ class Development(Base):
     SESSION_COOKIE_NAME = "conversations_sessionid"
 
     USE_SWAGGER = True
-    SESSION_CACHE_ALIAS = "session"
     CACHES = {
         "default": {
-            "BACKEND": "django.core.cache.backends.dummy.DummyCache",
-        },
-        "session": {
             "BACKEND": "django_redis.cache.RedisCache",
             "LOCATION": values.Value(
                 "redis://redis:6379/2",