Compare commits

...

2 Commits

Author SHA1 Message Date
charles fc58b43b9b 🧪(backend) tests
I add tests to test the app.
2025-12-17 17:01:13 +01:00
charles 0807f51a97 (backend) implement FindRagBackend
We want to be able to use Find api in rag tools.
I add a new rag backend class to do so.
2025-12-17 10:26:22 +01:00
15 changed files with 405 additions and 28 deletions
+11
View File
@@ -71,6 +71,9 @@ services:
- "host.docker.internal:host-gateway"
ports:
- "8071:8000"
networks:
- default
- lasuite
volumes:
- ./src/backend:/app
- ./data/static:/data/static
@@ -89,6 +92,9 @@ services:
image: nginx:1.25
ports:
- "8083:8083"
networks:
- default
- lasuite
volumes:
- ./docker/files/etc/nginx/conf.d:/etc/nginx/conf.d:ro
depends_on:
@@ -177,3 +183,8 @@ services:
kc_postgresql:
condition: service_healthy
restart: true
networks:
lasuite:
name: lasuite-network
driver: bridge
+2 -2
View File
@@ -244,8 +244,8 @@ For Mistral AI models using the Etalab platform:
{
"models": [
{
"hrid": "mistral-large",
"model_name": "mistral-large-latest",
"hrid": "mistral-medium",
"model_name": "mistral-medium-2508",
"human_readable_name": "Mistral Large (Etalab)",
"provider_name": "mistral-etalab",
"profile": null,
+1
View File
@@ -357,6 +357,7 @@ The RAG backend performs semantic search to find the most relevant content:
rag_results = document_store.search(
query,
results_count=settings.BRAVE_RAG_WEB_SEARCH_CHUNK_NUMBER,
**kwargs,
)
```
@@ -213,13 +213,14 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
logger.debug(response.json())
response.raise_for_status()
def search(self, query, results_count: int = 4) -> RAGWebResults:
def search(self, query: str, results_count: int = 4, **kwargs) -> RAGWebResults:
"""
Perform a search using the Albert API based on the provided query.
Args:
query (str): The search query.
results_count (int): The number of results to return.
**kwargs: Additional arguments.
Returns:
RAGWebResults: The search results.
@@ -256,7 +257,7 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
),
)
async def asearch(self, query, results_count: int = 4) -> RAGWebResults:
async def asearch(self, query, results_count: int = 4, **kwargs) -> RAGWebResults:
"""
Perform an asynchronous search using the Albert API based on the provided query.
@@ -142,17 +142,27 @@ class BaseRagBackend:
"""
return await sync_to_async(self.delete_collection)()
def search(self, query, results_count: int = 4) -> RAGWebResults:
def search(self, query: str, results_count: int = 4, **kwargs) -> RAGWebResults:
"""
Search the collection for the given query.
Args:
query: The search query string.
results_count: Number of results to return.
**kwargs: Additional arguments. Expected: 'session' for OIDC authentication.
"""
raise NotImplementedError("Must be implemented in subclass.")
async def asearch(self, query, results_count: int = 4) -> RAGWebResults:
async def asearch(self, query: str, results_count: int = 4, **kwargs) -> RAGWebResults:
"""
Search the collection for the given query.
Search the collection for the given query asynchronously.
Args:
query: The search query string.
results_count: Number of results to return.
**kwargs: Additional arguments. Expected: 'session' for OIDC authentication.
"""
return await sync_to_async(self.search)(query=query, results_count=results_count)
return await sync_to_async(self.search)(query=query, results_count=results_count, **kwargs)
@classmethod
@contextmanager
@@ -0,0 +1,201 @@
"""Implementation of the Find API for RAG document search."""
import logging
from io import BytesIO
from multiprocessing.context import AuthenticationError
from typing import List, Optional
from urllib.parse import urljoin
from uuid import uuid4
from django.conf import settings
from django.core.exceptions import ImproperlyConfigured
from django.utils import timezone
import requests
from chat.agent_rag.constants import RAGWebResult, RAGWebResults, RAGWebUsage
from chat.agent_rag.document_converter.markitdown import DocumentConverter
from chat.agent_rag.document_rag_backends.base_rag_backend import BaseRagBackend
from utils.oicd import refresh_access_token
logger = logging.getLogger(__name__)
class FindRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-attributes
"""
This class is a placeholder for the Find API implementation.
It is designed to be used with the RAG (Retrieval-Augmented Generation) document search system.
It provides methods to:
- Parse documents and convert them to Markdown format:
+ Handle PDF parsing using the Albert API.
+ Use the DocumentConverter (markitdown) for other formats.
- Store parsed documents in the Find index.
- Perform a search operation using the Find API.
"""
def __init__(
self,
collection_id: Optional[str] = None,
read_only_collection_id: Optional[List[str]] = None,
):
# Initialize any necessary parameters or configurations here
super().__init__(collection_id, read_only_collection_id)
self._pdf_parser_endpoint = urljoin(settings.ALBERT_API_URL, "/v1/parse-beta")
self.api_key = settings.FIND_API_KEY
self.search_endpoint = "api/v1.0/documents/search/"
self.indexing_endpoint = "api/v1.0/documents/index/"
if not self.api_key:
raise ImproperlyConfigured("FIND_API_KEY must be set in Django settings.")
def create_collection(self, name: str, description: Optional[str] = None) -> str:
"""
init collection_id
"""
self.collection_id = self.collection_id or 1
return self.collection_id
# TODO
def delete_collection(self) -> None:
"""
Deletion not available
"""
logger.warning("deletion of collections is not yet supported in FindRagBackend")
# TODO: factor with albert api
def parse_pdf_document(self, name: str, content_type: str, content: BytesIO) -> str:
"""
Parse the PDF document content and return the text content.
This method should handle the logic to convert the PDF into
a format suitable for the Albert API.
"""
response = requests.post(
self._pdf_parser_endpoint,
headers={
"Authorization": f"Bearer {settings.ALBERT_API_KEY}",
},
files={
"file": (
name,
content,
content_type,
), # Use the name as the filename in the request
"output_format": (None, "markdown"), # Specify the output format as Markdown,
},
timeout=settings.ALBERT_API_PARSE_TIMEOUT,
)
response.raise_for_status()
return "\n\n".join(
document_page["content"] for document_page in response.json().get("data", [])
)
# TODO: factor with albert api
def parse_document(self, name: str, content_type: str, content: BytesIO):
"""
Parse the document and prepare it for the search operation.
This method should handle the logic to convert the document
into a format suitable for the Find API.
Args:
name (str): The name of the document.
content_type (str): The MIME type of the document (e.g., "application/pdf").
content (BytesIO): The content of the document as a BytesIO stream.
Returns:
str: The document content in Markdown format.
"""
# Implement the parsing logic here
if content_type == "application/pdf":
# Handle PDF parsing
markdown_content = self.parse_pdf_document(
name=name, content_type=content_type, content=content
)
else:
markdown_content = DocumentConverter().convert_raw(
name=name, content_type=content_type, content=content
)
return markdown_content
def store_document(self, name: str, content: str) -> None:
"""
index document in Find
Args:
name (str): The name of the document.
content (str): The content of the document in Markdown format.
"""
logger.debug("index document '%s' in Find", name)
response = requests.post(
urljoin(settings.FIND_API_URL, self.indexing_endpoint),
headers={"Authorization": f"Bearer {self.api_key}"},
json={
"id": str(uuid4()),
"title": str(name) or "",
"depth": 0,
"path": str(name) or "",
"numchild": 0,
"content": content or "",
"created_at": timezone.now().isoformat(),
"updated_at": timezone.now().isoformat(),
"tags": [f"collection-{self.collection_id}"],
"size": len(content.encode("utf-8")),
"users": [], # TODO
"groups": [],
"reach": "public",
"is_active": True,
},
timeout=settings.ALBERT_API_TIMEOUT,
)
response.raise_for_status()
def search(self, query: str, results_count: int = 4, **kwargs) -> RAGWebResults:
"""
Perform a search using the Find API.
Uses the user's OIDC token from the request session.
Args:
query: The search query.
results_count: Number of results to return.
**kwargs: Additional arguments. Expected: 'session' containing OIDC tokens.
Returns:
RAGWebResults: The search results.
"""
logger.debug("search documents in Find with query '%s'", query)
# TODO: factor session auth in a decorator
session = refresh_access_token(kwargs.get("session"))
oidc_access_token = session.get("oidc_access_token")
if not oidc_access_token:
raise AuthenticationError({"error": "Not authenticated"})
collection_ids = self.get_all_collection_ids()
response = requests.post(
urljoin(settings.FIND_API_URL, self.search_endpoint),
headers={"Authorization": f"Bearer {oidc_access_token}"},
json={
"q": query,
"tags": [f"collection:{collection_id}" for collection_id in collection_ids],
"k": 10,
},
timeout=settings.ALBERT_API_TIMEOUT,
)
logger.debug(response.json())
response.raise_for_status()
return RAGWebResults(
data=[
RAGWebResult(
url=result["_source"]["title.fr"],
content=result["_source"]["content.fr"],
score=result["_score"],
)
for result in response.json()
],
usage=RAGWebUsage(
prompt_tokens=0,
completion_tokens=0,
),
)
+10 -1
View File
@@ -89,6 +89,7 @@ class ContextDeps:
conversation: models.ChatConversation
user: User
session: Optional[Dict] = None
web_search_enabled: bool = False
@@ -103,7 +104,14 @@ def get_model_configuration(model_hrid: str):
class AIAgentService: # pylint: disable=too-many-instance-attributes
"""Service class for AI-related operations (Pydantic-AI edition)."""
def __init__(self, conversation: models.ChatConversation, user, model_hrid=None, language=None):
def __init__( # noqa: PLR0913 # pylint: disable=too-many-arguments,too-many-positional-arguments
self,
conversation: models.ChatConversation,
user,
session=None,
model_hrid=None,
language=None,
):
"""
Initialize the AI agent service.
@@ -133,6 +141,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
self._context_deps = ContextDeps(
conversation=conversation,
user=user,
session=session,
web_search_enabled=self._is_web_search_enabled,
)
@@ -38,9 +38,6 @@ def brave_settings(settings):
settings.BRAVE_SEARCH_EXTRA_SNIPPETS = True
settings.BRAVE_SUMMARIZATION_ENABLED = False
settings.BRAVE_CACHE_TTL = 3600
settings.RAG_DOCUMENT_SEARCH_BACKEND = (
"chat.agent_rag.document_rag_backends.albert_rag_backend.AlbertRagBackend"
)
settings.BRAVE_RAG_WEB_SEARCH_CHUNK_NUMBER = 5
@@ -7,6 +7,7 @@ import json
import logging
from io import BytesIO
from unittest import mock
from unittest.mock import Mock
from django.utils import formats, timezone
@@ -41,28 +42,59 @@ from chat.tests.utils import replace_uuids_with_placeholder
pytestmark = pytest.mark.django_db(transaction=True)
@pytest.fixture(autouse=True)
def ai_settings(settings):
@pytest.fixture(
autouse=True,
params=[
"chat.agent_rag.document_rag_backends.find_rag_backend.FindRagBackend",
"chat.agent_rag.document_rag_backends.albert_rag_backend.AlbertRagBackend",
],
)
def ai_settings(request, settings):
"""Fixture to set AI service URLs for testing."""
settings.AI_BASE_URL = "https://www.external-ai-service.com/"
settings.AI_API_KEY = "test-api-key"
settings.AI_MODEL = "test-model"
settings.AI_AGENT_INSTRUCTIONS = "You are a helpful test assistant :)"
# Enable Albert API for document search
settings.RAG_DOCUMENT_SEARCH_BACKEND = (
"chat.agent_rag.document_rag_backends.albert_rag_backend.AlbertRagBackend"
)
settings.ALBERT_API_URL = "https://albert.api.etalab.gouv.fr"
settings.ALBERT_API_KEY = "albert-api-key"
# enable on rag document search tool
settings.RAG_DOCUMENT_SEARCH_BACKEND = request.param
settings.RAG_WEB_SEARCH_PROMPT_UPDATE = (
"Based on the following document contents:\n\n{search_results}\n\n"
"Please answer the user's question: {user_prompt}"
)
settings.AI_BASE_URL = "https://www.external-ai-service.com/"
settings.AI_API_KEY = "test-api-key"
settings.AI_MODEL = "test-model"
settings.AI_AGENT_INSTRUCTIONS = "You are a helpful test assistant :)"
# Albert API settings
settings.ALBERT_API_URL = "https://albert.api.etalab.gouv.fr"
settings.ALBERT_API_KEY = "albert-api-key"
# Find API settings
settings.FIND_API_URL = "https://find.api.example.com"
settings.FIND_API_KEY = "find-api-key"
return settings
@pytest.fixture(autouse=True)
def mock_process_request():
"""Mock process_request to bypass authentication in tests."""
with mock.patch(
"lasuite.oidc_login.decorators.RefreshOIDCAccessToken.process_request"
) as mocked_process_request:
mocked_process_request.return_value = None
yield mocked_process_request
@pytest.fixture(autouse=True)
def mock_refresh_access_token():
"""Mock refresh_access_token to bypass token refresh in tests."""
with mock.patch("utils.oicd.refresh_access_token") as mocked_refresh_access_token:
mock_session = Mock(spec=httpx.Client)
mock_session.access_token = "mocked-access-token"
mocked_refresh_access_token.return_value = mock_session
yield mocked_refresh_access_token
@pytest.fixture(name="sample_pdf_content")
def fixture_sample_pdf_content():
"""Create a dummy PDF content as BytesIO."""
@@ -134,6 +166,32 @@ def fixture_mock_albert_api():
)
@pytest.fixture(name="mock_find_api")
def fixture_mock_find_api():
"""Fixture to mock the Find API endpoints."""
# Mock document indexing (Find API)
responses.post(
"https://find.api.example.com/api/v1.0/documents/index/",
json={"id": "456", "status": "indexed"},
status=status.HTTP_200_OK,
)
# Mock document search (Find API)
responses.post(
"https://find.api.example.com/api/v1.0/documents/search/",
json=[
{
"_source": {
"title.fr": "sample.pdf",
"content.fr": "This is the content of the PDF.",
},
"_score": 0.9,
}
],
status=status.HTTP_200_OK,
)
@pytest.fixture(name="mock_summarization_agent")
def fixture_mock_summarization_agent():
"""Mock the SummarizationAgent to return a fixed summary."""
@@ -216,9 +274,10 @@ def fixture_mock_openai_stream():
@responses.activate
@respx.mock
@freeze_time()
def test_post_conversation_with_document_upload( # pylint: disable=too-many-arguments,too-many-positional-arguments
def test_post_conversation_with_document_upload( # noqa: PLR0913 pylint: disable=too-many-arguments,too-many-positional-arguments
api_client,
mock_albert_api, # pylint: disable=unused-argument
mock_find_api, # pylint: disable=unused-argument
sample_pdf_content,
today_promt_date,
mock_ai_agent_service,
@@ -570,6 +629,7 @@ def test_post_conversation_with_document_upload_feature_disabled( # pylint: dis
def test_post_conversation_with_document_upload_summarize( # pylint: disable=too-many-arguments,too-many-positional-arguments # noqa: PLR0913
api_client,
mock_albert_api, # pylint: disable=unused-argument
mock_find_api, # pylint: disable=unused-argument
sample_pdf_content,
today_promt_date,
mock_ai_agent_service,
@@ -26,7 +26,7 @@ def add_document_rag_search_tool(agent: Agent) -> None:
document_store = document_store_backend(ctx.deps.conversation.collection_id)
rag_results = document_store.search(query)
rag_results = document_store.search(query, session=ctx.deps.session)
ctx.usage += RunUsage(
input_tokens=rag_results.usage.prompt_tokens,
+4
View File
@@ -7,11 +7,13 @@ from uuid import uuid4
from django.conf import settings
from django.core.files.storage import default_storage
from django.http import Http404, StreamingHttpResponse
from django.utils.decorators import method_decorator
import langfuse
import magic
import posthog
from lasuite.malware_detection import malware_detection
from lasuite.oidc_login.decorators import refresh_oidc_access_token
from rest_framework import decorators, filters, mixins, permissions, status, viewsets
from rest_framework.exceptions import MethodNotAllowed, PermissionDenied, ValidationError
from rest_framework.response import Response
@@ -122,6 +124,7 @@ class ChatViewSet( # pylint: disable=too-many-ancestors, abstract-method
self.permission_classes = []
return super().get_permissions()
@method_decorator(refresh_oidc_access_token)
@decorators.action(
methods=["post"],
detail=True,
@@ -173,6 +176,7 @@ class ChatViewSet( # pylint: disable=too-many-ancestors, abstract-method
ai_service = AIAgentService(
conversation=conversation,
user=self.request.user,
session=request.session,
model_hrid=model_hrid,
language=(
self.request.user.language
+1 -1
View File
@@ -22,7 +22,7 @@ def no_http_requests(monkeypatch):
Credits: https://blog.jerrycodes.com/no-http-requests/
"""
allowed_hosts = {"localhost", "minio", "minio:9000"}
allowed_hosts = {"localhost", "127.0.0.1", "minio", "minio:9000"}
original_urlopen = HTTPConnectionPool.urlopen
def urlopen_mock(self, method, url, *args, **kwargs):
@@ -0,0 +1,43 @@
{
"models": [
{
"hrid": "default-model",
"model_name": "settings.AI_MODEL",
"human_readable_name": "Default Model",
"provider_name": "default-provider",
"profile": null,
"settings": {},
"is_active": true,
"icon": [
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABwAAAAcCAMAAABF0y+mAAAAn1BMVEUALosAKoovTZjw8vb////+9/jlPUniAAz",
"iABUAGIWbpsTwq7HhAAAAI4dle7DrdX4AJohRaaboXWj7+/zn6On5//9NZaT29vfoWmVHYKDoUl/k5OUAIYddc6vpbHYCM47Y3+v53+LiFCUA",
"HIWnsckYPJHi6PL77O7jJjW3wdf1w8jre4QgQ5TZ2txwg7Pr3+I8WZ6OnsTuoamClL7tlZ5xz5y8AAAAzUlEQVR4AZ3RRQKDQBBEUSTu7h5c4",
"vc/W6Yp3KG2Dz4ynDdeEBvOmq12xx2E1u0B+4NOEocj4DgNJ1PgLAvni8WyBq5Yc71ubFJx23C2q4P7dRYejg1xzvCUgvz5guz11k7gXYKF/1",
"8oyiYuvHAYeVkhXCzolVStHcGDjiQzNmMQxsMI5rEJRdQSPZvbpE2E8aY6gC6Z+2Hg4dFA0Yb4YedNL/v4Fk8WJuwiGhrChJNXI210rnib9Fs",
"JlXRUC/HwTscPIXf/iklq/tjb/gHAdxkCUjAg2QAAAABJRU5ErkJggg=="
],
"system_prompt": "settings.AI_AGENT_INSTRUCTIONS",
"tools": "settings.AI_AGENT_TOOLS"
},
{
"hrid": "default-summarization-model",
"model_name": "settings.AI_MODEL",
"human_readable_name": "Default Summarization Model",
"provider_name": "default-provider",
"profile": null,
"settings": {},
"is_active": true,
"icon": null,
"system_prompt": "settings.SUMMARIZATION_SYSTEM_PROMPT",
"tools": []
}
],
"providers": [
{
"hrid": "default-provider",
"base_url": "settings.AI_BASE_URL",
"api_key": "settings.AI_API_KEY",
"kind": "mistral"
}
]
}
+11 -1
View File
@@ -713,7 +713,7 @@ class Base(BraveSettings, Configuration):
environ_prefix=None,
)
RAG_DOCUMENT_SEARCH_BACKEND = values.Value(
"chat.agent_rag.document_rag_backends.albert_rag_backend.AlbertRagBackend",
"chat.agent_rag.document_rag_backends.find_rag_backend.FindRagBackend",
environ_name="RAG_DOCUMENT_SEARCH_BACKEND",
environ_prefix=None,
)
@@ -841,6 +841,16 @@ USER QUESTION:
environ_prefix=None,
)
# Find
FIND_API_KEY = values.Value(
environ_name="FIND_API_KEY",
environ_prefix=None,
)
FIND_API_URL = values.Value(
environ_name="FIND_API_URL",
environ_prefix=None,
)
# Logging
# We want to make it easy to log to console but by default we log production
# to Sentry and don't want to log to console.
+30
View File
@@ -0,0 +1,30 @@
"""Utility functions for OIDC token management."""
from django.conf import settings
import requests
from lasuite.oidc_login.backends import get_oidc_refresh_token, store_tokens
def refresh_access_token(session):
"""Refresh the OIDC access token using the refresh token."""
response = requests.post(
settings.OIDC_OP_TOKEN_ENDPOINT,
data={
"grant_type": "refresh_token",
"client_id": settings.OIDC_RP_CLIENT_ID,
"client_secret": settings.OIDC_RP_CLIENT_SECRET,
"refresh_token": get_oidc_refresh_token(session),
},
timeout=5,
)
response.raise_for_status()
token_info = response.json()
store_tokens(
session,
access_token=token_info.get("access_token"),
id_token=None,
refresh_token=token_info.get("refresh_token"),
)
return session