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:
Quentin BEY
2026-02-02 12:34:19 +01:00
parent 23fa1d6b9e
commit 3106d5f25f
8 changed files with 87 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
+16 -10
View File
@@ -53,7 +53,7 @@ class BaseRagBackend(ABC):
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]
@@ -84,14 +84,15 @@ class BaseRagBackend(ABC):
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.
content (bytes): The content of the document as a bytes stream.
Returns:
str: The document content in Markdown format.
"""
raise NotImplementedError("Must be implemented in subclass.")
def store_document(self, name: str, content: str) -> None:
@abstractmethod
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.
@@ -99,10 +100,11 @@ class BaseRagBackend(ABC):
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.
@@ -110,23 +112,27 @@ class BaseRagBackend(ABC):
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: bytes, **kwargs
) -> str:
"""
Parse the document and store it in the Albert collection.
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.
content (bytes): The content of the document as a bytes 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:
@@ -151,7 +157,7 @@ class BaseRagBackend(ABC):
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.")
@@ -162,7 +168,7 @@ class BaseRagBackend(ABC):
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 str(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)
@@ -147,15 +147,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": [user_sub],
"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.
@@ -164,28 +165,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,
)
response.raise_for_status()
src/backend/chat/clients/pydantic_ai.py
+3 -1
View File
@@ -107,7 +107,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,
@@ -289,6 +289,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,
)
else:
# Remote URL
@@ -300,6 +301,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/"):
src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_upload.py
+25 -23
View File
@@ -78,10 +78,10 @@ def ai_settings(request, settings):
@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:
session = SessionStore()
session["oidc_access_token"] = "mocked-access-token"
mocked_refresh_access_token.return_value = session
yield mocked_refresh_access_token
@@ -103,10 +103,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"},
@@ -123,7 +131,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,
)
@@ -141,24 +149,20 @@ 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},
"usage": {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens},
},
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/",
@@ -172,10 +176,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,
@@ -267,8 +271,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,
@@ -597,8 +600,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
+19 -2
View File
@@ -37,11 +37,20 @@ 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_URL = "https://app-find"
settings.FIND_API_KEY = "find-api-key"
settings.AI_MODEL = "test-model"
settings.AI_AGENT_INSTRUCTIONS = "You are a helpful test assistant :)"
return settings
@@ -85,6 +94,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 +814,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/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