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

✨(backend) enhance Find API integration with user sub and tag

Browse Source 
I enhance Find API integration with user
access control and configuration options
This commit is contained in:
charles
2025-12-18 09:49:11 +01:00
parent 713b34fdcd
commit 0abe12382a
11 changed files with 132 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/llm-configuration.md
+1 -1
View File
@@ -246,7 +246,7 @@ For Mistral AI models using the Etalab platform:
{
"hrid": "mistral-medium",
"model_name": "mistral-medium-2508",
"human_readable_name": "Mistral Large (Etalab)",
"human_readable_name": "Mistral Medium (Etalab)",
"provider_name": "mistral-etalab",
"profile": null,
"settings": {
src/backend/chat/agent_rag/document_rag_backends/albert_rag_backend.py
+4 -2
View File
@@ -166,7 +166,7 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
return markdown_content
def store_document(self, name: str, content: str) -> None:
def store_document(self, name: str, content: str, **kwargs) -> None:
"""
Store the document content in the Albert collection.
This method should handle the logic to send the document content to the Albert API.
@@ -174,6 +174,7 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
Args:
name (str): The name of the document.
content (str): The content of the document in Markdown format.
**kwargs: Additional arguments.
"""
response = requests.post(
urljoin(self._base_url, self._documents_endpoint),
@@ -188,7 +189,7 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
logger.debug(response.json())
response.raise_for_status()
async def astore_document(self, name: str, content: str) -> None:
async def astore_document(self, name: str, content: str, **kwargs) -> None:
"""
Store the document content in the Albert collection.
This method should handle the logic to send the document content to the Albert API.
@@ -196,6 +197,7 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att
Args:
name (str): The name of the document.
content (str): The content of the document in Markdown format.
**kwargs: Additional arguments.
"""
async with httpx.AsyncClient(timeout=settings.ALBERT_API_TIMEOUT) as client:
response = await client.post(
src/backend/chat/agent_rag/document_rag_backends/base_rag_backend.py
+15 -10
View File
@@ -53,7 +53,7 @@ class BaseRagBackend:
collection_ids = []
if self.collection_id:
collection_ids.append(int(self.collection_id))
collection_ids.append(self.collection_id)
if self.read_only_collection_id:
collection_ids.extend(
[int(collection_id) for collection_id in self.read_only_collection_id]
@@ -90,7 +90,7 @@ class BaseRagBackend:
"""
raise NotImplementedError("Must be implemented in subclass.")
def store_document(self, name: str, content: str) -> None:
def store_document(self, name: str, content: str, **kwargs) -> None:
"""
Store the document content in the collection.
This method should handle the logic to send the document content to the API.
@@ -98,10 +98,11 @@ class BaseRagBackend:
Args:
name (str): The name of the document.
content (str): The content of the document in Markdown format.
**kwargs: Additional arguments. ex: "user_sub" for access control.
"""
raise NotImplementedError("Must be implemented in subclass.")
async def astore_document(self, name: str, content: str) -> None:
async def astore_document(self, name: str, content: str, **kwargs) -> None:
"""
Store the document content in the collection.
This method should handle the logic to send the document content to the API.
@@ -109,10 +110,13 @@ class BaseRagBackend:
Args:
name (str): The name of the document.
content (str): The content of the document in Markdown format.
**kwargs: Additional arguments. ex: "user_sub" for access control.
"""
return await sync_to_async(self.store_document)(name=name, content=content)
return await sync_to_async(self.store_document)(name=name, content=content, **kwargs)
def parse_and_store_document(self, name: str, content_type: str, content: BytesIO) -> str:
def parse_and_store_document(
self, name: str, content_type: str, content: BytesIO, **kwargs
) -> str:
"""
Parse the document and store it in the Albert collection.
@@ -120,12 +124,13 @@ class BaseRagBackend:
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.
**kwargs: Additional arguments. ex: "user_sub" for access control.
"""
if not self.collection_id:
raise RuntimeError("The RAG backend requires collection_id")
document_content = self.parse_document(name, content_type, content)
self.store_document(name, document_content)
self.store_document(name, document_content, **kwargs)
return document_content
def delete_collection(self) -> None:
@@ -145,22 +150,22 @@ class BaseRagBackend:
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.
**kwargs: Additional arguments. ex: 'session' for OIDC authentication.
"""
raise NotImplementedError("Must be implemented in subclass.")
async def asearch(self, query: str, results_count: int = 4, **kwargs) -> RAGWebResults:
"""
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.
**kwargs: Additional arguments. ex: 'session' for OIDC authentication.
"""
return await sync_to_async(self.search)(query=query, results_count=results_count, **kwargs)
src/backend/chat/agent_rag/document_rag_backends/find_rag_backend.py
+15 -20
View File
@@ -1,8 +1,8 @@
"""Implementation of the Find API for RAG document search."""
import logging
import uuid
from io import BytesIO
from multiprocessing.context import AuthenticationError
from typing import List, Optional
from urllib.parse import urljoin
from uuid import uuid4
@@ -16,7 +16,7 @@ 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
from utils.oidc import with_fresh_access_token
logger = logging.getLogger(__name__)
@@ -49,14 +49,13 @@ class FindRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-attri
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:
def create_collection(self, name: str, description: Optional[str] = None) -> uuid.UUID:
"""
init collection_id
"""
self.collection_id = self.collection_id or 1
self.collection_id = self.collection_id or uuid.uuid4()
return self.collection_id
# TODO
def delete_collection(self) -> None:
"""
Deletion not available
@@ -119,13 +118,14 @@ class FindRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-attri
return markdown_content
def store_document(self, name: str, content: str) -> None:
def store_document(self, name: str, content: str, **kwargs) -> None:
"""
index document in Find
Args:
name (str): The name of the document.
content (str): The content of the document in Markdown format.
user_sub (str): The user subject identifier for access control.
"""
logger.debug("index document '%s' in Find", name)
response = requests.post(
@@ -142,15 +142,16 @@ class FindRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-attri
"updated_at": timezone.now().isoformat(),
"tags": [f"collection-{self.collection_id}"],
"size": len(content.encode("utf-8")),
"users": [], # TODO
"users": [kwargs["user_sub"]] if "user_sub" in kwargs else [],
"groups": [],
"reach": "public",
"reach": "authenticated",
"is_active": True,
},
timeout=settings.ALBERT_API_TIMEOUT,
timeout=settings.FIND_API_TIMEOUT,
)
response.raise_for_status()
@with_fresh_access_token
def search(self, query: str, results_count: int = 4, **kwargs) -> RAGWebResults:
"""
Perform a search using the Find API.
@@ -159,28 +160,22 @@ class FindRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-attri
Args:
query: The search query.
results_count: Number of results to return.
**kwargs: Additional arguments. Expected: 'session' containing OIDC tokens.
**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}"},
headers={"Authorization": f"Bearer {kwargs["session"].get("oidc_access_token")}"},
json={
"q": query,
"tags": [f"collection:{collection_id}" for collection_id in collection_ids],
"k": 10,
"tags": [f"collection-{collection_id}" for collection_id in self.get_all_collection_ids()],
"k": results_count,
},
timeout=settings.ALBERT_API_TIMEOUT,
timeout=settings.FIND_API_TIMEOUT,
)
logger.debug(response.json())
response.raise_for_status()
src/backend/chat/clients/pydantic_ai.py
+7 -3
View File
@@ -106,7 +106,7 @@ 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__( # noqa: PLR0913 # pylint: disable=too-many-arguments,too-many-positional-arguments
def __init__( # pylint: disable=too-many-arguments,too-many-positional-arguments
self,
conversation: models.ChatConversation,
user,
@@ -242,7 +242,9 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
# --------------------------------------------------------------------- #
# Core agent runner
# --------------------------------------------------------------------- #
async def parse_input_documents(self, documents: List[BinaryContent | DocumentUrl]):
async def parse_input_documents(
self, documents: List[BinaryContent | DocumentUrl], user_sub: str
):
"""
Parse and store input documents in the conversation's document store.
"""
@@ -286,6 +288,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
name=document.identifier,
content_type=document.media_type,
content=document_data,
user_sub=user_sub,
)
else:
# Remote URL
@@ -295,6 +298,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
name=document.identifier,
content_type=document.media_type,
content=document.data,
user_sub=self.user.sub,
)
if not document.media_type.startswith("text/"):
@@ -432,7 +436,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes
)
try:
await self.parse_input_documents(input_documents)
await self.parse_input_documents(input_documents, user_sub=self.user.sub)
except Exception as exc: # pylint: disable=broad-except
logger.exception("Error parsing input documents: %s", exc)
yield events_v4.ToolResultPart(
src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_upload.py
+22 -22
View File
@@ -88,10 +88,8 @@ def mock_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
with mock.patch("utils.oidc.refresh_access_token") as mocked_refresh_access_token:
mocked_refresh_access_token.return_value = Mock(spec=httpx.Client)
yield mocked_refresh_access_token
@@ -113,10 +111,18 @@ def fixture_sample_pdf_content():
return BytesIO(pdf_data)
@pytest.fixture(name="mock_albert_api")
def fixture_mock_albert_api():
@pytest.fixture(name="mock_document_api")
def fixture_mock_document_api():
"""Fixture to mock the Albert API endpoints."""
# Mock collection creation
document_name = "sample.pdf"
document_content = "This is the content of the PDF."
prompt_tokens = 10
completion_tokens = 20
search_method = "semantic"
search_score = 0.9
responses.post(
"https://albert.api.etalab.gouv.fr/v1/collections",
json={"id": "123", "name": "test-collection"},
@@ -133,7 +139,7 @@ def fixture_mock_albert_api():
"metadata": {"document_name": "sample.pdf"},
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 20},
"usage": {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens},
},
status=status.HTTP_200_OK,
)
@@ -151,13 +157,13 @@ def fixture_mock_albert_api():
json={
"data": [
{
"method": "semantic",
"method": search_method,
"chunk": {
"id": 123,
"content": "This is the content of the PDF.",
"metadata": {"document_name": "sample.pdf"},
"content": document_content,
"metadata": {"document_name": document_name},
},
"score": 0.9,
"score": search_score,
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 20},
@@ -165,10 +171,6 @@ def fixture_mock_albert_api():
status=status.HTTP_200_OK,
)
@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/",
@@ -182,10 +184,10 @@ def fixture_mock_find_api():
json=[
{
"_source": {
"title.fr": "sample.pdf",
"content.fr": "This is the content of the PDF.",
"title.fr": document_name,
"content.fr": document_content,
},
"_score": 0.9,
"_score": search_score,
}
],
status=status.HTTP_200_OK,
@@ -277,8 +279,7 @@ def fixture_mock_openai_stream():
def test_post_conversation_with_document_upload(
# 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
mock_document_api, # pylint: disable=unused-argument
sample_pdf_content,
today_promt_date,
mock_ai_agent_service,
@@ -607,8 +608,7 @@ def test_post_conversation_with_document_upload_feature_disabled(
@freeze_time()
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
mock_document_api, # pylint: disable=unused-argument
sample_pdf_content,
today_promt_date,
mock_ai_agent_service,
src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_url.py
+29 -2
View File
@@ -4,6 +4,7 @@ import uuid
# pylint: disable=too-many-lines
from io import BytesIO
from unittest import mock
from django.core.files.storage import default_storage
from django.utils import formats, timezone
@@ -37,11 +38,19 @@ 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.RAG_DOCUMENT_SEARCH_BACKEND = request.param
settings.AI_BASE_URL = "https://www.external-ai-service.com/"
settings.AI_API_KEY = "test-api-key"
settings.FIND_API_KEY = "find-api-key"
settings.AI_MODEL = "test-model"
settings.AI_AGENT_INSTRUCTIONS = "You are a helpful test assistant :)"
return settings
@@ -58,6 +67,16 @@ def fixture_sample_document_content():
)
@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
@responses.activate
@freeze_time()
def test_post_conversation_with_local_pdf_document_url(
@@ -85,6 +104,10 @@ def test_post_conversation_with_local_pdf_document_url(
json={"id": "document_id", "object": "document"},
status=200,
)
responses.post(
"https://app-find/api/v1.0/documents/index/",
status=200,
)
chat_conversation = ChatConversationFactory(owner__language="en-us")
api_client.force_authenticate(user=chat_conversation.owner)
@@ -801,6 +824,10 @@ def test_post_conversation_with_local_not_pdf_document_url(
json={"id": "document_id", "object": "document"},
status=200,
)
responses.post(
"https://app-find/api/v1.0/documents/index/",
status=200,
)
chat_conversation = ChatConversationFactory(owner__language="en-us")
api_client.force_authenticate(user=chat_conversation.owner)
src/backend/chat/tests/views/chat/conversations/test_conversation_with_image_url.py
+11
View File
@@ -1,6 +1,7 @@
"""Unit tests for chat conversation actions with image URL."""
import uuid
from unittest import mock
from django.utils import formats, timezone
@@ -53,6 +54,16 @@ def fixture_sample_image_content():
)
@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
@freeze_time("2025-10-18T20:48:20.286204Z")
def test_post_conversation_with_local_image_url(
api_client,
src/backend/chat/tools/web_search_brave.py
+2 -2
View File
@@ -127,7 +127,7 @@ async def _extract_and_summarize_snippets_async(query: str, url: str) -> List[st
return []
async def _fetch_and_store_async(url: str, document_store) -> None:
async def _fetch_and_store_async(url: str, document_store, **kwargs) -> None:
"""Fetch, extract and store text content from the URL in the document store."""
try:
@@ -136,7 +136,7 @@ async def _fetch_and_store_async(url: str, document_store) -> None:
logger.debug("Fetched document: %s", document)
if document:
await document_store.astore_document(url, document)
await document_store.astore_document(url, document, **kwargs)
except DocumentFetchError as e:
logger.warning("Failed to fetch and store %s: %s", url, e)
# Continue with other documents
src/backend/conversations/settings.py
+8 -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.find_rag_backend.FindRagBackend",
"chat.agent_rag.document_rag_backends.albert_rag_backend.AlbertRagBackend",
environ_name="RAG_DOCUMENT_SEARCH_BACKEND",
environ_prefix=None,
)
@@ -843,13 +843,20 @@ USER QUESTION:
# Find
FIND_API_KEY = values.Value(
None,
environ_name="FIND_API_KEY",
environ_prefix=None,
)
FIND_API_URL = values.Value(
"https://app-find/api",
environ_name="FIND_API_URL",
environ_prefix=None,
)
FIND_API_TIMEOUT = values.PositiveIntegerValue(
default=30, # seconds
environ_name="FIND_API_TIMEOUT",
environ_prefix=None,
)
# Logging
# We want to make it easy to log to console but by default we log production
src/backend/utils/oicd.py → src/backend/utils/oidc.py
+18
View File
@@ -1,9 +1,12 @@
"""Utility functions for OIDC token management."""
from functools import wraps
from django.conf import settings
import requests
from lasuite.oidc_login.backends import get_oidc_refresh_token, store_tokens
from rest_framework.exceptions import AuthenticationFailed
def refresh_access_token(session):
@@ -28,3 +31,18 @@ def refresh_access_token(session):
refresh_token=token_info.get("refresh_token"),
)
return session
def with_fresh_access_token(func):
"""
Decorator to handle OIDC token refresh and extraction.
Expects 'session' in kwargs and update it with the fresh token.
"""
@wraps(func)
def wrapper(*args, **kwargs):
session = kwargs.get("session")
if session is None:
raise AuthenticationFailed({"error": "Session is required but not provided"})
kwargs["session"] = refresh_access_token(session)
return func(*args, **kwargs)
return wrapper