diff --git a/docs/llm-configuration.md b/docs/llm-configuration.md index 55aeba9..6706b13 100644 --- a/docs/llm-configuration.md +++ b/docs/llm-configuration.md @@ -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": { diff --git a/src/backend/chat/agent_rag/document_rag_backends/albert_rag_backend.py b/src/backend/chat/agent_rag/document_rag_backends/albert_rag_backend.py index 172df52..9753f4a 100644 --- a/src/backend/chat/agent_rag/document_rag_backends/albert_rag_backend.py +++ b/src/backend/chat/agent_rag/document_rag_backends/albert_rag_backend.py @@ -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( diff --git a/src/backend/chat/agent_rag/document_rag_backends/base_rag_backend.py b/src/backend/chat/agent_rag/document_rag_backends/base_rag_backend.py index deee2ce..32ff068 100644 --- a/src/backend/chat/agent_rag/document_rag_backends/base_rag_backend.py +++ b/src/backend/chat/agent_rag/document_rag_backends/base_rag_backend.py @@ -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: @@ -147,22 +153,22 @@ class BaseRagBackend(ABC): 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) diff --git a/src/backend/chat/agent_rag/document_rag_backends/find_rag_backend.py b/src/backend/chat/agent_rag/document_rag_backends/find_rag_backend.py index a3e3e8a..b2b14cf 100644 --- a/src/backend/chat/agent_rag/document_rag_backends/find_rag_backend.py +++ b/src/backend/chat/agent_rag/document_rag_backends/find_rag_backend.py @@ -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() diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py index f84b34b..fc71034 100644 --- a/src/backend/chat/clients/pydantic_ai.py +++ b/src/backend/chat/clients/pydantic_ai.py @@ -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/"): diff --git a/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_upload.py b/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_upload.py index f0941ab..74aae36 100644 --- a/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_upload.py +++ b/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_upload.py @@ -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, diff --git a/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_url.py b/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_url.py index 0f9f726..bb90233 100644 --- a/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_url.py +++ b/src/backend/chat/tests/views/chat/conversations/test_conversation_with_document_url.py @@ -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) diff --git a/src/backend/chat/tools/web_search_brave.py b/src/backend/chat/tools/web_search_brave.py index efba941..55b3521 100644 --- a/src/backend/chat/tools/web_search_brave.py +++ b/src/backend/chat/tools/web_search_brave.py @@ -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