diff --git a/src/backend/chat/agents/conversation.py b/src/backend/chat/agents/conversation.py index 963efc5..e22ad57 100644 --- a/src/backend/chat/agents/conversation.py +++ b/src/backend/chat/agents/conversation.py @@ -117,20 +117,14 @@ class ConversationAgent(BaseAgent): """Dynamic instruction function to set the expected language to use.""" return f"Answer in {get_language_name(language).lower()}." if language else "" - def get_web_search_tool_name(self) -> str | None: + def is_web_search_configured(self) -> bool: """ - Get the name of the web search tool if available. + Return True when a web search backend is configured on this model. - If several are available, return the first one found. - - Warning, this says the tool is available, not that - it (the tool/feature) is enabled for the current conversation. + This does not mean web search is enabled for the current conversation + (feature flags and runtime deps still apply). """ - for toolset in self.toolsets: - for tool in toolset.tools.values(): - if tool.name.startswith("web_search_"): - return tool.name - return None + return bool(getattr(self.configuration, "web_search", None)) @dataclasses.dataclass(init=False) diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py index 0be56eb..3b3b04c 100644 --- a/src/backend/chat/clients/pydantic_ai.py +++ b/src/backend/chat/clients/pydantic_ai.py @@ -251,6 +251,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes session=session, web_search_enabled=self._is_web_search_enabled and self._is_smart_search_enabled, ) + self._web_search_tool_registered = False self.conversation_agent = ConversationAgent( model_hrid=self.model_hrid, @@ -440,8 +441,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes logger.warning("Web search is forced but the feature is disabled, ignoring.") return False - web_search_tool_name = self.conversation_agent.get_web_search_tool_name() - if not web_search_tool_name: + if not self.conversation_agent.is_web_search_configured(): logger.warning("Web search is forced but no web search tool is available, ignoring.") return False @@ -450,9 +450,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes @self.conversation_agent.instructions def force_web_search_prompt() -> str: """Dynamic system prompt function to force web search.""" - return ( - f"You must call the {web_search_tool_name} tool before answering the user request." - ) + return "You must call the web_search tool before answering the user request." return True @@ -699,6 +697,33 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes """Wrap the document_summarize tool to provide context and add the tool.""" return await document_summarize(ctx, *args, **kwargs) + def _setup_web_search_tool(self) -> None: + """Register model-specific web search tool when configured.""" + if self._web_search_tool_registered: + return + configuration = self.conversation_agent.configuration + if not getattr(configuration, "web_search", None): + return + + async def only_if_web_search_enabled(ctx, tool_def): + """Prepare function to include a tool only if web search is enabled in the context.""" + return tool_def if ctx.deps.web_search_enabled else None + + web_search_impl = import_string(configuration.web_search) + + @self.conversation_agent.tool( + name="web_search", + retries=1, + prepare=only_if_web_search_enabled, + description="Search the web for up-to-date information", + ) + @functools.wraps(web_search_impl) + async def web_search(ctx: RunContext, *args, **kwargs) -> ToolReturn: + """Wrap the web_search tool to provide context and add the tool.""" + return await web_search_impl(ctx, *args, **kwargs) + + self._web_search_tool_registered = True + async def _handle_input_documents( self, input_documents: List[BinaryContent | DocumentUrl], @@ -1016,6 +1041,7 @@ class AIAgentService: # pylint: disable=too-many-instance-attributes conversation_has_documents = doc_result.has_documents await self._agent_stop_streaming(force_cache_check=True) + self._setup_web_search_tool() self._setup_web_search(force_web_search) if await self._check_should_enable_rag(conversation_has_documents): diff --git a/src/backend/chat/llm_configuration.py b/src/backend/chat/llm_configuration.py index 1f04290..e9e25ee 100644 --- a/src/backend/chat/llm_configuration.py +++ b/src/backend/chat/llm_configuration.py @@ -125,6 +125,7 @@ class LLModel(BaseModel): supports_streaming: bool | None = None system_prompt: SettingEnvValue tools: list[str] + web_search: SettingEnvValue | None = None @field_validator("tools", mode="before") @classmethod @@ -134,6 +135,14 @@ class LLModel(BaseModel): return _get_setting_or_env_or_value(value) return value + @field_validator("web_search", mode="before") + @classmethod + def validate_web_search(cls, value: str | None) -> str | None: + """Convert web_search path if it's a setting or environment variable.""" + if isinstance(value, str): + return _get_setting_or_env_or_value(value) + return value + @model_validator(mode="after") def check_provider_or_provider_name(self) -> Self: """ diff --git a/src/backend/chat/tests/agents/test_build_conversation_agent.py b/src/backend/chat/tests/agents/test_build_conversation_agent.py index ccc16a5..d076580 100644 --- a/src/backend/chat/tests/agents/test_build_conversation_agent.py +++ b/src/backend/chat/tests/agents/test_build_conversation_agent.py @@ -3,14 +3,12 @@ # pylint:disable=protected-access import pytest -import responses from freezegun import freeze_time from pydantic_ai import Agent from pydantic_ai.models.openai import OpenAIChatModel -from pydantic_ai.models.test import TestModel from chat.agents.conversation import ConversationAgent -from chat.clients.pydantic_ai import ContextDeps +from chat.llm_configuration import LLModel, LLMProvider @pytest.fixture(autouse=True) @@ -87,47 +85,30 @@ def test_add_dynamic_system_prompt(): assert agent._instructions[2]() == "Answer in french." -def test_agent_get_web_search_tool_name(settings): - """Test the web_search_available method.""" - settings.AI_AGENT_TOOLS = ["get_current_weather", "web_search_albert_rag"] +def test_agent_is_web_search_configured(): + """Test whether web search backend is configured on the model.""" agent = ConversationAgent(model_hrid="default-model") - assert agent.get_web_search_tool_name() == "web_search_albert_rag" + assert agent.is_web_search_configured() is False - settings.AI_AGENT_TOOLS = ["get_current_weather"] + +def test_agent_is_web_search_configured_when_defined_in_model_config(settings): + """Web search is configured when LLModel.web_search is set.""" + settings.LLM_CONFIGURATIONS = { + "default-model": LLModel( + hrid="default-model", + model_name="model-123", + human_readable_name="Default Model", + is_active=True, + icon=None, + system_prompt="You are a helpful assistant", + tools=[], + web_search="chat.tools.web_search_brave.web_search_brave_llm_context", + provider=LLMProvider( + hrid="default-provider", + base_url="https://api.llm.com/v1/", + api_key="test-key", + ), + ), + } agent = ConversationAgent(model_hrid="default-model") - assert agent.get_web_search_tool_name() is None - - settings.AI_AGENT_TOOLS = ["get_current_weather", "web_search_tavily", "web_search_albert_rag"] - agent = ConversationAgent(model_hrid="default-model") - assert agent.get_web_search_tool_name() == "web_search_tavily" - - -@responses.activate -def test_web_search_tool_avalability(settings): - """Test the web search tool availability according to context.""" - responses.add( - responses.POST, - "https://api.tavily.com/search", - json={"results": []}, - status=200, - ) - context_deps = ContextDeps(conversation=None, user=None, web_search_enabled=True) - - # No tools (context allows web search, but no tool configured) - agent = ConversationAgent(model_hrid="default-model") - with agent.override(model=TestModel(), deps=context_deps): - response = agent.run_sync("What tools do you have?") - assert response.output == "success (no tool calls)" - - # Tool configured, context allows web search - settings.AI_AGENT_TOOLS = ["web_search_tavily"] - agent = ConversationAgent(model_hrid="default-model") # re-init to pick up new settings - with agent.override(model=TestModel(), deps=context_deps): - response = agent.run_sync("What tools do you have?") - assert response.output == '{"web_search_tavily":[]}' - - # Tool configured, context disables web search - context_deps.web_search_enabled = False - with agent.override(model=TestModel(), deps=context_deps): - response = agent.run_sync("What tools do you have?") - assert response.output == "success (no tool calls)" + assert agent.is_web_search_configured() is True diff --git a/src/backend/chat/tests/clients/pydantic_ai/test_smart_web_search.py b/src/backend/chat/tests/clients/pydantic_ai/test_smart_web_search.py index e7964a8..43d97d6 100644 --- a/src/backend/chat/tests/clients/pydantic_ai/test_smart_web_search.py +++ b/src/backend/chat/tests/clients/pydantic_ai/test_smart_web_search.py @@ -23,7 +23,8 @@ def _llm_config_with_websearch(settings): is_active=True, icon=None, system_prompt="You are an amazing assistant.", - tools=["web_search_brave_with_document_backend"], + tools=[], + web_search="chat.tools.web_search_brave.web_search_brave_llm_context", provider=LLMProvider( hrid="unused", base_url="https://example.com", @@ -48,6 +49,7 @@ def test_smart_search_disabled_suppresses_tool_at_runtime(_llm_config_with_webse if not service._is_smart_search_enabled and service._is_web_search_enabled: service._context_deps.web_search_enabled = False + service._setup_web_search_tool() with service.conversation_agent.override(model=TestModel(), deps=service._context_deps): response = service.conversation_agent.run_sync("Search the web for something.") @@ -65,10 +67,11 @@ def test_smart_search_enabled_tool_is_called(_llm_config_with_websearch): assert service._is_smart_search_enabled is True assert service._context_deps.web_search_enabled is True + service._setup_web_search_tool() with service.conversation_agent.override(model=TestModel(), deps=service._context_deps): response = service.conversation_agent.run_sync("Search the web for something.") - assert "web_search_brave_with_document_backend" in response.output + assert "web_search" in response.output def test_force_websearch_overrides_smart_search_disabled(_llm_config_with_websearch): @@ -82,14 +85,16 @@ def test_force_websearch_overrides_smart_search_disabled(_llm_config_with_websea assert service._is_smart_search_enabled is False assert service._context_deps.web_search_enabled is False + # Match _run_agent: register the tool first, then enable deps + forced prompt. + service._setup_web_search_tool() service._setup_web_search(force_web_search=True) - web_search_tool_name = service.conversation_agent.get_web_search_tool_name() + assert service.conversation_agent.is_web_search_configured() is True assert service._context_deps.web_search_enabled is True assert any( - callable(instr) and web_search_tool_name in instr() + callable(instr) and "web_search" in instr() for instr in service.conversation_agent._instructions ) with service.conversation_agent.override(model=TestModel(), deps=service._context_deps): response = service.conversation_agent.run_sync("Search the web for something.") - assert "web_search_brave_with_document_backend" in response.output + assert "web_search" in response.output diff --git a/src/backend/chat/tests/tools/test_web_search_brave.py b/src/backend/chat/tests/tools/test_web_search_brave.py index debcfa2..f5681b4 100644 --- a/src/backend/chat/tests/tools/test_web_search_brave.py +++ b/src/backend/chat/tests/tools/test_web_search_brave.py @@ -21,13 +21,14 @@ from chat.tools.web_search_brave import ( _extract_and_summarize_snippets_async, _fetch_and_extract_async, _fetch_and_store_async, - _query_brave_api_async, + _query_brave_llm_context_api_async, format_tool_return, web_search_brave, web_search_brave_with_document_backend, ) -BRAVE_URL = "https://api.search.brave.com/res/v1/web/search" +BRAVE_WEB_SEARCH_URL = "https://api.search.brave.com/res/v1/web/search" +BRAVE_LLM_CONTEXT_URL = "https://api.search.brave.com/res/v1/llm/context" @pytest.fixture(autouse=True) @@ -67,7 +68,7 @@ def fixture_mocked_context(): @respx.mock async def test_agent_web_search_brave_success_with_extra_snippets(mocked_context): """Test when the Brave search returns results with extra_snippets.""" - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=200, json={ @@ -122,7 +123,7 @@ async def test_agent_web_search_brave_success_with_extra_snippets(mocked_context @respx.mock async def test_agent_web_search_brave_success_without_extra_snippets(mocked_context): """Test when the Brave search returns results without extra_snippets.""" - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=200, json={ @@ -168,7 +169,7 @@ async def test_agent_web_search_brave_success_without_extra_snippets_summarizati """Test when the Brave search returns results without extra_snippets with summarization.""" settings.BRAVE_SUMMARIZATION_ENABLED = True - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=200, json={ @@ -216,7 +217,7 @@ async def test_agent_web_search_brave_success_without_extra_snippets_summarizati @respx.mock async def test_agent_web_search_brave_empty_results(mocked_context): """Test when the Brave search returns no results.""" - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=200, json={"web": {"results": []}}, @@ -233,7 +234,7 @@ async def test_agent_web_search_brave_empty_results(mocked_context): @respx.mock async def test_agent_web_search_brave_http_error(mocked_context): """Test handling of HTTP errors from Brave API.""" - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=500, json={"error": "Internal Server Error"}, @@ -256,7 +257,7 @@ async def test_agent_web_search_brave_params_exclude_none(settings, mocked_conte settings.BRAVE_SEARCH_COUNTRY = None settings.BRAVE_SEARCH_LANG = None - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=200, json={"web": {"results": []}}, @@ -290,7 +291,7 @@ async def test_agent_web_search_brave_params_exclude_none(settings, mocked_conte @respx.mock async def test_agent_web_search_brave_concurrent_processing(mocked_context): """Test concurrent processing with asyncio.gather.""" - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=200, json={ @@ -401,7 +402,7 @@ async def test_extract_and_summarize_snippets_summarization_failure(settings): @respx.mock async def test_web_search_brave_with_document_backend_success(mocked_context): """Test web_search_brave_with_document_backend with successful RAG search.""" - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=200, json={ @@ -452,7 +453,7 @@ async def test_web_search_brave_with_document_backend_success(mocked_context): @respx.mock async def test_web_search_brave_with_document_backend_fetch_error(mocked_context): """Test web_search_brave_with_document_backend when document fetching fails.""" - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=200, json={ @@ -495,7 +496,7 @@ async def test_web_search_brave_with_document_backend_no_matching_rag_results(mo This is actually a problematic scenario, but we want to ensure graceful handling. """ - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=200, json={ @@ -565,37 +566,37 @@ async def test_fetch_and_store_empty_document(): @pytest.mark.asyncio @respx.mock async def test_query_brave_api_missing_web_key(): - """Test _query_brave_api_async when response doesn't contain 'web' key.""" - respx.get(BRAVE_URL).mock( + """Test _query_brave_llm_context_api_async when response doesn't contain 'web' key.""" + respx.get(BRAVE_LLM_CONTEXT_URL).mock( return_value=httpx.Response( status_code=200, json={"error": "no web results"}, ) ) - result = await _query_brave_api_async("query") + result = await _query_brave_llm_context_api_async("query") assert not result @pytest.mark.asyncio @respx.mock async def test_query_brave_api_missing_results_key(): - """Test _query_brave_api_async when 'web' exists but 'results' is missing.""" - respx.get(BRAVE_URL).mock( + """Test _query_brave_llm_context_api_async when 'web' exists but 'results' is missing.""" + respx.get(BRAVE_LLM_CONTEXT_URL).mock( return_value=httpx.Response( status_code=200, json={"web": {}}, ) ) - result = await _query_brave_api_async("query") + result = await _query_brave_llm_context_api_async("query") assert not result @pytest.mark.asyncio @respx.mock async def test_query_brave_api_rate_limit(): - """Test _query_brave_api_async handles 429 rate limit errors.""" - respx.get(BRAVE_URL).mock( + """Test _query_brave_llm_context_api_async handles 429 rate limit errors.""" + respx.get(BRAVE_LLM_CONTEXT_URL).mock( return_value=httpx.Response( status_code=429, json={"error": "Rate limit exceeded"}, @@ -603,7 +604,7 @@ async def test_query_brave_api_rate_limit(): ) with pytest.raises(ModelRetry) as exc: - await _query_brave_api_async("query") + await _query_brave_llm_context_api_async("query") assert "rate limited" in str(exc.value).lower() @@ -611,8 +612,8 @@ async def test_query_brave_api_rate_limit(): @pytest.mark.asyncio @respx.mock async def test_query_brave_api_client_error(): - """Test _query_brave_api_async handles 4xx client errors.""" - respx.get(BRAVE_URL).mock( + """Test _query_brave_llm_context_api_async handles 4xx client errors.""" + respx.get(BRAVE_LLM_CONTEXT_URL).mock( return_value=httpx.Response( status_code=400, json={"error": "Bad request"}, @@ -620,7 +621,7 @@ async def test_query_brave_api_client_error(): ) with pytest.raises(ModelCannotRetry) as exc: - await _query_brave_api_async("query") + await _query_brave_llm_context_api_async("query") assert "client error" in str(exc.value).lower() @@ -628,11 +629,11 @@ async def test_query_brave_api_client_error(): @pytest.mark.asyncio @respx.mock async def test_query_brave_api_timeout(): - """Test _query_brave_api_async handles timeout errors.""" - respx.get(BRAVE_URL).mock(side_effect=httpx.TimeoutException("Request timed out")) + """Test _query_brave_llm_context_api_async handles timeout errors.""" + respx.get(BRAVE_LLM_CONTEXT_URL).mock(side_effect=httpx.TimeoutException("Request timed out")) with pytest.raises(ModelRetry) as exc: - await _query_brave_api_async("query") + await _query_brave_llm_context_api_async("query") assert "timed out" in str(exc.value).lower() @@ -640,11 +641,11 @@ async def test_query_brave_api_timeout(): @pytest.mark.asyncio @respx.mock async def test_query_brave_api_generic_http_error(): - """Test _query_brave_api_async handles generic HTTP errors.""" - respx.get(BRAVE_URL).mock(side_effect=httpx.ConnectError("Connection failed")) + """Test _query_brave_llm_context_api_async handles generic HTTP errors.""" + respx.get(BRAVE_LLM_CONTEXT_URL).mock(side_effect=httpx.ConnectError("Connection failed")) with pytest.raises(ModelRetry) as exc: - await _query_brave_api_async("query") + await _query_brave_llm_context_api_async("query") assert "connection error" in str(exc.value).lower() @@ -652,11 +653,11 @@ async def test_query_brave_api_generic_http_error(): @pytest.mark.asyncio @respx.mock async def test_query_brave_api_unexpected_error(): - """Test _query_brave_api_async handles unexpected errors.""" - respx.get(BRAVE_URL).mock(side_effect=ValueError("Unexpected value error")) + """Test _query_brave_llm_context_api_async handles unexpected errors.""" + respx.get(BRAVE_LLM_CONTEXT_URL).mock(side_effect=ValueError("Unexpected value error")) with pytest.raises(ModelCannotRetry) as exc: - await _query_brave_api_async("query") + await _query_brave_llm_context_api_async("query") assert "unexpected error" in str(exc.value).lower() @@ -768,7 +769,7 @@ async def test_fetch_and_store_extraction_error(): @respx.mock async def test_web_search_brave_mixed_results(mocked_context): """Test web_search_brave with mixed results (some with snippets, some without).""" - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=200, json={ @@ -812,7 +813,7 @@ async def test_web_search_brave_mixed_results(mocked_context): @respx.mock async def test_web_search_brave_extraction_fails_for_all(mocked_context): """Test web_search_brave when extraction fails for all results without snippets.""" - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=200, json={ @@ -845,7 +846,7 @@ async def test_web_search_brave_extraction_fails_for_all(mocked_context): @respx.mock async def test_web_search_brave_model_cannot_retry_exception(mocked_context): """Test web_search_brave handling of ModelCannotRetry from API.""" - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=403, json={"error": "Forbidden"}, @@ -863,7 +864,7 @@ async def test_web_search_brave_model_cannot_retry_exception(mocked_context): @respx.mock async def test_web_search_brave_unexpected_exception(mocked_context): """Test web_search_brave handling of unexpected exceptions.""" - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=200, json={ @@ -894,7 +895,7 @@ async def test_web_search_brave_unexpected_exception(mocked_context): @respx.mock async def test_web_search_brave_with_document_backend_empty_rag_results(mocked_context): """Test web_search_brave_with_document_backend when RAG search returns empty results.""" - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=200, json={ @@ -929,7 +930,7 @@ async def test_web_search_brave_with_document_backend_empty_rag_results(mocked_c @respx.mock async def test_web_search_brave_with_document_backend_store_exception(mocked_context): """Test web_search_brave_with_document_backend when document store raises exception.""" - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=200, json={ @@ -958,7 +959,7 @@ async def test_web_search_brave_with_document_backend_store_exception(mocked_con @respx.mock async def test_web_search_brave_with_document_backend_unexpected_exception(mocked_context): """Test web_search_brave_with_document_backend handling of unexpected exceptions.""" - respx.get(BRAVE_URL).mock(side_effect=TypeError("Unexpected type error")) + respx.get(BRAVE_WEB_SEARCH_URL).mock(side_effect=TypeError("Unexpected type error")) result = await web_search_brave_with_document_backend(mocked_context, "error query") assert result == ( @@ -971,7 +972,7 @@ async def test_web_search_brave_with_document_backend_unexpected_exception(mocke @respx.mock async def test_web_search_brave_with_document_backend_model_cannot_retry(mocked_context): """Test web_search_brave_with_document_backend handling of ModelCannotRetry.""" - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=401, json={"error": "Unauthorized"}, @@ -989,7 +990,7 @@ async def test_web_search_brave_with_document_backend_model_cannot_retry(mocked_ @respx.mock async def test_web_search_brave_with_document_backend_rag_search_params(mocked_context): """Test web_search_brave_with_document_backend passes correct parameters to RAG search.""" - respx.get(BRAVE_URL).mock( + respx.get(BRAVE_WEB_SEARCH_URL).mock( return_value=httpx.Response( status_code=200, json={ diff --git a/src/backend/chat/tools/__init__.py b/src/backend/chat/tools/__init__.py index a9769a6..2e01ff6 100644 --- a/src/backend/chat/tools/__init__.py +++ b/src/backend/chat/tools/__init__.py @@ -1,46 +1,13 @@ """Tools for the chat agent.""" -from pydantic_ai import Tool, ToolDefinition - +from pydantic_ai import Tool # noqa: I001 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, web_search_brave_with_document_backend -from .web_search_tavily import web_search_tavily - - -async def only_if_web_search_enabled(ctx, tool_def: ToolDefinition) -> ToolDefinition | None: - """Prepare function to include a tool only if web search is enabled in the context.""" - return tool_def if ctx.deps.web_search_enabled else None def get_pydantic_tools_by_name(name: str) -> Tool: """Get a tool by its name.""" tool_dict = { "get_current_weather": Tool(get_current_weather, takes_ctx=False), - "web_search_brave": Tool( - web_search_brave, - takes_ctx=True, - prepare=only_if_web_search_enabled, - max_retries=2, - ), - "web_search_brave_with_document_backend": Tool( - web_search_brave_with_document_backend, - takes_ctx=True, - prepare=only_if_web_search_enabled, - max_retries=2, - ), - "web_search_tavily": Tool( - web_search_tavily, - takes_ctx=False, - prepare=only_if_web_search_enabled, - max_retries=2, - ), - "web_search_albert_rag": Tool( - web_search_albert_rag, - takes_ctx=True, - prepare=only_if_web_search_enabled, - max_retries=2, - ), } return tool_dict[name] # will raise on purpose if name is not found diff --git a/src/backend/chat/tools/web_search_brave.py b/src/backend/chat/tools/web_search_brave.py index cd32dd4..0e35b13 100644 --- a/src/backend/chat/tools/web_search_brave.py +++ b/src/backend/chat/tools/web_search_brave.py @@ -142,23 +142,33 @@ async def _fetch_and_store_async(url: str, document_store, **kwargs) -> None: # Continue with other documents -async def _query_brave_api_async(query: str) -> List[dict]: - """Query the Brave Search API and return the raw results.""" - url = "https://api.search.brave.com/res/v1/web/search" +def _normalize_llm_context_results(json_response: dict) -> List[dict]: + """Normalize Brave LLM context payload into our common result shape.""" + generic_results = json_response.get("grounding", {}).get("generic", []) or [] + normalized_results: List[dict] = [] + for item in generic_results: + item_url = item.get("url") + if not item_url: + continue + + normalized_results.append( + { + "url": item_url, + # Fallback to URL if no title is provided + "title": item.get("title") or item_url, + # `snippets` is already a list + "snippets": item.get("snippets") or [], + } + ) + return normalized_results + + +async def _query_brave_api_with_endpoint_async(url: str, data: dict) -> List[dict]: + """Query a Brave endpoint and return raw results normalized to our schema.""" headers = { "Accept": "application/json", "X-Subscription-Token": settings.BRAVE_API_KEY, } - data = { - "q": query, - "country": settings.BRAVE_SEARCH_COUNTRY, - "search_lang": settings.BRAVE_SEARCH_LANG, - "count": settings.BRAVE_MAX_RESULTS, - "safesearch": settings.BRAVE_SEARCH_SAFE_SEARCH, - "spellcheck": settings.BRAVE_SEARCH_SPELLCHECK, - "result_filter": "web,faq,query", - "extra_snippets": settings.BRAVE_SEARCH_EXTRA_SNIPPETS, - } params = {k: v for k, v in data.items() if v is not None} try: @@ -167,6 +177,12 @@ async def _query_brave_api_async(query: str) -> List[dict]: response.raise_for_status() json_response = response.json() + # LLM context API: results are under `grounding.generic` + # See: https://api-dashboard.search.brave.com/documentation/services/llm-context + if "grounding" in json_response: + return _normalize_llm_context_results(json_response) + + # Fallback for classic web search JSON shape # https://api-dashboard.search.brave.com/app/documentation/web-search/responses#Result return json_response.get("web", {}).get("results", []) @@ -209,44 +225,99 @@ async def _query_brave_api_async(query: str) -> List[dict]: ) from e +async def _query_brave_llm_context_api_async(query: str) -> List[dict]: + """Query Brave LLM context endpoint and return normalized results.""" + logger.debug("Using LLM context endpoint") + return await _query_brave_api_with_endpoint_async( + "https://api.search.brave.com/res/v1/llm/context", + { + "q": query, + "country": settings.BRAVE_SEARCH_COUNTRY, + "search_lang": settings.BRAVE_SEARCH_LANG, + "count": settings.BRAVE_MAX_RESULTS, + "safesearch": settings.BRAVE_SEARCH_SAFE_SEARCH, + "spellcheck": settings.BRAVE_SEARCH_SPELLCHECK, + "result_filter": "web,faq,query", + "extra_snippets": settings.BRAVE_SEARCH_EXTRA_SNIPPETS, + "maximum_number_of_urls": settings.BRAVE_MAX_RESULTS, + "maximum_number_of_tokens": settings.BRAVE_MAX_TOKENS, + "maximum_number_of_snippets": settings.BRAVE_MAX_SNIPPETS, + "maximum_number_of_snippets_per_url": settings.BRAVE_MAX_SNIPPETS_PER_URL, + }, + ) + + +async def _query_brave_web_search_api_async(query: str) -> List[dict]: + """Query Brave classic web search endpoint and return normalized results.""" + logger.debug("Using classic web search endpoint") + return await _query_brave_api_with_endpoint_async( + "https://api.search.brave.com/res/v1/web/search", + { + "q": query, + "country": settings.BRAVE_SEARCH_COUNTRY, + "search_lang": settings.BRAVE_SEARCH_LANG, + "count": settings.BRAVE_MAX_RESULTS, + "safesearch": settings.BRAVE_SEARCH_SAFE_SEARCH, + "spellcheck": settings.BRAVE_SEARCH_SPELLCHECK, + "result_filter": "web,faq,query", + "extra_snippets": settings.BRAVE_SEARCH_EXTRA_SNIPPETS, + }, + ) + + def format_tool_return(raw_search_results: List[dict]) -> ToolReturn: - """Format the raw search results into a ToolReturn object.""" + """Build the tool payload from Brave results. + + Keep only sources that have non-empty snippets and prefer `snippets` over + `extra_snippets` when both are present. + """ + formatted_results = {} + sources = set() + + for idx, result in enumerate(raw_search_results): + logger.debug("Formatting result: %s", result) + snippets = result.get("snippets") or result.get("extra_snippets") or [] + if not snippets: + continue + + formatted_results[str(idx)] = { + "url": result["url"], + "title": result["title"], + "snippets": snippets, + } + sources.add(result["url"]) + return ToolReturn( - # Format return value "mistral-like": https://docs.mistral.ai/capabilities/citations/ - return_value={ - str(idx): { - "url": result["url"], - "title": result["title"], - "snippets": result.get("extra_snippets", []), - } - for idx, result in enumerate(raw_search_results) - if result.get("extra_snippets", []) - }, - metadata={ - "sources": { - result["url"] for result in raw_search_results if result.get("extra_snippets", []) - } - }, + return_value=formatted_results, + metadata={"sources": sources}, ) @last_model_retry_soft_fail async def web_search_brave(_ctx: RunContext, query: str) -> ToolReturn: """ - Search the web for up-to-date information + Search the web for up-to-date information. + This function use the classic websearch endpoint of the Brave API. + URLs are then fetched and extracted using trafilatura. + The extracted text is then summarized using the LLM summarization agent. + The results are then formatted and returned. Args: _ctx (RunContext): The run context, used by the wrapper. query (str): The query to search for. """ + logger.debug("Starting classic web search without RAG backend for query: %s", query) try: - raw_search_results = await _query_brave_api_async(query) + raw_search_results = await _query_brave_web_search_api_async(query) await sync_to_async(reset_caches)() # Clear trafilatura caches to avoid memory bloat/leaks - # Parallelize fetch/extract for results that don't include extra_snippets + # Parallelize fetch/extract only for results that don't already include any snippets + # (neither Brave `snippets` nor `extra_snippets`). to_process = [ - (idx, r) for idx, r in enumerate(raw_search_results) if not r.get("extra_snippets") + (idx, r) + for idx, r in enumerate(raw_search_results) + if not r.get("extra_snippets") and not r.get("snippets") ] if to_process: @@ -283,18 +354,45 @@ async def web_search_brave(_ctx: RunContext, query: str) -> ToolReturn: ) from exc +@last_model_retry_soft_fail +async def web_search_brave_llm_context(_ctx: RunContext, query: str) -> ToolReturn: + """ + Search the web using Brave LLM context endpoint (no RAG post-processing). + This function use the LLM context endpoint of the Brave API. + The results are then formatted and returned. + """ + logger.debug("Starting web search with LLM context endpoint for query: %s", query) + try: + raw_search_results = await _query_brave_llm_context_api_async(query) + formatted_result = format_tool_return(raw_search_results) + if not formatted_result.return_value: + raise ModelRetry("No valid search results were extracted from Brave LLM context.") + return formatted_result + except (ModelCannotRetry, ModelRetry): + raise + except Exception as exc: + logger.exception("Unexpected error in web_search_brave_llm_context: %s", exc) + raise ModelCannotRetry( + f"An unexpected error occurred during web search: {type(exc).__name__}. " + "You must explain this to the user and not try to answer based on your knowledge." + ) from exc + + @last_model_retry_soft_fail async def web_search_brave_with_document_backend(ctx: RunContext, query: str) -> ToolReturn: """ - Search the web for up-to-date information using RAG backend + Search the web for up-to-date information using RAG backend. + URLs are then fetched and extracted using trafilatura. + The extracted text is then stored in a temporary document store for RAG search. + The RAG search is then performed and the results are returned. Args: ctx (RunContext): The run context containing the conversation. query (str): The query to search for. """ - logger.info("Starting web search with RAG backend for query: %s", query) + logger.debug("Starting web search with RAG backend for query: %s", query) try: - raw_search_results = await _query_brave_api_async(query) + raw_search_results = await _query_brave_web_search_api_async(query) # Clear trafilatura caches in thread pool to avoid blocking loop = asyncio.get_event_loop() @@ -328,7 +426,7 @@ async def web_search_brave_with_document_backend(ctx: RunContext, query: str) -> session=ctx.deps.session, user_sub=ctx.deps.user.sub, ) - logger.info("RAG search returned: %s", rag_results) + logger.debug("RAG search returned: %s", rag_results) ctx.usage += RunUsage( input_tokens=rag_results.usage.prompt_tokens, diff --git a/src/backend/conversations/brave_settings.py b/src/backend/conversations/brave_settings.py index 32fab75..1a6f1bd 100644 --- a/src/backend/conversations/brave_settings.py +++ b/src/backend/conversations/brave_settings.py @@ -74,3 +74,20 @@ class BraveSettings: environ_name="BRAVE_SEARCH_EXTRA_SNIPPETS", environ_prefix=None, ) + + # LLM context endpoint limits + BRAVE_MAX_TOKENS = values.IntegerValue( + default=8192, + environ_name="BRAVE_MAX_TOKENS", + environ_prefix=None, + ) + BRAVE_MAX_SNIPPETS = values.IntegerValue( + default=50, + environ_name="BRAVE_MAX_SNIPPETS", + environ_prefix=None, + ) + BRAVE_MAX_SNIPPETS_PER_URL = values.IntegerValue( + default=10, + environ_name="BRAVE_MAX_SNIPPETS_PER_URL", + environ_prefix=None, + )