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

📝(service-public): use existing RAG implementation and upgrade for metadata

Browse Source
This commit is contained in:
camilleAND
2025-10-22 16:47:13 +02:00
parent 3b664c2a44
commit dd9d760659
4 changed files with 53 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/backend/chat/agent_rag/constants.py
+1
View File
@@ -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):
src/backend/chat/agent_rag/document_rag_backends/albert_rag_backend.py
+12 -3
View File
@@ -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
],
src/backend/chat/agent_rag/web_search/albert_api.py
+1
View File
@@ -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
],
src/backend/chat/tools/service_public.py
+39 -76
View File
@@ -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},
)