diff --git a/CHANGELOG.md b/CHANGELOG.md index c83a5dc..df7bb3e 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/backend/chat/clients/async_to_sync.py b/src/backend/chat/clients/async_to_sync.py index 9436dbc..1684234 100644 --- a/src/backend/chat/clients/async_to_sync.py +++ b/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: diff --git a/src/backend/chat/clients/exceptions.py b/src/backend/chat/clients/exceptions.py index 7dbb6b4..374d787 100644 --- a/src/backend/chat/clients/exceptions.py +++ b/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) diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py index 16153fc..065e3fe 100644 --- a/src/backend/chat/clients/pydantic_ai.py +++ b/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(), diff --git a/src/backend/chat/tests/clients/pydantic_ai/test_stream_methods.py b/src/backend/chat/tests/clients/pydantic_ai/test_stream_methods.py index e4e81e0..465ea5e 100644 --- a/src/backend/chat/tests/clients/pydantic_ai/test_stream_methods.py +++ b/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): diff --git a/src/backend/chat/tests/views/chat/conversations/test_stop_steaming.py b/src/backend/chat/tests/views/chat/conversations/test_stop_steaming.py new file mode 100644 index 0000000..b0d75b9 --- /dev/null +++ b/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 diff --git a/src/backend/chat/views.py b/src/backend/chat/views.py index eca4786..b77a2d1 100644 --- a/src/backend/chat/views.py +++ b/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) diff --git a/src/backend/conversations/settings.py b/src/backend/conversations/settings.py index 03f51d4..b7991cc 100755 --- a/src/backend/conversations/settings.py +++ b/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",