✨(back) manage streaming with the ai service

We want to handle both streaming or not when interacting with the AI
backend service.

Author: lunika

src/backend/core/api/viewsets.py

@@ -1361,8 +1361,20 @@ def ai_proxy(self, request, *args, **kwargs):
         serializer = serializers.AIProxySerializer(data=request.data)
         serializer.is_valid(raise_exception=True)
 
-        response = AIService().proxy(request.data)
-        return drf.response.Response(response, status=drf.status.HTTP_200_OK)
+        ai_service = AIService()
+
+        if settings.AI_STREAM:
+            return StreamingHttpResponse(
+                ai_service.stream(request.data),
+                content_type="text/event-stream",
+                status=drf.status.HTTP_200_OK,
+            )
+    
+        ai_response = ai_service.proxy(request.data)
+        return drf.response.Response(
+            ai_response.model_dump(),
+            status=drf.status.HTTP_200_OK,
+        )
 
     @drf.decorators.action(
         detail=True,
@@ -1821,6 +1833,7 @@ def get(self, request):
             "AI_BOT",
             "AI_FEATURE_ENABLED",
             "AI_MODEL",
+            "AI_STREAM",
             "COLLABORATION_WS_URL",
             "COLLABORATION_WS_NOT_CONNECTED_READY_ONLY",
             "CRISP_WEBSITE_ID",

src/backend/core/services/ai_services.py

@@ -1,6 +1,7 @@
 """AI services."""
 
 import logging
+from typing import Generator
 
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
@@ -23,9 +24,15 @@ def __init__(self):
             raise ImproperlyConfigured("AI configuration not set")
         self.client = OpenAI(base_url=settings.AI_BASE_URL, api_key=settings.AI_API_KEY)
 
-    def proxy(self, data: dict) -> dict:
+    def proxy(self, data: dict, stream: bool = False) -> Generator[str, None, None]:
         """Proxy AI API requests to the configured AI provider."""
-        data["stream"] = False
+        data["stream"] = stream
+        return self.client.chat.completions.create(**data)
 
-        response = self.client.chat.completions.create(**data)
-        return response.model_dump()
+    def stream(self, data: dict) -> Generator[str, None, None]:
+        """Stream AI API requests to the configured AI provider."""
+        stream = self.proxy(data, stream=True)
+        for chunk in stream:
+            yield f"data: {chunk.model_dump_json()}\n\n"
+
+        yield "data: [DONE]\n\n"

src/backend/core/tests/test_api_config.py

@@ -21,6 +21,7 @@
     AI_BOT={"name": "Test Bot", "color": "#000000"},
     AI_FEATURE_ENABLED=False,
     AI_MODEL="test-model",
+    AI_STREAM=False,
     COLLABORATION_WS_URL="http://testcollab/",
     COLLABORATION_WS_NOT_CONNECTED_READY_ONLY=True,
     CRISP_WEBSITE_ID="123",
@@ -46,6 +47,7 @@ def test_api_config(is_authenticated):
         "AI_BOT": {"name": "Test Bot", "color": "#000000"},
         "AI_FEATURE_ENABLED": False,
         "AI_MODEL": "test-model",
+        "AI_STREAM": False,
         "COLLABORATION_WS_URL": "http://testcollab/",
         "COLLABORATION_WS_NOT_CONNECTED_READY_ONLY": True,
         "CRISP_WEBSITE_ID": "123",

src/backend/core/tests/test_services_ai_services.py

@@ -2,10 +2,9 @@
 Test ai API endpoints in the impress core app.
 """
 
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 from django.core.exceptions import ImproperlyConfigured
-from django.test.utils import override_settings
 
 import pytest
 from openai import OpenAIError
@@ -15,6 +14,15 @@
 pytestmark = pytest.mark.django_db
 
 
+@pytest.fixture(autouse=True)
+def ai_settings(settings):
+    """Fixture to set AI settings."""
+    settings.AI_MODEL = "llama"
+    settings.AI_BASE_URL = "http://example.com"
+    settings.AI_API_KEY = "test-key"
+    settings.AI_FEATURE_ENABLED = True
+
+
 @pytest.mark.parametrize(
     "setting_name, setting_value",
     [
@@ -23,22 +31,19 @@
         ("AI_MODEL", None),
     ],
 )
-def test_api_ai_setting_missing(setting_name, setting_value):
+def test_services_ai_setting_missing(setting_name, setting_value, settings):
     """Setting should be set"""
+    setattr(settings, setting_name, setting_value)
 
-    with override_settings(**{setting_name: setting_value}):
-        with pytest.raises(
-            ImproperlyConfigured,
-            match="AI configuration not set",
-        ):
-            AIService()
+    with pytest.raises(
+        ImproperlyConfigured,
+        match="AI configuration not set",
+    ):
+        AIService()
 
 
-@override_settings(
-    AI_BASE_URL="http://example.com", AI_API_KEY="test-key", AI_MODEL="test-model"
-)
 @patch("openai.resources.chat.completions.Completions.create")
-def test_api_ai__client_error(mock_create):
+def test_services_ai_proxy_client_error(mock_create):
     """Fail when the client raises an error"""
 
     mock_create.side_effect = OpenAIError("Mocked client error")
@@ -50,15 +55,11 @@ def test_api_ai__client_error(mock_create):
         AIService().proxy({"messages": [{"role": "user", "content": "hello"}]})
 
 
-@override_settings(
-    AI_BASE_URL="http://example.com", AI_API_KEY="test-key", AI_MODEL="test-model"
-)
 @patch("openai.resources.chat.completions.Completions.create")
-def test_api_ai__success(mock_create):
+def test_services_ai_proxy_success(mock_create):
     """The AI request should work as expect when called with valid arguments."""
 
-    mock_response = MagicMock()
-    mock_response.model_dump.return_value = {
+    mock_create.return_value = {
         "id": "chatcmpl-test",
         "object": "chat.completion",
         "created": 1234567890,
@@ -71,7 +72,6 @@ def test_api_ai__success(mock_create):
             }
         ],
     }
-    mock_create.return_value = mock_response
 
     response = AIService().proxy({"messages": [{"role": "user", "content": "hello"}]})
 
@@ -89,3 +89,47 @@ def test_api_ai__success(mock_create):
         ],
     }
     assert response == expected_response
+    mock_create.assert_called_once_with(
+        messages=[{"role": "user", "content": "hello"}], stream=False
+    )
+
+
+@patch("openai.resources.chat.completions.Completions.create")
+def test_services_ai_proxy_with_stream(mock_create):
+    """The AI request should work as expect when called with valid arguments."""
+
+    mock_create.return_value = {
+        "id": "chatcmpl-test",
+        "object": "chat.completion",
+        "created": 1234567890,
+        "model": "test-model",
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": "Salut"},
+                "finish_reason": "stop",
+            }
+        ],
+    }
+
+    response = AIService().proxy(
+        {"messages": [{"role": "user", "content": "hello"}]}, stream=True
+    )
+
+    expected_response = {
+        "id": "chatcmpl-test",
+        "object": "chat.completion",
+        "created": 1234567890,
+        "model": "test-model",
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": "Salut"},
+                "finish_reason": "stop",
+            }
+        ],
+    }
+    assert response == expected_response
+    mock_create.assert_called_once_with(
+        messages=[{"role": "user", "content": "hello"}], stream=True
+    )

src/backend/impress/settings.py

@@ -629,6 +629,9 @@ class Base(Configuration):
         default=False, environ_name="AI_FEATURE_ENABLED", environ_prefix=None
     )
     AI_MODEL = values.Value(None, environ_name="AI_MODEL", environ_prefix=None)
+    AI_STREAM = values.BooleanValue(
+        default=False, environ_name="AI_STREAM", environ_prefix=None
+    )
     AI_USER_RATE_THROTTLE_RATES = {
         "minute": 3,
         "hour": 50,

src/frontend/apps/e2e/__tests__/app-impress/common.ts

@@ -9,6 +9,7 @@ export const CONFIG = {
   },
   AI_FEATURE_ENABLED: true,
   AI_MODEL: 'llama',
+  AI_STREAM: false,
   CRISP_WEBSITE_ID: null,
   COLLABORATION_WS_URL: 'ws://localhost:4444/collaboration/ws/',
   COLLABORATION_WS_NOT_CONNECTED_READY_ONLY: false,

src/frontend/apps/impress/src/core/config/api/useConfig.tsx

@@ -15,6 +15,7 @@ export interface ConfigResponse {
   AI_BOT: { name: string; color: string };
   AI_FEATURE_ENABLED?: boolean;
   AI_MODEL?: string;
+  AI_STREAM: boolean;
   COLLABORATION_WS_URL?: string;
   COLLABORATION_WS_NOT_CONNECTED_READY_ONLY?: boolean;
   CRISP_WEBSITE_ID?: string;

src/frontend/apps/impress/src/features/docs/doc-editor/components/AI/useAI.tsx

@@ -33,7 +33,7 @@ export const useAI = (docId: Doc['id'], aiAllowed: boolean) => {
     const model = openai.chat(conf.AI_MODEL);
 
     const extension = createAIExtension({
-      stream: false,
+      stream: conf.AI_STREAM,
       model,
       agentCursor: conf?.AI_BOT,
       promptBuilder: promptBuilder(llmFormats.html.defaultPromptBuilder),