From b62fffc69ded9bffacfc040aa968c339f5f8d7f7 Mon Sep 17 00:00:00 2001 From: charles Date: Wed, 10 Dec 2025 15:31:31 +0100 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8(backend)=20implement=20FindRagBackend?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit We want to be able to use Find api in rag tools. I add a new rag backend class to do so. --- compose.yml | 11 + docs/llm-configuration.md | 4 +- docs/tools/web_search_brave.md | 1 + .../albert_rag_backend.py | 5 +- .../document_rag_backends/base_rag_backend.py | 18 +- .../document_rag_backends/find_rag_backend.py | 202 ++++++++++++++++++ src/backend/chat/clients/pydantic_ai.py | 4 +- src/backend/chat/tools/document_search_rag.py | 2 +- src/backend/chat/views.py | 5 + src/backend/conftest.py | 2 +- .../conversations/configuration/llm/dev.json | 43 ++++ src/backend/conversations/settings.py | 12 +- src/backend/utils/oicd.py | 26 +++ 13 files changed, 323 insertions(+), 12 deletions(-) create mode 100644 src/backend/chat/agent_rag/document_rag_backends/find_rag_backend.py create mode 100644 src/backend/conversations/configuration/llm/dev.json create mode 100644 src/backend/utils/oicd.py diff --git a/compose.yml b/compose.yml index c72dca0..c003776 100644 --- a/compose.yml +++ b/compose.yml @@ -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 diff --git a/docs/llm-configuration.md b/docs/llm-configuration.md index 8168e8e..55aeba9 100644 --- a/docs/llm-configuration.md +++ b/docs/llm-configuration.md @@ -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, diff --git a/docs/tools/web_search_brave.md b/docs/tools/web_search_brave.md index 8a00021..ced1822 100644 --- a/docs/tools/web_search_brave.md +++ b/docs/tools/web_search_brave.md @@ -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, ) ``` 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 8af560b..172df52 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 @@ -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. 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 78d25cd..f920bf9 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 @@ -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 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 new file mode 100644 index 0000000..e028890 --- /dev/null +++ b/src/backend/chat/agent_rag/document_rag_backends/find_rag_backend.py @@ -0,0 +1,202 @@ +"""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(f"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 = kwargs.get("session") + refresh_access_token(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, + ), + ) diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py index 71ac2dd..32c1484 100644 --- a/src/backend/chat/clients/pydantic_ai.py +++ b/src/backend/chat/clients/pydantic_ai.py @@ -91,6 +91,7 @@ class ContextDeps: conversation: models.ChatConversation user: User + session: Optional[Dict] = None web_search_enabled: bool = False @@ -105,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__(self, conversation: models.ChatConversation, user, model_hrid=None, language=None): + def __init__(self, conversation: models.ChatConversation, user, session=None, model_hrid=None, language=None): """ Initialize the AI agent service. @@ -135,6 +136,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, ) diff --git a/src/backend/chat/tools/document_search_rag.py b/src/backend/chat/tools/document_search_rag.py index 46827fa..92e311a 100644 --- a/src/backend/chat/tools/document_search_rag.py +++ b/src/backend/chat/tools/document_search_rag.py @@ -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, diff --git a/src/backend/chat/views.py b/src/backend/chat/views.py index 2e09011..3a93735 100644 --- a/src/backend/chat/views.py +++ b/src/backend/chat/views.py @@ -11,7 +11,10 @@ from django.http import Http404, StreamingHttpResponse import langfuse import magic import posthog +from django.utils.decorators import method_decorator from lasuite.malware_detection import malware_detection +from lasuite.oidc_login.backends import get_oidc_refresh_token +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 @@ -123,6 +126,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, @@ -174,6 +178,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 diff --git a/src/backend/conftest.py b/src/backend/conftest.py index 81fa0bc..371b69c 100644 --- a/src/backend/conftest.py +++ b/src/backend/conftest.py @@ -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): diff --git a/src/backend/conversations/configuration/llm/dev.json b/src/backend/conversations/configuration/llm/dev.json new file mode 100644 index 0000000..a004952 --- /dev/null +++ b/src/backend/conversations/configuration/llm/dev.json @@ -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" + } + ] +} diff --git a/src/backend/conversations/settings.py b/src/backend/conversations/settings.py index 56d2fe3..9868625 100755 --- a/src/backend/conversations/settings.py +++ b/src/backend/conversations/settings.py @@ -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. diff --git a/src/backend/utils/oicd.py b/src/backend/utils/oicd.py new file mode 100644 index 0000000..1f602e6 --- /dev/null +++ b/src/backend/utils/oicd.py @@ -0,0 +1,26 @@ +import requests +from django.conf import settings +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") + )