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

✨(backend) implement FindRagBackend

Browse Source 
We want to be able to use Find api in rag tools.
I add a new rag backend class to do so.
This commit is contained in:
charles
2025-12-10 15:31:31 +01:00
parent 3232da72c5
commit b62fffc69d
13 changed files with 323 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
compose.yml
+11
View File
@@ -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
docs/llm-configuration.md
+2 -2
View File
@@ -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,
docs/tools/web_search_brave.md
+1
View File
@@ -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,
)
```
src/backend/chat/agent_rag/document_rag_backends/albert_rag_backend.py
+3 -2
View File
@@ -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.
src/backend/chat/agent_rag/document_rag_backends/base_rag_backend.py
+14 -4
View File
@@ -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
src/backend/chat/agent_rag/document_rag_backends/find_rag_backend.py
+202
View File
@@ -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,
),
)
src/backend/chat/clients/pydantic_ai.py
+3 -1
View File
@@ -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,
)
src/backend/chat/tools/document_search_rag.py
+1 -1
View File
@@ -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,
src/backend/chat/views.py
+5
View File
@@ -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
src/backend/conftest.py
+1 -1
View File
@@ -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):
src/backend/conversations/configuration/llm/dev.json
+43
View File
@@ -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"
}
]
}
src/backend/conversations/settings.py
+11 -1
View File
@@ -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.
src/backend/utils/oicd.py
+26
View File
@@ -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")
)