From dd9d760659b8d1be54d94a8b04b78a5cb42c3339 Mon Sep 17 00:00:00 2001 From: camilleAND Date: Wed, 22 Oct 2025 16:37:49 +0200 Subject: [PATCH] =?UTF-8?q?=F0=9F=93=9D(service-public):=20use=20existing?= =?UTF-8?q?=20RAG=20implementation=20and=20upgrade=20for=20metadata?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/chat/agent_rag/constants.py | 1 + .../albert_rag_backend.py | 15 ++- .../chat/agent_rag/web_search/albert_api.py | 1 + src/backend/chat/tools/service_public.py | 115 ++++++------------ 4 files changed, 53 insertions(+), 79 deletions(-) diff --git a/src/backend/chat/agent_rag/constants.py b/src/backend/chat/agent_rag/constants.py index 0fb7d2c..f258874 100644 --- a/src/backend/chat/agent_rag/constants.py +++ b/src/backend/chat/agent_rag/constants.py @@ -22,6 +22,7 @@ class RAGWebResult(BaseModel): score: float = Field( ..., description="Relevance score of the web result, typically between 0 and 1." ) + metadata: dict = Field(..., description="Metadata of the web result.") class RAGWebResults(BaseModel): 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 4de73bb..d548646 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 @@ -3,7 +3,7 @@ import json import logging from io import BytesIO -from typing import Optional +from typing import List, Optional from urllib.parse import urljoin from django.conf import settings @@ -150,22 +150,30 @@ 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, results_count: int = 4, collections: Optional[List[int]] = None) -> 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. + collections (Optional[List[int]]): List of collection IDs to search in. + If None, uses the current collection_id. Returns: RAGWebResults: The search results. """ + # Use provided collections or fall back to current collection_id + if collections is not None: + collection_list = collections + else: + collection_list = [int(self.collection_id)] + response = requests.post( urljoin(self._base_url, self._search_endpoint), headers=self._headers, json={ - "collections": [int(self.collection_id)], + "collections": collection_list, "prompt": query, "score_threshold": 0.6, "k": results_count, # Number of chunks to return from the search @@ -182,6 +190,7 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att url=result.chunk.metadata["document_name"], content=result.chunk.content, score=result.score, + metadata=result.chunk.metadata, ) for result in searches.data ], diff --git a/src/backend/chat/agent_rag/web_search/albert_api.py b/src/backend/chat/agent_rag/web_search/albert_api.py index 9c56db6..f92a13c 100644 --- a/src/backend/chat/agent_rag/web_search/albert_api.py +++ b/src/backend/chat/agent_rag/web_search/albert_api.py @@ -103,6 +103,7 @@ class AlbertWebSearchManager(BaseWebSearchManager): url=self._clean_url(result.chunk.metadata["document_name"]), content=result.chunk.content, score=result.score, + metadata=result.chunk.metadata, ) for result in searches.data ], diff --git a/src/backend/chat/tools/service_public.py b/src/backend/chat/tools/service_public.py index 6a6a060..36b9718 100644 --- a/src/backend/chat/tools/service_public.py +++ b/src/backend/chat/tools/service_public.py @@ -1,55 +1,22 @@ """Service Public RAG search tool using Albert API pre-defined collections. -This tool reuses the existing Albert API RAG by directly querying the search -endpoint with a fixed set of curated collections (e.g. Service-Public, Travail-Emploi). +This tool reuses the existing AlbertRagBackend to query curated collections +(e.g. Service-Public, Travail-Emploi) without creating temporary collections. """ -from typing import List -import json import logging -from urllib.parse import urljoin +from typing import List -import requests from django.conf import settings +from django.utils.module_loading import import_string from pydantic_ai import RunContext, RunUsage from pydantic_ai.messages import ToolReturn - logger = logging.getLogger(__name__) - # Default curated collections (Albert IDs) DEFAULT_COLLECTION_IDS: List[int] = [784, 785] # travail-emploi, service-public -PROMPT_PREFIX = "Voilà les informations trouvées, résume les pour répondre à la question de l'utilisateur, à la fin de ta réponse, ajoutes une section sources avec les urls des sources si présentes: " - - -def _albert_search_with_collections(query: str, collections: List[int]) -> dict: - """Call Albert search with explicit collections. - - Returns a dict compatible with existing RAG result mapping tooling. - """ - base_url = settings.ALBERT_API_URL - api_key = settings.ALBERT_API_KEY - endpoint = urljoin(base_url, "v1/search") - - # Minimal payload aligned with Albert API - payload = { - "collections": collections, - "prompt": query, - # Reasonable defaults; can be made configurable later if needed - "k": getattr(settings, "RAG_WEB_SEARCH_CHUNK_NUMBER", 10), - "web_search": False, - } - headers = { - "accept": "application/json", - "Authorization": f"Bearer {api_key}", - "Content-Type": "application/json", - } - - response = requests.post(endpoint, headers=headers, json=payload, timeout=settings.ALBERT_API_TIMEOUT) - response.raise_for_status() - return response.json() - +INSTRUCTIONS = "Voilà les informations trouvées, résume les pour répondre à la question de l'utilisateur, à la fin de ta réponse, ajoutes une section sources avec les urls des sources si présentes: \n" async def service_public(ctx: RunContext, query: str) -> ToolReturn: """Search curated Service-Public collections on Albert and return snippets. @@ -59,45 +26,41 @@ async def service_public(ctx: RunContext, query: str) -> ToolReturn: query: The user query to search within curated collections """ try: - json_response = _albert_search_with_collections(query, DEFAULT_COLLECTION_IDS) + # Use AlbertRagBackend to search in specific collections + backend_class = import_string(settings.RAG_DOCUMENT_SEARCH_BACKEND) + backend = backend_class() + + # Search in the curated collections + rag_results = backend.search(query, collections=DEFAULT_COLLECTION_IDS) + + # Convert to compact format for the model using your logic + compact = [] + sources = [] + for result in rag_results.data: + # AlbertRagBackend.search() returns RAGWebResult objects with {url, content, score, metadata} + compact.append( + { + "title": result.metadata["document_name"], + "snippet": result.content, + "url": result.metadata["url"], + } + ) + if result.metadata["url"]: + sources.append(result.metadata["url"]) + + # Update run usage + ctx.usage += RunUsage( + input_tokens=rag_results.usage.prompt_tokens, + output_tokens=rag_results.usage.completion_tokens, + ) + + return ToolReturn( + return_value=INSTRUCTIONS + str(compact), + content='', + metadata={"sources": list(set(sources))}, + ) + except Exception as exc: # pylint: disable=broad-except logger.exception("Albert Service Public search failed: %s", exc) return ToolReturn(return_value=[], content="", metadata={"error": str(exc)}) - # Map to a compact structure that the model can consume easily - data = json_response.get("data", []) - usage_obj = json_response.get("usage", {}) - - compact = [] - sources = [] - for item in data: - # Albert returns an object with fields: score, chunk{ content, metadata{ document_name } } - chunk = item.get("chunk", {}) - metadata = chunk.get("metadata", {}) - compact.append( - { - "title": metadata.get("document_name"), - "snippet": chunk.get("content"), - "url": metadata.get("url"), - } - ) - if metadata.get("document_name"): - sources.append(metadata["document_name"]) - - # Update run usage if available - if usage_obj: - try: - ctx.usage += RunUsage( - input_tokens=usage_obj.get("prompt_tokens", 0), - output_tokens=usage_obj.get("completion_tokens", 0), - ) - except Exception: # noqa: BLE001 - # Non-blocking if shape changes - pass - - return ToolReturn( - return_value=compact, - content="", # let the model consume return_value; avoid injecting in UI - metadata={"sources": sources}, - ) -