From dcec57719f9206a01de862a989783f56898cb32c Mon Sep 17 00:00:00 2001 From: Quentin BEY Date: Wed, 12 Nov 2025 23:40:16 +0100 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8(RAG)=20add=20generic=20collection=20R?= =?UTF-8?q?AG=20tools?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This allows to deploy generic RAG tools with predefined collections to allow specific document database for some users. --- CHANGELOG.md | 5 + .../albert_rag_backend.py | 18 +- .../document_rag_backends/base_rag_backend.py | 46 +- src/backend/chat/clients/pydantic_ai.py | 2 + ...d_document_rag_search_tool_from_setting.py | 66 +++ .../tools/test_document_generic_search_rag.py | 403 ++++++++++++++++++ .../chat/tools/document_generic_search_rag.py | 137 ++++++ src/backend/conftest.py | 15 + src/backend/conversations/settings.py | 5 + src/backend/core/feature_flags/flags.py | 7 + 10 files changed, 696 insertions(+), 8 deletions(-) create mode 100644 src/backend/chat/tests/clients/pydantic_ai/test_add_document_rag_search_tool_from_setting.py create mode 100644 src/backend/chat/tests/tools/test_document_generic_search_rag.py create mode 100644 src/backend/chat/tools/document_generic_search_rag.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 7c0fac4..9f4208b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,11 @@ and this project adheres to ## [Unreleased] +### Added + +- ✨(RAG) add generic collection RAG tools #159 + + ## [0.0.8] - 2025-11-10 ### Fixed 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 aa5bf68..8af560b 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 @@ -33,9 +33,13 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att - Perform a search operation using the Albert API. """ - def __init__(self, collection_id: Optional[str] = None): + 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) + super().__init__(collection_id, read_only_collection_id) self._base_url = settings.ALBERT_API_URL self._headers = { "Authorization": f"Bearer {settings.ALBERT_API_KEY}", @@ -220,11 +224,13 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att Returns: RAGWebResults: The search results. """ + collection_ids = self.get_all_collection_ids() # might raise RuntimeError + response = requests.post( urljoin(self._base_url, self._search_endpoint), headers=self._headers, json={ - "collections": [int(self.collection_id)], + "collections": collection_ids, "prompt": query, "score_threshold": 0.6, "k": results_count, # Number of chunks to return from the search @@ -261,12 +267,14 @@ class AlbertRagBackend(BaseRagBackend): # pylint: disable=too-many-instance-att Returns: RAGWebResults: The search results. """ + collection_ids = self.get_all_collection_ids() # might raise RuntimeError + async with httpx.AsyncClient(timeout=settings.ALBERT_API_TIMEOUT) as client: response = await client.post( urljoin(self._base_url, self._search_endpoint), headers=self._headers, json={ - "collections": [int(self.collection_id)], + "collections": collection_ids, "prompt": query, "score_threshold": 0.6, "k": results_count, # Number of chunks to return from the search 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 0ca4fa9..78d25cd 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 @@ -3,7 +3,7 @@ import logging from contextlib import asynccontextmanager, contextmanager from io import BytesIO -from typing import Optional +from typing import List, Optional from asgiref.sync import sync_to_async @@ -15,11 +15,51 @@ logger = logging.getLogger(__name__) class BaseRagBackend: """Base class for RAG backends.""" - def __init__(self, collection_id: Optional[str] = None): - """Backend settings.""" + def __init__( + self, + collection_id: Optional[str] = None, + read_only_collection_id: Optional[List[str]] = None, + ): + """ + Backend settings. + + Collection ID is required for RAG operations, where you want to manage the collection + lifecycle (create/delete). + Read-only collection IDs can be used to access existing collections + without managing their lifecycle. + + Collection ID and read-only collection IDs are separated in the implementation to prevent + unwanted actions. + + Args: + collection_id (Optional[str]): The collection ID for managing the collection lifecycle. + read_only_collection_id (Optional[List[str]]): List of read-only collection IDs. + """ self.collection_id = collection_id + self.read_only_collection_id = read_only_collection_id or [] self._default_collection_description = "Temporary collection for RAG document search" + def get_all_collection_ids(self) -> List[str]: + """ + Get all collection IDs, including the main collection ID and read-only collection IDs. + + Returns: + List[str]: List of all collection IDs. + Raises: + RuntimeError: If neither collection_id nor read_only_collection_id is provided. + """ + if not self.collection_id and not self.read_only_collection_id: + raise RuntimeError("The RAG backend requires collection_id or read_only_collection_id") + + collection_ids = [] + if self.collection_id: + collection_ids.append(int(self.collection_id)) + if self.read_only_collection_id: + collection_ids.extend( + [int(collection_id) for collection_id in self.read_only_collection_id] + ) + return collection_ids + def create_collection(self, name: str, description: Optional[str] = None) -> str: """ Create a temporary collection for the search operation. diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py index ebc8d94..60d8fd5 100644 --- a/src/backend/chat/clients/pydantic_ai.py +++ b/src/backend/chat/clients/pydantic_ai.py @@ -72,6 +72,7 @@ from chat.clients.pydantic_ui_message_converter import ( ui_message_to_user_content, ) from chat.mcp_servers import get_mcp_servers +from chat.tools.document_generic_search_rag import add_document_rag_search_tool_from_setting from chat.tools.document_search_rag import add_document_rag_search_tool from chat.tools.document_summarize import document_summarize from chat.vercel_ai_sdk.core import events_v4, events_v5 @@ -140,6 +141,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes instrument=self._store_analytics, deps_type=ContextDeps, ) + add_document_rag_search_tool_from_setting(self.conversation_agent, self.user) @property def _stop_cache_key(self): diff --git a/src/backend/chat/tests/clients/pydantic_ai/test_add_document_rag_search_tool_from_setting.py b/src/backend/chat/tests/clients/pydantic_ai/test_add_document_rag_search_tool_from_setting.py new file mode 100644 index 0000000..7ded789 --- /dev/null +++ b/src/backend/chat/tests/clients/pydantic_ai/test_add_document_rag_search_tool_from_setting.py @@ -0,0 +1,66 @@ +"""Unit tests for add_document_rag_search_tool_from_setting integration with AIAgentService.""" + +import pytest + +from core.factories import UserFactory + +from chat.clients.pydantic_ai import AIAgentService +from chat.factories import ChatConversationFactory +from chat.llm_configuration import LLModel, LLMProvider + +pytestmark = pytest.mark.django_db() + + +def test_ai_agent_service_adds_rag_tools_from_settings(settings): + """Test that AIAgentService adds RAG tools from SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS.""" + settings.LLM_CONFIGURATIONS = { + "default-model": LLModel( + hrid="default-model", + model_name="amazing-llm", + human_readable_name="Amazing LLM", + is_active=True, + icon=None, + system_prompt="You are an amazing assistant.", + tools=[], + provider=LLMProvider(hrid="unused", base_url="https://example.com", api_key="key"), + ), + } + settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = { + "legal_documents": { + "collection_ids": [100, 101, 102], + "feature_flag_value": "enabled", + "tool_description": "Use this tool to search legal documents and laws.", + }, + "french_public_services": { + "collection_ids": [784, 785], + "feature_flag_value": "enabled", + "tool_description": ( + "Use this tool when the user asks for information about French public services." + ), + }, + } + + user = UserFactory() + conversation = ChatConversationFactory(owner=user) + + # Create the service + service = AIAgentService(conversation, user=user) + + # Check that tools were added to the conversation_agent + agent_tools = service.conversation_agent._function_toolset.tools # pylint: disable=protected-access + + assert "legal_documents" in agent_tools + assert "french_public_services" in agent_tools + + # Verify tool names and descriptions + assert agent_tools["legal_documents"].name == "legal_documents" + assert ( + agent_tools["legal_documents"].description + == "Use this tool to search legal documents and laws." + ) + + assert agent_tools["french_public_services"].name == "french_public_services" + assert ( + agent_tools["french_public_services"].description + == "Use this tool when the user asks for information about French public services." + ) diff --git a/src/backend/chat/tests/tools/test_document_generic_search_rag.py b/src/backend/chat/tests/tools/test_document_generic_search_rag.py new file mode 100644 index 0000000..3cea7a4 --- /dev/null +++ b/src/backend/chat/tests/tools/test_document_generic_search_rag.py @@ -0,0 +1,403 @@ +""" +Unit tests for document generic search RAG tool functionality. +""" + +import json +import logging + +import httpx +import pytest +import responses +import respx +from asgiref.sync import sync_to_async +from pydantic_ai import Agent, RunContext, RunUsage + +from core.factories import UserFactory + +from chat.tools.document_generic_search_rag import ( + add_document_rag_search_tool_from_setting, + get_specific_rag_search_tool_config, +) + +pytestmark = pytest.mark.django_db() + + +def test_get_specific_rag_search_tool_config_with_disabled_features(settings): + """Test get_specific_rag_search_tool_config returns tools for enabled features.""" + user = UserFactory() + + settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = { + "french_public_services": { + "collection_ids": [784, 785], + "feature_flag_value": "disabled", + "tool_description": ( + "Use this tool when the user asks for information about French public services, " + "the French labor market, employment laws, social benefits, or " + "assistance with administrative procedures." + ), + }, + "legal_documents": { + "collection_ids": [100, 101, 102], + "feature_flag_value": "disabled", + "tool_description": "Use this tool to search French legal documents and laws.", + "rag_backend_name": "chat.tests.tools.test_document_generic_search_rag.MockRagBackend", + }, + } + + # The fixture tools are disabled by default + assert get_specific_rag_search_tool_config(user) == {} + + +def test_get_specific_rag_search_tool_config_with_enabled_features(settings): + """Test get_specific_rag_search_tool_config returns tools for enabled features.""" + user = UserFactory() + + settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = { + "french_public_services": { + "collection_ids": [784, 785], + "feature_flag_value": "enabled", + "tool_description": ( + "Use this tool when the user asks for information about French public services, " + "the French labor market, employment laws, social benefits, or " + "assistance with administrative procedures." + ), + }, + "legal_documents": { + "collection_ids": [100, 101, 102], + "feature_flag_value": "enabled", + "tool_description": "Use this tool to search French legal documents and laws.", + }, + } + + assert get_specific_rag_search_tool_config(user) == { + "french_public_services": { + "collection_ids": [784, 785], + "feature_flag_value": "enabled", + "tool_description": "Use this tool when the user " + "asks for information about " + "French public services, the " + "French labor market, " + "employment laws, social " + "benefits, or assistance with " + "administrative procedures.", + }, + "legal_documents": { + "collection_ids": [100, 101, 102], + "feature_flag_value": "enabled", + "tool_description": "Use this tool to search French legal documents and laws.", + }, + } + + +@responses.activate +def test_get_specific_rag_search_tool_config_with_dynamic_features(settings, posthog): + """Test get_specific_rag_search_tool_config with dynamic features.""" + user = UserFactory() + + responses.post( + f"{posthog.host}/flags/?v=2", + json={"flags": {"legal-documents": {"enabled": True}}}, + status=200, + ) + + settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = { + "french_public_services": { + "collection_ids": [784, 785], + "feature_flag_value": "dynamic", + "tool_description": ( + "Use this tool when the user asks for information about French public services, " + "the French labor market, employment laws, social benefits, or " + "assistance with administrative procedures." + ), + }, + "legal_documents": { + "collection_ids": [100, 101, 102], + "feature_flag_value": "dynamic", + "tool_description": "Use this tool to search French legal documents and laws.", + }, + } + + assert get_specific_rag_search_tool_config(user) == { + "legal_documents": { + "collection_ids": [100, 101, 102], + "feature_flag_value": "dynamic", + "tool_description": "Use this tool to search French legal documents and laws.", + } + } + + +def test_add_document_rag_search_tool_from_setting_adds_tools(settings): + """Test that add_document_rag_search_tool_from_setting adds tools to the agent.""" + settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = { + "legal_documents": { + "collection_ids": [100, 101, 102], + "feature_flag_value": "enabled", + "tool_description": "Use this tool to search legal documents and laws.", + }, + } + + user = UserFactory() + + agent = Agent("test") + assert len(agent._function_toolset.tools) == 0 # pylint: disable=protected-access + + add_document_rag_search_tool_from_setting(agent, user) + + # Check that tools were added + assert len(agent._function_toolset.tools) == 1 # pylint: disable=protected-access + assert agent._function_toolset.tools["legal_documents"].name == "legal_documents" # pylint: disable=protected-access + assert ( + agent._function_toolset.tools["legal_documents"].description # pylint: disable=protected-access + == "Use this tool to search legal documents and laws." + ) + assert agent._function_toolset.tools["legal_documents"].function_schema.json_schema == { # pylint: disable=protected-access + "additionalProperties": False, + "properties": { + "query": {"description": "The query to search information about.", "type": "string"} + }, + "required": ["query"], + "type": "object", + } + + +def test_add_document_rag_search_tool_with_invalid_backend(settings, caplog): + """Test that invalid backend import is handled gracefully.""" + caplog.set_level(logging.WARNING, logger="chat.tools.document_generic_search_rag") + + settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = { + "legal_documents": { + "collection_ids": [100, 101, 102], + "rag_backend_name": "non.existent.Backend", + "feature_flag_value": "enabled", + "tool_description": "Use this tool to search legal documents and laws.", + }, + } + user = UserFactory() + agent = Agent("test") + + add_document_rag_search_tool_from_setting(agent, user) + + # Tool should not be added due to import error + assert len(agent._function_toolset.tools) == 0 # pylint: disable=protected-access + + # Check that warning was logged + assert len(caplog.records) == 1 + assert "Could not import RAG backend non.existent.Backend" in caplog.records[0].message + + +def test_add_document_rag_search_tool_with_missing_collection_ids(settings, caplog): + """Test that missing collection_ids is handled gracefully.""" + caplog.set_level(logging.WARNING, logger="chat.tools.document_generic_search_rag") + + settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = { + "legal_documents": { + "feature_flag_value": "enabled", + "tool_description": "Use this tool to search legal documents and laws.", + }, + } + user = UserFactory() + agent = Agent("test") + + add_document_rag_search_tool_from_setting(agent, user) + + # Tool should not be added due to import error + assert len(agent._function_toolset.tools) == 0 # pylint: disable=protected-access + + # Check that warning was logged + assert len(caplog.records) == 1 + assert "No collection IDs provided for tool legal_documents" in caplog.records[0].message + + +def test_add_document_rag_search_tool_with_missing_tool_description(settings, caplog): + """Test that missing tool_description is handled gracefully.""" + caplog.set_level(logging.WARNING, logger="chat.tools.document_generic_search_rag") + + settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = { + "legal_documents": { + "collection_ids": [100, 101, 102], + "feature_flag_value": "enabled", + }, + } + user = UserFactory() + agent = Agent("test") + + add_document_rag_search_tool_from_setting(agent, user) + + # Tool should not be added due to import error + assert len(agent._function_toolset.tools) == 0 # pylint: disable=protected-access + + # Check that warning was logged + assert len(caplog.records) == 1 + assert "No tool description provided for tool legal_documents" in caplog.records[0].message + + +@respx.mock +def test_document_search_rag_tool_execution(settings): + """Test that the generated RAG tool executes correctly.""" + search_mock = respx.post("https://albert.api.etalab.gouv.fr/v1/search").mock( + return_value=httpx.Response( + status_code=200, + json={ + "data": [ + { + "method": "semantic", + "chunk": { + "id": 1, + "content": "Relevant content snippet.", + "metadata": {"document_name": "doc1.txt"}, + }, + "score": 0.9, + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 20}, + }, + ) + ) + settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = { + "legal_documents": { + "collection_ids": [100, 101, 102], + "feature_flag_value": "enabled", + "tool_description": "Use this tool to search legal documents and laws.", + }, + "legal_documents_2": { + "collection_ids": [200], + "feature_flag_value": "enabled", + "tool_description": "Use this tool to search legal documents and laws.", + }, + } + user = UserFactory() + agent = Agent(model="test") + + add_document_rag_search_tool_from_setting(agent, user) + + result = agent.run_sync("What information can you find about French services?") + + # Verify the result + assert json.loads(result.output) == { + "legal_documents": {"0": {"snippets": "Relevant content snippet.", "url": "doc1.txt"}}, + "legal_documents_2": {"0": {"snippets": "Relevant content snippet.", "url": "doc1.txt"}}, + } + + assert len(search_mock.calls) == 2 + assert json.loads(search_mock.calls[0].request.content) == { + "collections": [100, 101, 102], + "k": 4, + "prompt": "a", + "score_threshold": 0.6, + } + assert json.loads(search_mock.calls[1].request.content) == { + "collections": [200], + "k": 4, + "prompt": "a", + "score_threshold": 0.6, + } + + +def test_get_specific_rag_search_tool_config_with_empty_settings(settings): + """Test get_specific_rag_search_tool_config with empty SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS.""" + settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = {} + + user = UserFactory() + config = get_specific_rag_search_tool_config(user) + + assert config == {} + + +@pytest.mark.asyncio +@respx.mock +async def test_add_document_rag_search_tool_function_call(settings): + """Test the function behavior.""" + search_mock = respx.post("https://albert.api.etalab.gouv.fr/v1/search").mock( + return_value=httpx.Response( + status_code=200, + json={ + "data": [ + { + "method": "semantic", + "chunk": { + "id": 1, + "content": "Relevant content snippet.", + "metadata": {"document_name": "doc1.txt"}, + }, + "score": 0.9, + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 20}, + }, + ) + ) + settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = { + "legal_documents": { + "collection_ids": [100, 101, 102], + "feature_flag_value": "enabled", + "tool_description": "Use this tool to search legal documents and laws.", + }, + } + + user = await sync_to_async(UserFactory)() + + agent = Agent("test") + add_document_rag_search_tool_from_setting(agent, user) + + result = await agent._function_toolset.tools["legal_documents"].function( # pylint: disable=protected-access + RunContext(model="test", usage=RunUsage(), deps={}), + query="Find information about French laws.", + ) + + assert result.return_value == { + "0": {"snippets": "Relevant content snippet.", "url": "doc1.txt"} + } + assert result.metadata == {"sources": {"doc1.txt"}} + assert len(search_mock.calls) == 1 + assert json.loads(search_mock.calls[0].request.content) == { + "collections": [100, 101, 102], + "k": 4, + "prompt": "Find information about French laws.", + "score_threshold": 0.6, + } + + +@pytest.mark.asyncio +@respx.mock +async def test_document_search_rag_http_status_error(settings, caplog): + """Test that HTTPStatusError is properly handled and logged.""" + caplog.set_level(logging.ERROR, logger="chat.tools.document_generic_search_rag") + + # Mock the API to return a 500 error + respx.post("https://albert.api.etalab.gouv.fr/v1/search").mock( + return_value=httpx.Response( + status_code=500, + json={"error": "Internal server error"}, + ) + ) + + settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = { + "legal_documents": { + "collection_ids": [100, 101, 102], + "feature_flag_value": "enabled", + "tool_description": "Use this tool to search legal documents and laws.", + }, + } + + user = await sync_to_async(UserFactory)() + agent = Agent("test") + add_document_rag_search_tool_from_setting(agent, user) + + # Call the tool function and expect a ModelRetry to be raised and caught + tool_result = await agent._function_toolset.tools["legal_documents"].function( # pylint: disable=protected-access + RunContext(model="test", usage=RunUsage(), deps={}), + query="Find information about French laws.", + ) + + # Verify the exception message + assert tool_result == ( + "Document search service is currently unavailable: Server error '500 Internal " + "Server Error' for url 'https://albert.api.etalab.gouv.fr/v1/search'\n" + "For more information check: " + "https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/500 You must " + "explain this to the user and not try to answer based on your knowledge." + ) + + # Verify that error was logged + assert "RAG document search failed for tool legal_documents" in caplog.records[0].message + assert "Document search service is currently unavailable" in caplog.records[1].message diff --git a/src/backend/chat/tools/document_generic_search_rag.py b/src/backend/chat/tools/document_generic_search_rag.py new file mode 100644 index 0000000..02eed37 --- /dev/null +++ b/src/backend/chat/tools/document_generic_search_rag.py @@ -0,0 +1,137 @@ +""" +Helpers to add RAG document search tools to an agent based on settings. + +The purpose is to provide a generic way to add multiple RAG document search tools +to an agent based on configuration in settings. Each tool can target specific +document collections and have its own description. + +Our use case implies that different users might have access to different document collections, +so the tools added to the agent are also user-specific. +""" + +import logging + +from django.conf import settings +from django.contrib.auth import get_user_model +from django.utils.module_loading import import_string + +from httpx import HTTPStatusError +from pydantic_ai import Agent, ModelRetry, RunContext, RunUsage +from pydantic_ai.messages import ToolReturn + +from core.feature_flags.helpers import is_feature_enabled + +from chat.tools.utils import last_model_retry_soft_fail + +logger = logging.getLogger(__name__) + +User = get_user_model() + + +def get_specific_rag_search_tool_config(user: User) -> dict: + """ + Get the specific RAG search tool configuration from settings. + + Settings example: + SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = { + "french_public_services": { + "collection_ids": [784, 785], + "feature_flag_value": "disabled", + "tool_description": ( + "Use this tool when the user asks for information about French public services, " + "the French labor market, employment laws, social benefits, or " + "assistance with administrative procedures." + ), + }, + } + """ + return { + tool_name: tool_config + for tool_name, tool_config in settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS.items() + if is_feature_enabled(user, tool_name) + } + + +def _create_document_search_rag(agent, name, description, backend, ids): + """Factory function to create a document search RAG tool.""" + + @agent.tool( + name=name, + retries=1, + require_parameter_descriptions=True, + description=description, + ) + @last_model_retry_soft_fail + async def document_search_rag(ctx: RunContext, query: str) -> ToolReturn: + """ + Args: + ctx (RunContext): The run context containing the conversation. + query (str): The query to search information about. + """ + document_store = backend(read_only_collection_id=ids) + + try: + rag_results = await document_store.asearch(query) + except HTTPStatusError as exc: + logger.error( + "RAG document search failed for tool %s with error: %s", name, exc, exc_info=True + ) + raise ModelRetry(f"Document search service is currently unavailable: {exc}") from exc + + ctx.usage += RunUsage( + input_tokens=rag_results.usage.prompt_tokens, + output_tokens=rag_results.usage.completion_tokens, + ) + + return ToolReturn( + return_value={ + str(idx): { + "url": result.url, + "snippets": result.content, + } + for idx, result in enumerate(rag_results.data) + }, + metadata={"sources": {result.url for result in rag_results.data}}, + ) + + return document_search_rag + + +def add_document_rag_search_tool_from_setting(agent: Agent, user: User) -> None: + """ + This function takes a configuration setting and generates specific search RAG tools and add + it to the agent. + + Args: + agent (Agent): The agent to which the tool will be added. + user (User): The user for whom the tool is being added. + """ + + for tool_name, tool_config in get_specific_rag_search_tool_config(user).items(): + document_store_backend_name = tool_config.get( + "rag_backend_name", settings.RAG_DOCUMENT_SEARCH_BACKEND + ) + try: + document_store_backend = import_string(document_store_backend_name) + except ImportError as exc: + logger.warning( + "Could not import RAG backend %s: %s", + document_store_backend_name, + exc, + exc_info=True, + ) + continue # Skip if the backend is not available + + collection_ids = tool_config.get("collection_ids", []) + if not collection_ids: + logger.warning("No collection IDs provided for tool %s, skipping.", tool_name) + continue # Skip if no collection IDs are provided + + tool_description = tool_config.get("tool_description") + if not tool_description: + logger.warning("No tool description provided for tool %s, skipping.", tool_name) + continue # Skip if no tool description is provided + + _create_document_search_rag( + agent, tool_name, tool_description, document_store_backend, collection_ids + ) diff --git a/src/backend/conftest.py b/src/backend/conftest.py index 7495890..81fa0bc 100644 --- a/src/backend/conftest.py +++ b/src/backend/conftest.py @@ -1,5 +1,6 @@ """Global fixtures for the backend tests.""" +import posthog import pytest from rest_framework.test import APIClient from urllib3.connectionpool import HTTPConnectionPool @@ -41,3 +42,17 @@ def feature_flags_fixture(settings): """ settings.FEATURE_FLAGS = settings.FEATURE_FLAGS.model_copy(deep=True) yield settings.FEATURE_FLAGS + + +@pytest.fixture(name="posthog", scope="function") +def posthog_fixture(settings): + """Mock PostHog in tests to avoid real network calls.""" + settings.POSTHOG_KEY = {"id": "132456", "host": "https://eu.i.posthog-test.com"} + + posthog.api_key = settings.POSTHOG_KEY["id"] + posthog.host = settings.POSTHOG_KEY["host"] + + yield posthog + + posthog.api_key = None + posthog.host = None diff --git a/src/backend/conversations/settings.py b/src/backend/conversations/settings.py index 759dc72..682f99c 100755 --- a/src/backend/conversations/settings.py +++ b/src/backend/conversations/settings.py @@ -717,6 +717,11 @@ class Base(BraveSettings, Configuration): environ_name="RAG_DOCUMENT_SEARCH_BACKEND", environ_prefix=None, ) + SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS = values.DictValue( + default={}, + environ_name="SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS", + environ_prefix=None, + ) # Web search RAG_WEB_SEARCH_PROMPT_UPDATE = values.Value( diff --git a/src/backend/core/feature_flags/flags.py b/src/backend/core/feature_flags/flags.py index abc9f97..a255f0a 100644 --- a/src/backend/core/feature_flags/flags.py +++ b/src/backend/core/feature_flags/flags.py @@ -2,6 +2,7 @@ from enum import StrEnum +from django.conf import settings from django.utils.text import slugify from pydantic import BaseModel, ConfigDict @@ -43,3 +44,9 @@ class FeatureFlags(BaseModel): # features web_search: FeatureToggle = FeatureToggle.DISABLED document_upload: FeatureToggle = FeatureToggle.DISABLED + + def __getattr__(self, name: str): + """Dynamically get specific RAG document search tool feature flags from settings.""" + if config := settings.SPECIFIC_RAG_DOCUMENT_SEARCH_TOOLS.get(name): + return FeatureToggle[config.get("feature_flag_value", "DISABLED").upper()] + return super().__getattr__(name)