✨(RAG) add generic collection RAG tools
This allows to deploy generic RAG tools with predefined collections to allow specific document database for some users.
This commit is contained in:
Quentin BEY
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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):
|
||||
|
||||
+66
@@ -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."
|
||||
)
|
||||
@@ -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
|
||||
@@ -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
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user