🐛(web-search) support web search using document storage
This allows to not call an LLM for summary but instead use a vectorized search to return only useful snippets.
This commit is contained in:
Quentin BEY
@@ -1,12 +1,20 @@
|
||||
"""Tests for the Brave web search tool."""
|
||||
|
||||
from unittest.mock import patch
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import pytest
|
||||
import responses
|
||||
from pydantic_ai import RunContext, RunUsage
|
||||
|
||||
from chat.tools.web_search_brave import web_search_brave
|
||||
from chat.tools.web_search_brave import (
|
||||
_extract_and_summarize_snippets,
|
||||
_fetch_and_extract,
|
||||
_fetch_and_store,
|
||||
_query_brave_api,
|
||||
web_search_brave,
|
||||
web_search_brave_with_document_backend,
|
||||
)
|
||||
|
||||
BRAVE_URL = "https://api.search.brave.com/res/v1/web/search"
|
||||
|
||||
@@ -22,6 +30,12 @@ def brave_settings(settings):
|
||||
settings.BRAVE_SEARCH_SAFE_SEARCH = "moderate"
|
||||
settings.BRAVE_SEARCH_SPELLCHECK = True
|
||||
settings.BRAVE_SEARCH_EXTRA_SNIPPETS = True
|
||||
settings.BRAVE_SUMMARIZATION_ENABLED = False
|
||||
settings.BRAVE_MAX_WORKERS = 2
|
||||
settings.BRAVE_CACHE_TTL = 3600
|
||||
settings.RAG_DOCUMENT_SEARCH_BACKEND = (
|
||||
"chat.agent_rag.document_rag_backends.albert_rag_backend.AlbertRagBackend"
|
||||
)
|
||||
|
||||
|
||||
@responses.activate
|
||||
@@ -216,3 +230,348 @@ def test_agent_web_search_brave_params_exclude_none(settings):
|
||||
|
||||
# Empty body for GET request
|
||||
assert not brave_request.body
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_agent_web_search_brave_parallel_processing(settings):
|
||||
"""Test parallel processing with multiple workers."""
|
||||
settings.BRAVE_MAX_WORKERS = 2
|
||||
|
||||
responses.add(
|
||||
responses.GET,
|
||||
BRAVE_URL,
|
||||
json={
|
||||
"web": {
|
||||
"results": [
|
||||
{"url": "https://example.com/1", "title": "Result 1"},
|
||||
{"url": "https://example.com/2", "title": "Result 2"},
|
||||
]
|
||||
}
|
||||
},
|
||||
status=200,
|
||||
)
|
||||
|
||||
with patch("chat.tools.web_search_brave.fetch_url") as mock_fetch:
|
||||
mock_fetch.return_value = "<html><body>Content</body></html>"
|
||||
tool_return = web_search_brave("parallel query")
|
||||
|
||||
assert len(tool_return.return_value) == 2
|
||||
assert mock_fetch.call_count == 2
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_agent_web_search_brave_single_worker(settings):
|
||||
"""Test processing with single worker (no ThreadPoolExecutor overhead)."""
|
||||
settings.BRAVE_MAX_WORKERS = 1
|
||||
|
||||
responses.add(
|
||||
responses.GET,
|
||||
BRAVE_URL,
|
||||
json={
|
||||
"web": {
|
||||
"results": [
|
||||
{"url": "https://example.com/single", "title": "Single Result"},
|
||||
]
|
||||
}
|
||||
},
|
||||
status=200,
|
||||
)
|
||||
|
||||
with patch("chat.tools.web_search_brave.fetch_url") as mock_fetch:
|
||||
mock_fetch.return_value = "<html><body>Single Content</body></html>"
|
||||
tool_return = web_search_brave("single worker query")
|
||||
|
||||
assert len(tool_return.return_value) == 1
|
||||
assert tool_return.return_value[0]["extra_snippets"] == ["Single Content"]
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_fetch_and_extract_with_cache():
|
||||
"""Test caching mechanism in _fetch_and_extract."""
|
||||
with patch("chat.tools.web_search_brave.cache") as mock_cache:
|
||||
with patch("chat.tools.web_search_brave.fetch_url") as mock_fetch:
|
||||
# First call - cache miss
|
||||
mock_cache.get.return_value = None
|
||||
mock_fetch.return_value = "<html><body>Cached Content</body></html>"
|
||||
|
||||
result1 = _fetch_and_extract("https://example.com/cache")
|
||||
|
||||
assert result1 == "Cached Content"
|
||||
mock_cache.get.assert_called_once()
|
||||
mock_cache.set.assert_called_once()
|
||||
|
||||
# Second call - cache hit
|
||||
mock_cache.get.return_value = "Cached Content"
|
||||
result2 = _fetch_and_extract("https://example.com/cache")
|
||||
|
||||
assert result2 == "Cached Content"
|
||||
# fetch_url should still be called only once (from first call)
|
||||
assert mock_fetch.call_count == 1
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_extract_and_summarize_snippets_empty_document():
|
||||
"""Test _extract_and_summarize_snippets when extraction returns empty string."""
|
||||
with patch("chat.tools.web_search_brave.fetch_url") as mock_fetch:
|
||||
with patch("chat.tools.web_search_brave.extract") as mock_extract:
|
||||
mock_fetch.return_value = "<html><body></body></html>"
|
||||
mock_extract.return_value = ""
|
||||
|
||||
result = _extract_and_summarize_snippets("query", "https://example.com/empty")
|
||||
|
||||
assert not result
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_extract_and_summarize_snippets_summarization_failure(settings):
|
||||
"""Test _extract_and_summarize_snippets when summarization fails."""
|
||||
settings.BRAVE_SUMMARIZATION_ENABLED = True
|
||||
|
||||
with patch("chat.tools.web_search_brave.fetch_url") as mock_fetch:
|
||||
with patch("chat.tools.web_search_brave.llm_summarize") as mock_summarize:
|
||||
mock_fetch.return_value = "<html><body>Content</body></html>"
|
||||
mock_summarize.side_effect = Exception("Summarization error")
|
||||
|
||||
result = _extract_and_summarize_snippets("query", "https://example.com/error")
|
||||
|
||||
assert not result
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_web_search_brave_with_document_backend_success():
|
||||
"""Test web_search_brave_with_document_backend with successful RAG search."""
|
||||
responses.add(
|
||||
responses.GET,
|
||||
BRAVE_URL,
|
||||
json={
|
||||
"web": {
|
||||
"results": [
|
||||
{"url": "https://example.com/doc1", "title": "Document 1"},
|
||||
{"url": "https://example.com/doc2", "title": "Document 2"},
|
||||
]
|
||||
}
|
||||
},
|
||||
status=200,
|
||||
)
|
||||
|
||||
mock_ctx = Mock(spec=RunContext)
|
||||
mock_ctx.usage = RunUsage(input_tokens=0, output_tokens=0)
|
||||
|
||||
mock_document_store = MagicMock()
|
||||
mock_rag_result1 = Mock(url="https://example.com/doc1", content="RAG Content 1")
|
||||
mock_rag_result2 = Mock(url="https://example.com/doc2", content="RAG Content 2")
|
||||
mock_rag_results = Mock(
|
||||
data=[mock_rag_result1, mock_rag_result2], usage=Mock(prompt_tokens=10, completion_tokens=5)
|
||||
)
|
||||
mock_document_store.search.return_value = mock_rag_results
|
||||
|
||||
mock_backend_class = MagicMock()
|
||||
mock_backend_class.temporary_collection.return_value.__enter__.return_value = (
|
||||
mock_document_store
|
||||
)
|
||||
|
||||
with patch("chat.tools.web_search_brave.import_string", return_value=mock_backend_class):
|
||||
with patch("chat.tools.web_search_brave.fetch_url") as mock_fetch:
|
||||
mock_fetch.return_value = "<html><body>Document content</body></html>"
|
||||
|
||||
tool_return = web_search_brave_with_document_backend(mock_ctx, "rag query")
|
||||
|
||||
assert len(tool_return.return_value) == 2
|
||||
assert tool_return.return_value[0]["link"] == "https://example.com/doc1"
|
||||
assert tool_return.return_value[0]["extra_snippets"] == ["RAG Content 1"]
|
||||
assert tool_return.return_value[1]["link"] == "https://example.com/doc2"
|
||||
assert tool_return.return_value[1]["extra_snippets"] == ["RAG Content 2"]
|
||||
assert tool_return.metadata["sources"] == {
|
||||
"https://example.com/doc1",
|
||||
"https://example.com/doc2",
|
||||
}
|
||||
|
||||
# Verify usage was updated
|
||||
assert mock_ctx.usage.input_tokens == 10
|
||||
assert mock_ctx.usage.output_tokens == 5
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_web_search_brave_with_document_backend_single_worker(settings):
|
||||
"""Test web_search_brave_with_document_backend with single worker."""
|
||||
settings.BRAVE_MAX_WORKERS = 1
|
||||
|
||||
responses.add(
|
||||
responses.GET,
|
||||
BRAVE_URL,
|
||||
json={
|
||||
"web": {
|
||||
"results": [
|
||||
{"url": "https://example.com/single", "title": "Single Doc"},
|
||||
]
|
||||
}
|
||||
},
|
||||
status=200,
|
||||
)
|
||||
|
||||
mock_ctx = Mock(spec=RunContext)
|
||||
mock_ctx.usage = RunUsage(input_tokens=0, output_tokens=0)
|
||||
|
||||
mock_document_store = MagicMock()
|
||||
mock_rag_result = Mock(url="https://example.com/single", content="Single Content")
|
||||
mock_rag_results = Mock(
|
||||
data=[mock_rag_result], usage=Mock(prompt_tokens=5, completion_tokens=3)
|
||||
)
|
||||
mock_document_store.search.return_value = mock_rag_results
|
||||
|
||||
mock_backend_class = MagicMock()
|
||||
mock_backend_class.temporary_collection.return_value.__enter__.return_value = (
|
||||
mock_document_store
|
||||
)
|
||||
|
||||
with patch("chat.tools.web_search_brave.import_string", return_value=mock_backend_class):
|
||||
with patch("chat.tools.web_search_brave._fetch_and_store") as mock_store:
|
||||
tool_return = web_search_brave_with_document_backend(mock_ctx, "single query")
|
||||
|
||||
assert len(tool_return.return_value) == 1
|
||||
mock_store.assert_called_once()
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_web_search_brave_with_document_backend_fetch_error(settings):
|
||||
"""Test web_search_brave_with_document_backend when document fetching fails (multi-worker)."""
|
||||
settings.BRAVE_MAX_WORKERS = 2
|
||||
|
||||
responses.add(
|
||||
responses.GET,
|
||||
BRAVE_URL,
|
||||
json={
|
||||
"web": {
|
||||
"results": [
|
||||
{"url": "https://example.com/error", "title": "Error Doc"},
|
||||
{"url": "https://example.com/ok", "title": "OK Doc"},
|
||||
]
|
||||
}
|
||||
},
|
||||
status=200,
|
||||
)
|
||||
|
||||
mock_ctx = Mock(spec=RunContext)
|
||||
mock_ctx.usage = RunUsage(input_tokens=0, output_tokens=0)
|
||||
|
||||
mock_document_store = MagicMock()
|
||||
mock_rag_results = Mock(data=[], usage=Mock(prompt_tokens=0, completion_tokens=0))
|
||||
mock_document_store.search.return_value = mock_rag_results
|
||||
|
||||
mock_backend_class = MagicMock()
|
||||
mock_backend_class.temporary_collection.return_value.__enter__.return_value = (
|
||||
mock_document_store
|
||||
)
|
||||
|
||||
with patch("chat.tools.web_search_brave.import_string", return_value=mock_backend_class):
|
||||
with patch("chat.tools.web_search_brave._fetch_and_store") as mock_store:
|
||||
# First call fails, second succeeds
|
||||
mock_store.side_effect = [Exception("Fetch error"), None]
|
||||
|
||||
tool_return = web_search_brave_with_document_backend(mock_ctx, "error query")
|
||||
|
||||
# Should complete despite error (error is caught and logged in multi-worker path)
|
||||
assert tool_return.return_value == []
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_web_search_brave_with_document_backend_no_matching_rag_results():
|
||||
"""
|
||||
Test when RAG returns results that don't match any search results.
|
||||
|
||||
This is actually a problematic scenario, but we want to ensure graceful handling.
|
||||
"""
|
||||
responses.add(
|
||||
responses.GET,
|
||||
BRAVE_URL,
|
||||
json={
|
||||
"web": {
|
||||
"results": [
|
||||
{"url": "https://example.com/doc1", "title": "Document 1"},
|
||||
]
|
||||
}
|
||||
},
|
||||
status=200,
|
||||
)
|
||||
|
||||
mock_ctx = Mock(spec=RunContext)
|
||||
mock_ctx.usage = RunUsage(input_tokens=0, output_tokens=0)
|
||||
|
||||
mock_document_store = MagicMock()
|
||||
# RAG result with different URL
|
||||
mock_rag_result = Mock(url="https://different.com/doc", content="Different Content")
|
||||
mock_rag_results = Mock(
|
||||
data=[mock_rag_result], usage=Mock(prompt_tokens=5, completion_tokens=3)
|
||||
)
|
||||
mock_document_store.search.return_value = mock_rag_results
|
||||
|
||||
mock_backend_class = MagicMock()
|
||||
mock_backend_class.temporary_collection.return_value.__enter__.return_value = (
|
||||
mock_document_store
|
||||
)
|
||||
|
||||
with patch("chat.tools.web_search_brave.import_string", return_value=mock_backend_class):
|
||||
with patch("chat.tools.web_search_brave.fetch_url") as mock_fetch:
|
||||
mock_fetch.return_value = "<html><body>Content</body></html>"
|
||||
|
||||
tool_return = web_search_brave_with_document_backend(mock_ctx, "query")
|
||||
|
||||
# No results should be returned since RAG URL doesn't match search results
|
||||
assert tool_return.return_value == []
|
||||
assert tool_return.metadata["sources"] == set()
|
||||
|
||||
|
||||
def test_fetch_and_store():
|
||||
"""Test _fetch_and_store function."""
|
||||
mock_document_store = MagicMock()
|
||||
|
||||
with patch("chat.tools.web_search_brave._fetch_and_extract") as mock_extract:
|
||||
mock_extract.return_value = "Extracted document content"
|
||||
|
||||
_fetch_and_store("https://example.com/doc", mock_document_store)
|
||||
|
||||
mock_extract.assert_called_once_with("https://example.com/doc")
|
||||
mock_document_store.store_document.assert_called_once_with(
|
||||
"https://example.com/doc", "Extracted document content"
|
||||
)
|
||||
|
||||
|
||||
def test_fetch_and_store_empty_document():
|
||||
"""Test _fetch_and_store when extraction returns empty document."""
|
||||
mock_document_store = MagicMock()
|
||||
|
||||
with patch("chat.tools.web_search_brave._fetch_and_extract") as mock_extract:
|
||||
mock_extract.return_value = ""
|
||||
|
||||
_fetch_and_store("https://example.com/empty", mock_document_store)
|
||||
|
||||
# Should not store empty document
|
||||
mock_document_store.store_document.assert_not_called()
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_query_brave_api_missing_web_key():
|
||||
"""Test _query_brave_api when response doesn't contain 'web' key."""
|
||||
responses.add(
|
||||
responses.GET,
|
||||
BRAVE_URL,
|
||||
json={"error": "no web results"},
|
||||
status=200,
|
||||
)
|
||||
|
||||
result = _query_brave_api("query")
|
||||
assert not result
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_query_brave_api_missing_results_key():
|
||||
"""Test _query_brave_api when 'web' exists but 'results' is missing."""
|
||||
responses.add(
|
||||
responses.GET,
|
||||
BRAVE_URL,
|
||||
json={"web": {}},
|
||||
status=200,
|
||||
)
|
||||
|
||||
result = _query_brave_api("query")
|
||||
assert not result
|
||||
|
||||
@@ -4,7 +4,7 @@ from pydantic_ai import Tool, ToolDefinition
|
||||
|
||||
from .fake_current_weather import get_current_weather
|
||||
from .web_seach_albert_rag import web_search_albert_rag
|
||||
from .web_search_brave import web_search_brave
|
||||
from .web_search_brave import web_search_brave, web_search_brave_with_document_backend
|
||||
from .web_search_tavily import web_search_tavily
|
||||
|
||||
|
||||
@@ -20,6 +20,11 @@ def get_pydantic_tools_by_name(name: str) -> Tool:
|
||||
"web_search_brave": Tool(
|
||||
web_search_brave, takes_ctx=False, prepare=only_if_web_search_enabled
|
||||
),
|
||||
"web_search_brave_with_document_backend": Tool(
|
||||
web_search_brave_with_document_backend,
|
||||
takes_ctx=True,
|
||||
prepare=only_if_web_search_enabled,
|
||||
),
|
||||
"web_search_tavily": Tool(
|
||||
web_search_tavily, takes_ctx=False, prepare=only_if_web_search_enabled
|
||||
),
|
||||
|
||||
@@ -1,13 +1,16 @@
|
||||
"""Web search tool using Brave for the chat agent."""
|
||||
|
||||
import logging
|
||||
import uuid
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from typing import List
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.cache import cache
|
||||
from django.utils.module_loading import import_string
|
||||
|
||||
import requests
|
||||
from pydantic_ai import RunContext, RunUsage
|
||||
from pydantic_ai.messages import ToolReturn
|
||||
from trafilatura import extract, fetch_url
|
||||
from trafilatura.meta import reset_caches
|
||||
@@ -81,6 +84,13 @@ def _extract_and_summarize_snippets(query: str, url: str) -> List[str]:
|
||||
return [snippet] if snippet else []
|
||||
|
||||
|
||||
def _fetch_and_store(url: str, document_store) -> None:
|
||||
"""Fetch, extract and store text content from the URL in the document store."""
|
||||
document = _fetch_and_extract(url)
|
||||
if document:
|
||||
document_store.store_document(url, document)
|
||||
|
||||
|
||||
def _query_brave_api(query: str) -> List[dict]:
|
||||
"""Query the Brave Search API and return the raw results."""
|
||||
url = "https://api.search.brave.com/res/v1/web/search"
|
||||
@@ -155,3 +165,68 @@ def web_search_brave(query: str) -> ToolReturn:
|
||||
],
|
||||
metadata={"sources": {result["url"] for result in raw_search_results}},
|
||||
)
|
||||
|
||||
|
||||
def web_search_brave_with_document_backend(ctx: RunContext, query: str) -> ToolReturn:
|
||||
"""
|
||||
Search the web for up-to-date information
|
||||
|
||||
Args:
|
||||
ctx (RunContext): The run context containing the conversation.
|
||||
query (str): The query to search for.
|
||||
"""
|
||||
raw_search_results = _query_brave_api(query)
|
||||
|
||||
reset_caches() # Clear trafilatura caches to avoid memory bloat/leaks
|
||||
|
||||
# Store documents in a temporary document store for RAG search
|
||||
document_store_backend = import_string(settings.RAG_DOCUMENT_SEARCH_BACKEND)
|
||||
with document_store_backend.temporary_collection(f"tmp-{uuid.uuid4()}") as document_store:
|
||||
max_workers = min(settings.BRAVE_MAX_WORKERS, len(raw_search_results))
|
||||
if max_workers == 1:
|
||||
for result in raw_search_results:
|
||||
# Fetch and extract document content
|
||||
_fetch_and_store(result["url"], document_store)
|
||||
else:
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
futures = [
|
||||
executor.submit(_fetch_and_store, result["url"], document_store)
|
||||
for result in raw_search_results
|
||||
]
|
||||
for future in as_completed(futures):
|
||||
try:
|
||||
future.result()
|
||||
except Exception as e: # pylint: disable=broad-except
|
||||
logger.exception("Error fetching/storing document: %s", e)
|
||||
|
||||
rag_results = document_store.search(query)
|
||||
|
||||
ctx.usage += RunUsage(
|
||||
input_tokens=rag_results.usage.prompt_tokens,
|
||||
output_tokens=rag_results.usage.completion_tokens,
|
||||
)
|
||||
|
||||
# Map RAG results back to raw search results to include extra_snippets
|
||||
# Suboptimal O(N^2) but N is small...
|
||||
for rag_result in rag_results.data:
|
||||
for result in raw_search_results:
|
||||
if result["url"] == rag_result.url:
|
||||
result.setdefault("extra_snippets", []).append(rag_result.content)
|
||||
break
|
||||
|
||||
return ToolReturn(
|
||||
return_value=[
|
||||
{
|
||||
"link": result["url"],
|
||||
"title": result["title"],
|
||||
"extra_snippets": result.get("extra_snippets", []),
|
||||
}
|
||||
for result in raw_search_results
|
||||
if result.get("extra_snippets", [])
|
||||
],
|
||||
metadata={
|
||||
"sources": {
|
||||
result["url"] for result in raw_search_results if result.get("extra_snippets", [])
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user