electron-rare/docs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

✨(back) manage streaming with the ai service

Browse Source 
We want to handle both streaming or not when interacting with the AI
backend service.
This commit is contained in:
Manuel Raynaud
2025-06-12 11:08:23 +02:00
committed by Anthony LC
parent 0ec06e81d6
commit 96a759400a
8 changed files with 98 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/backend/core/api/viewsets.py
+15 -2
View File
@@ -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",
src/backend/core/services/ai_services.py
+11 -4
View File
@@ -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"
src/backend/core/tests/test_api_config.py
+2
View File
@@ -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"],
src/backend/core/tests/test_services_ai_services.py
+64 -20
View File
@@ -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
)
src/backend/impress/settings.py
+3
View File
@@ -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,
src/frontend/apps/e2e/__tests__/app-impress/utils-common.ts
+1
View File
@@ -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,
src/frontend/apps/impress/src/core/config/api/useConfig.tsx
+1
View File
@@ -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[];
src/frontend/apps/impress/src/features/docs/doc-editor/components/AI/useAI.tsx
+1 -1
View File
@@ -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),