diff --git a/src/backend/core/api/viewsets.py b/src/backend/core/api/viewsets.py index 720d4b88..09563d8a 100644 --- a/src/backend/core/api/viewsets.py +++ b/src/backend/core/api/viewsets.py @@ -1650,8 +1650,20 @@ class DocumentViewSet( 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, + ) def _reject_invalid_ips(self, ips): """ @@ -2295,6 +2307,7 @@ class ConfigView(drf.views.APIView): "AI_BOT", "AI_FEATURE_ENABLED", "AI_MODEL", + "AI_STREAM", "COLLABORATION_WS_URL", "COLLABORATION_WS_NOT_CONNECTED_READY_ONLY", "CONVERSION_FILE_EXTENSIONS_ALLOWED", diff --git a/src/backend/core/services/ai_services.py b/src/backend/core/services/ai_services.py index 4f78249b..9a9c40e2 100644 --- a/src/backend/core/services/ai_services.py +++ b/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 @@ -26,9 +27,15 @@ class AIService: 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" diff --git a/src/backend/core/tests/test_api_config.py b/src/backend/core/tests/test_api_config.py index 587a4ad4..5e785b9c 100644 --- a/src/backend/core/tests/test_api_config.py +++ b/src/backend/core/tests/test_api_config.py @@ -22,6 +22,7 @@ pytestmark = pytest.mark.django_db 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", @@ -49,6 +50,7 @@ def test_api_config(is_authenticated): "AI_FEATURE_ENABLED": False, "AI_MODEL": "test-model", "AI_FEATURE_ENABLED": False, + "AI_STREAM": False, "COLLABORATION_WS_URL": "http://testcollab/", "COLLABORATION_WS_NOT_CONNECTED_READY_ONLY": True, "CONVERSION_FILE_EXTENSIONS_ALLOWED": [".docx", ".md"], diff --git a/src/backend/core/tests/test_services_ai_services.py b/src/backend/core/tests/test_services_ai_services.py index f6fa3703..5818dc32 100644 --- a/src/backend/core/tests/test_services_ai_services.py +++ b/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 @@ from core.services.ai_services import AIService 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 @@ pytestmark = pytest.mark.django_db ("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 + ) diff --git a/src/backend/impress/settings.py b/src/backend/impress/settings.py index d93ccf25..e5de62f2 100755 --- a/src/backend/impress/settings.py +++ b/src/backend/impress/settings.py @@ -696,6 +696,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, diff --git a/src/frontend/apps/e2e/__tests__/app-impress/utils-common.ts b/src/frontend/apps/e2e/__tests__/app-impress/utils-common.ts index 6c76ae70..620ee3fa 100644 --- a/src/frontend/apps/e2e/__tests__/app-impress/utils-common.ts +++ b/src/frontend/apps/e2e/__tests__/app-impress/utils-common.ts @@ -10,6 +10,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: true, diff --git a/src/frontend/apps/impress/src/core/config/api/useConfig.tsx b/src/frontend/apps/impress/src/core/config/api/useConfig.tsx index 9148f550..f326e580 100644 --- a/src/frontend/apps/impress/src/core/config/api/useConfig.tsx +++ b/src/frontend/apps/impress/src/core/config/api/useConfig.tsx @@ -18,6 +18,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; CONVERSION_FILE_EXTENSIONS_ALLOWED: string[]; diff --git a/src/frontend/apps/impress/src/features/docs/doc-editor/components/AI/useAI.tsx b/src/frontend/apps/impress/src/features/docs/doc-editor/components/AI/useAI.tsx index d9200f4b..2d20ad3b 100644 --- a/src/frontend/apps/impress/src/features/docs/doc-editor/components/AI/useAI.tsx +++ b/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),