From e85af923abc16c3eaad1542d3b29a7cf6fc82321 Mon Sep 17 00:00:00 2001 From: George Pickett <297992784+georgeatparallel@users.noreply.github.com> Date: Tue, 20 Oct 2026 10:33:13 +0000 Subject: [PATCH 1/2] Add keyless Parallel Search MCP for unconfigured web search --- README.md | 19 + example.env | 8 +- prometheus/configuration/config.py | 7 +- .../nodes/issue_bug_analyzer_node.py | 12 +- .../issue_documentation_analyzer_node.py | 5 +- .../nodes/issue_feature_analyzer_node.py | 12 +- .../nodes/issue_question_analyzer_node.py | 5 +- prometheus/tools/web_search.py | 85 ++++- pyproject.toml | 4 +- tests/tools/test_parallel_search.py | 351 ++++++++++++++++++ tests/tools/test_web_search.py | 26 ++ 11 files changed, 504 insertions(+), 30 deletions(-) create mode 100644 tests/tools/test_parallel_search.py diff --git a/README.md b/README.md index 167a5e9..1fed279 100644 --- a/README.md +++ b/README.md @@ -271,6 +271,25 @@ Verify at [http://localhost:7474](http://localhost:7474) --- +### Web search + +Issue-analysis agents use [Parallel Search MCP](https://docs.parallel.ai/integrations/mcp/search-mcp) +when no Tavily key is configured. Search runs in Fast mode without an account or API key. +The anonymous service has usage limits; requests can return quota errors. + +`PROMETHEUS_WEB_SEARCH_PROVIDER` defaults to `auto`: an existing +`PROMETHEUS_TAVILY_API_KEY` keeps Tavily selected. Set the provider to `parallel` +or `tavily` to choose explicitly, or `disabled` to remove web search from the agents. +An explicitly selected Tavily provider requires a Tavily key. Provider failures +are reported without switching providers. `PROMETHEUS_WEB_SEARCH_TIMEOUT` bounds +Parallel requests, including connection setup (30 seconds by default). + +Parallel uses server-managed result counts and source selection. Custom result +limits and domain filters remain available through Tavily; requesting them with +Parallel raises an error. The agent-facing tool accepts a `query` with either provider. +Queries are sent to the selected external service; use `disabled` when external +web search is not permitted. + ## 🤝 Contributing We welcome contributions! Please see our **[Contributing Guidelines](CONTRIBUTING.md)** for details on how to get started. diff --git a/example.env b/example.env index b4d0181..8b3f7c3 100644 --- a/example.env +++ b/example.env @@ -33,8 +33,12 @@ PROMETHEUS_ADVANCED_MODEL_TEMPERATURE=0.5 PROMETHEUS_BASE_MODEL_TEMPERATURE=0.5 -# Tavily API settings (You can sign up for a free API key at https://www.tavily.com/) -PROMETHEUS_TAVILY_API_KEY=your_tavily_api_key +# Web search: auto uses a configured Tavily key, otherwise free Parallel Search MCP. +# Use parallel or tavily to select explicitly, or disabled to remove search tools. +PROMETHEUS_WEB_SEARCH_PROVIDER=auto +PROMETHEUS_WEB_SEARCH_TIMEOUT=30 +# Optional Tavily API key (https://www.tavily.com/) +# PROMETHEUS_TAVILY_API_KEY=your_tavily_api_key # Database settings PROMETHEUS_DATABASE_URL=postgresql+asyncpg://postgres:password@postgres:5432/postgres diff --git a/prometheus/configuration/config.py b/prometheus/configuration/config.py index 0dec2e2..e42cf27 100644 --- a/prometheus/configuration/config.py +++ b/prometheus/configuration/config.py @@ -1,5 +1,6 @@ from typing import List, Literal, Optional +from pydantic import Field from pydantic_settings import BaseSettings, SettingsConfigDict @@ -62,8 +63,10 @@ class Settings(BaseSettings): # Default normal user repository number DEFAULT_USER_REPOSITORY_LIMIT: int = 5 - # tool for Websearch - TAVILY_API_KEY: str + # Web search: retain configured Tavily access; otherwise use keyless Parallel. + WEB_SEARCH_PROVIDER: Literal["auto", "parallel", "tavily", "disabled"] = "auto" + WEB_SEARCH_TIMEOUT: float = Field(default=30.0, gt=0, le=300) + TAVILY_API_KEY: Optional[str] = None # Athena semantic memory service ATHENA_BASE_URL: Optional[str] = None diff --git a/prometheus/lang_graph/nodes/issue_bug_analyzer_node.py b/prometheus/lang_graph/nodes/issue_bug_analyzer_node.py index 7fe8529..7dcad23 100644 --- a/prometheus/lang_graph/nodes/issue_bug_analyzer_node.py +++ b/prometheus/lang_graph/nodes/issue_bug_analyzer_node.py @@ -49,12 +49,7 @@ class IssueBugAnalyzerNode: Tools available: - web_search: Searches the web for technical information to aid in bug analysis and resolution. -When using the web_search tool, it's recommended to include these parameters: - - include_domains: ['stackoverflow.com', 'github.com', 'developer.mozilla.org', 'learn.microsoft.com', 'fastapi.tiangolo.com' - 'docs.python.org', 'pydantic.dev', 'pypi.org', 'readthedocs.org', 'docs.djangoproject.com','flask.palletsprojects.com'] - - search_depth: "advanced" - - Make sure to explicitly pass these parameters in your tool call. +Pass a specific technical query to web_search. Important: - Do NOT provide actual code snippets or diffs @@ -72,16 +67,19 @@ def __init__(self, model: BaseChatModel): self.model = model self.system_prompt = SystemMessage(self.SYS_PROMPT) self.tools = self._init_tools() - self.model_with_tools = model.bind_tools(self.tools) + self.model_with_tools = model.bind_tools(self.tools) if self.tools else model self._logger = logging.getLogger(f"thread-{threading.get_ident()}.{__name__}") def _init_tools(self): """Initializes tools for the node.""" tools = [] + if self.web_search_tool.provider == "disabled": + return tools web_search_fn = functools.partial(self.web_search_tool.web_search) web_search_tool = StructuredTool.from_function( func=web_search_fn, + coroutine=self.web_search_tool.aweb_search, name=self.web_search_tool.web_search.__name__, description=self.web_search_tool.web_search_spec.description, args_schema=self.web_search_tool.web_search_spec.input_schema, diff --git a/prometheus/lang_graph/nodes/issue_documentation_analyzer_node.py b/prometheus/lang_graph/nodes/issue_documentation_analyzer_node.py index ee1e9bb..1dda2f2 100644 --- a/prometheus/lang_graph/nodes/issue_documentation_analyzer_node.py +++ b/prometheus/lang_graph/nodes/issue_documentation_analyzer_node.py @@ -44,17 +44,20 @@ def __init__(self, model: BaseChatModel): self.system_prompt = SystemMessage(self.SYS_PROMPT) self.web_search_tool = WebSearchTool() self.tools = self._init_tools() - self.model_with_tools = model.bind_tools(self.tools) + self.model_with_tools = model.bind_tools(self.tools) if self.tools else model self._logger = logging.getLogger(f"thread-{threading.get_ident()}.{__name__}") def _init_tools(self): """Initializes tools for the node.""" tools = [] + if self.web_search_tool.provider == "disabled": + return tools web_search_fn = functools.partial(self.web_search_tool.web_search) web_search_tool = StructuredTool.from_function( func=web_search_fn, + coroutine=self.web_search_tool.aweb_search, name=self.web_search_tool.web_search.__name__, description=self.web_search_tool.web_search_spec.description, args_schema=self.web_search_tool.web_search_spec.input_schema, diff --git a/prometheus/lang_graph/nodes/issue_feature_analyzer_node.py b/prometheus/lang_graph/nodes/issue_feature_analyzer_node.py index c2dd6f0..141b90b 100644 --- a/prometheus/lang_graph/nodes/issue_feature_analyzer_node.py +++ b/prometheus/lang_graph/nodes/issue_feature_analyzer_node.py @@ -50,12 +50,7 @@ class IssueFeatureAnalyzerNode: Tools available: - web_search: Searches the web for technical information to aid in feature design and implementation. -When using the web_search tool, it's recommended to include these parameters: - - include_domains: ['stackoverflow.com', 'github.com', 'developer.mozilla.org', 'learn.microsoft.com', 'fastapi.tiangolo.com' - 'docs.python.org', 'pydantic.dev', 'pypi.org', 'readthedocs.org', 'docs.djangoproject.com','flask.palletsprojects.com'] - - search_depth: "advanced" - - Make sure to explicitly pass these parameters in your tool call. +Pass a specific technical query to web_search. Important: - Do NOT provide actual code snippets or diffs @@ -74,16 +69,19 @@ def __init__(self, model: BaseChatModel): self.model = model self.system_prompt = SystemMessage(self.SYS_PROMPT) self.tools = self._init_tools() - self.model_with_tools = model.bind_tools(self.tools) + self.model_with_tools = model.bind_tools(self.tools) if self.tools else model self._logger = logging.getLogger(f"thread-{threading.get_ident()}.{__name__}") def _init_tools(self): """Initializes tools for the node.""" tools = [] + if self.web_search_tool.provider == "disabled": + return tools web_search_fn = functools.partial(self.web_search_tool.web_search) web_search_tool = StructuredTool.from_function( func=web_search_fn, + coroutine=self.web_search_tool.aweb_search, name=self.web_search_tool.web_search.__name__, description=self.web_search_tool.web_search_spec.description, args_schema=self.web_search_tool.web_search_spec.input_schema, diff --git a/prometheus/lang_graph/nodes/issue_question_analyzer_node.py b/prometheus/lang_graph/nodes/issue_question_analyzer_node.py index 09ba24e..fd19647 100644 --- a/prometheus/lang_graph/nodes/issue_question_analyzer_node.py +++ b/prometheus/lang_graph/nodes/issue_question_analyzer_node.py @@ -86,17 +86,20 @@ def __init__(self, model: BaseChatModel): self.system_prompt = SystemMessage(self.SYS_PROMPT) self.web_search_tool = WebSearchTool() self.tools = self._init_tools() - self.model_with_tools = model.bind_tools(self.tools) + self.model_with_tools = model.bind_tools(self.tools) if self.tools else model self._logger = logging.getLogger(f"thread-{threading.get_ident()}.{__name__}") def _init_tools(self): """Initializes tools for the node.""" tools = [] + if self.web_search_tool.provider == "disabled": + return tools web_search_fn = functools.partial(self.web_search_tool.web_search) web_search_tool = StructuredTool.from_function( func=web_search_fn, + coroutine=self.web_search_tool.aweb_search, name=self.web_search_tool.web_search.__name__, description=self.web_search_tool.web_search_spec.description, args_schema=self.web_search_tool.web_search_spec.input_schema, diff --git a/prometheus/tools/web_search.py b/prometheus/tools/web_search.py index 78232db..9f1e7a6 100644 --- a/prometheus/tools/web_search.py +++ b/prometheus/tools/web_search.py @@ -1,8 +1,12 @@ +import asyncio import logging import threading from dataclasses import dataclass from typing import Annotated +from mcp import ClientSession +from mcp.client.streamable_http import streamablehttp_client +from mcp.types import Implementation from pydantic import BaseModel, Field from tavily import InvalidAPIKeyError, TavilyClient, UsageLimitExceededError @@ -17,7 +21,7 @@ class ToolSpec: class WebSearchInput(BaseModel): - """Base parameters for Tavily search.""" + """Parameters shared by the web search providers.""" query: Annotated[str, Field(description="Search query")] @@ -79,8 +83,12 @@ def __init__(self): self._logger = logging.getLogger(f"thread-{threading.get_ident()}.{__name__}") tavily_api_key = settings.TAVILY_API_KEY - if tavily_api_key is None: - self._logger.warning("Tavily API key is not set") + provider = settings.WEB_SEARCH_PROVIDER + self.provider = ( + ("tavily" if tavily_api_key else "parallel") if provider == "auto" else provider + ) + self.timeout = settings.WEB_SEARCH_TIMEOUT + if self.provider != "tavily" or not tavily_api_key: tavily_client = None else: tavily_client = TavilyClient(api_key=tavily_api_key) @@ -89,21 +97,30 @@ def __init__(self): def web_search( self, query: str, - max_results: int = 5, + max_results: int | None = None, include_domains=None, - exclude_domains: list[str] = None, + exclude_domains: list[str] | None = None, ) -> str: """Search the web for technical information to aid in bug analysis and resolution. Args: query: Search query string. - max_results: Maximum number of results to return (default: 5). - include_domains: List of domains to include in search. - exclude_domains: List of domains to exclude from search. + max_results: Tavily result limit (default: 5); Parallel uses server settings. + include_domains: Tavily domains to include; unsupported on Parallel. + exclude_domains: Tavily domains to exclude; unsupported on Parallel. Returns: Formatted search results as a string. """ + if self.provider == "parallel": + return asyncio.run( + self.aweb_search(query, max_results, include_domains, exclude_domains) + ) + if self.provider == "disabled": + raise WebSearchToolException("Web search is disabled") + if self.tavily_client is None: + raise WebSearchToolException("Tavily requires PROMETHEUS_TAVILY_API_KEY") + # Set default include domains if not provided if include_domains is None: include_domains = [ @@ -123,7 +140,7 @@ def web_search( try: response = self.tavily_client.search( query=query, - max_results=max_results, + max_results=5 if max_results is None else max_results, search_depth="advanced", include_answer=True, include_domains=include_domains or [], # Convert None to an empty list @@ -140,3 +157,53 @@ def web_search( format_response = format_results(response) self._logger.info(f"web_search format_response: {format_response}") return format_response + + async def aweb_search( + self, + query: str, + max_results: int | None = None, + include_domains=None, + exclude_domains: list[str] | None = None, + ) -> str: + """Search asynchronously, allowing Parallel requests to be cancelled.""" + if self.provider != "parallel": + return await asyncio.to_thread( + self.web_search, query, max_results, include_domains, exclude_domains + ) + if max_results is not None or include_domains or exclude_domains: + raise WebSearchToolException( + "Parallel uses server-managed search settings. Select the tavily provider " + "to use domain filters or custom result limits." + ) + try: + return await asyncio.wait_for(self._parallel_search(query), timeout=self.timeout) + except TimeoutError: + raise WebSearchToolException("Parallel search timed out") from None + except WebSearchToolException: + raise + except Exception: + # SDK transport errors can contain request details. Keep the public error concise. + raise WebSearchToolException("Parallel Search MCP request failed") from None + + async def _parallel_search(self, query: str) -> str: + # Each call owns and closes its transport. Do not attach a process/client-wide + # conversation ID: analyzer instances are reused across independent issue runs. + headers = {"User-Agent": "Prometheus/1.3 (+https://github.com/EuniAI/Prometheus)"} + async with streamablehttp_client("https://search.parallel.ai/mcp", headers=headers) as ( + read, + write, + _, + ): + async with ClientSession( + read, write, client_info=Implementation(name="Prometheus", version="1.3") + ) as session: + await session.initialize() + result = await session.call_tool( + "web_search", {"objective": query, "search_queries": [query]} + ) + text = "\n".join(block.text for block in result.content if block.type == "text") + if result.isError: + raise WebSearchToolException(f"Parallel Search MCP returned an error: {text}") + if not text.strip(): + raise WebSearchToolException("Parallel Search MCP returned no text results") + return text diff --git a/pyproject.toml b/pyproject.toml index e502859..f446819 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,9 @@ dependencies = [ "sqlmodel==0.0.24", "asyncpg", "pyjwt==2.6.0", - "mcp>=1.4.1", + "mcp>=1.10.1,<2", + # MCP versions compatible with the pinned PyJWT require Pydantic <2.12. + "pydantic>=2.7,<2.12", "tavily-python>=0.5.1", "langchain-mcp-adapters>=0.1.9", "httpx==0.28.1", diff --git a/tests/tools/test_parallel_search.py b/tests/tools/test_parallel_search.py new file mode 100644 index 0000000..42be299 --- /dev/null +++ b/tests/tools/test_parallel_search.py @@ -0,0 +1,351 @@ +"""Exercise keyless search through the analyzer's real tool and MCP transports.""" + +import asyncio +import json +from types import SimpleNamespace +from unittest.mock import Mock + +import httpx +import pytest +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage +from langgraph.prebuilt import ToolNode + +from prometheus.configuration.config import Settings +from prometheus.exceptions.web_search_tool_exception import WebSearchToolException +from prometheus.lang_graph.nodes.issue_bug_analyzer_node import IssueBugAnalyzerNode +from prometheus.lang_graph.nodes.issue_documentation_analyzer_node import ( + IssueDocumentationAnalyzerNode, +) +from prometheus.lang_graph.nodes.issue_feature_analyzer_node import IssueFeatureAnalyzerNode +from prometheus.lang_graph.nodes.issue_question_analyzer_node import IssueQuestionAnalyzerNode +from prometheus.tools import web_search + + +@pytest.fixture +def search_settings(monkeypatch): + config = SimpleNamespace( + WEB_SEARCH_PROVIDER="auto", TAVILY_API_KEY=None, WEB_SEARCH_TIMEOUT=2.0 + ) + monkeypatch.setattr(web_search, "settings", config) + return config + + +@pytest.fixture +def mcp_transport(monkeypatch): + """Keep the SDK protocol and request headers, replacing only HTTP I/O.""" + requests = [] + response = { + "content": [ + { + "type": "text", + "text": json.dumps( + { + "results": [ + { + "title": "Python errors", + "url": "https://docs.python.org/3/tutorial/errors.html", + "excerpts": ["Exceptions can be handled with try and except."], + } + ] + } + ), + } + ], + "isError": False, + } + + async def handle(request): + requests.append(request) + if request.method != "POST": + return httpx.Response(405) + if "_http_status" in response: + return httpx.Response(response["_http_status"], text="fixture HTTP failure") + payload = json.loads(request.content) + method = payload["method"] + if "id" not in payload: + return httpx.Response(202) + if method == "initialize": + result = { + "protocolVersion": "2025-03-26", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "fixture", "version": "1"}, + } + elif method == "tools/list": + result = {"tools": [{"name": "web_search", "inputSchema": {"type": "object"}}]} + else: + assert method == "tools/call" + result = response + return httpx.Response(200, json={"jsonrpc": "2.0", "id": payload["id"], "result": result}) + + original = web_search.streamablehttp_client + + def connect(url, **kwargs): + def factory(headers=None, timeout=None, auth=None): + return httpx.AsyncClient( + headers=headers, + timeout=timeout, + auth=auth, + transport=httpx.MockTransport(handle), + follow_redirects=True, + ) + + return original(url, **kwargs, httpx_client_factory=factory) + + monkeypatch.setattr(web_search, "streamablehttp_client", connect) + return requests, response + + +def model(): + llm = Mock() + llm.bind_tools.return_value = llm + return llm + + +def test_fresh_settings_without_search_credentials(monkeypatch): + # Supply unrelated required application settings while isolating dotenv and search credentials. + for name in ("PROMETHEUS_TAVILY_API_KEY", "PROMETHEUS_WEB_SEARCH_PROVIDER"): + monkeypatch.delenv(name, raising=False) + values = { + name: "fixture" for name, field in Settings.model_fields.items() if field.is_required() + } + values.update( + ENVIRONMENT="local", + BACKEND_CORS_ORIGINS=[], + ENABLE_AUTHENTICATION=False, + LOGGING_LEVEL="INFO", + NEO4J_BATCH_SIZE=1, + KNOWLEDGE_GRAPH_MAX_AST_DEPTH=1, + KNOWLEDGE_GRAPH_CHUNK_SIZE=100, + KNOWLEDGE_GRAPH_CHUNK_OVERLAP=1, + ) + config = Settings(_env_file=None, **values) + assert config.TAVILY_API_KEY is None + assert config.WEB_SEARCH_PROVIDER == "auto" + + +def test_analyzer_dispatch_and_consumption(search_settings, mcp_transport): + requests, _ = mcp_transport + llm = model() + node = IssueBugAnalyzerNode(llm) + assert node.web_search_tool.provider == "parallel" + llm.invoke.side_effect = [ + AIMessage( + content="", + tool_calls=[ + { + "name": "web_search", + "args": {"query": "Python ValueError exception handling"}, + "id": "s1", + } + ], + ), + AIMessage(content="Use try and except based on the documentation."), + ] + initial = HumanMessage(content="Analyze this error") + state = node({"issue_bug_analyzer_messages": [initial]}) + dispatched = ToolNode(node.tools, messages_key="issue_bug_analyzer_messages").invoke(state) + result = dispatched["issue_bug_analyzer_messages"][0] + assert isinstance(result, ToolMessage) + assert "https://docs.python.org/3/tutorial/errors.html" in result.content + assert "Exceptions can be handled" in result.content + final = node( + {"issue_bug_analyzer_messages": [initial, *state["issue_bug_analyzer_messages"], result]} + ) + assert "Use try and except" in final["issue_bug_analyzer_messages"][0].content + assert llm.invoke.call_args.args[0][-1] is result + calls = [json.loads(r.content) for r in requests if r.method == "POST"] + assert {c["method"] for c in calls} >= {"initialize", "tools/call"} + call = next(c for c in calls if c["method"] == "tools/call") + assert call["params"] == { + "name": "web_search", + "arguments": { + "objective": "Python ValueError exception handling", + "search_queries": ["Python ValueError exception handling"], + }, + } + for request in requests: + assert str(request.url) == "https://search.parallel.ai/mcp" + assert ( + request.headers["user-agent"] + == "Prometheus/1.3 (+https://github.com/EuniAI/Prometheus)" + ) + assert "authorization" not in request.headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "node_class", + [ + IssueBugAnalyzerNode, + IssueFeatureAnalyzerNode, + IssueDocumentationAnalyzerNode, + IssueQuestionAnalyzerNode, + ], +) +async def test_async_tool_dispatch(search_settings, mcp_transport, node_class): + node = node_class(model()) + result = await node.tools[0].ainvoke({"query": "Python ValueError exception handling"}) + assert "https://docs.python.org/3/tutorial/errors.html" in result + + +@pytest.mark.parametrize( + "provider,key,expected", + [ + ("auto", "saved-key", "tavily"), + ("auto", "", "parallel"), + ("parallel", "saved-key", "parallel"), + ("tavily", None, "tavily"), + ("disabled", "saved-key", "disabled"), + ], +) +def test_provider_choices(search_settings, monkeypatch, provider, key, expected): + search_settings.WEB_SEARCH_PROVIDER = provider + search_settings.TAVILY_API_KEY = key + tavily = Mock() + monkeypatch.setattr(web_search, "TavilyClient", tavily) + tool = web_search.WebSearchTool() + assert tool.provider == expected + assert tavily.called is (expected == "tavily" and bool(key)) + if provider == "tavily": + with pytest.raises(WebSearchToolException, match="requires PROMETHEUS_TAVILY_API_KEY"): + tool.web_search("Python errors") + assert not tavily.called + + +@pytest.mark.parametrize( + "node_class", + [ + IssueBugAnalyzerNode, + IssueFeatureAnalyzerNode, + IssueDocumentationAnalyzerNode, + IssueQuestionAnalyzerNode, + ], +) +def test_disabled_has_no_tool_or_network(search_settings, monkeypatch, node_class): + search_settings.WEB_SEARCH_PROVIDER = "disabled" + connect = Mock(side_effect=AssertionError("network must not run")) + monkeypatch.setattr(web_search, "streamablehttp_client", connect) + llm = model() + node = node_class(llm) + assert node.tools == [] + llm.bind_tools.assert_not_called() + with pytest.raises(WebSearchToolException, match="disabled"): + node.web_search_tool.web_search("Python errors") + connect.assert_not_called() + + +@pytest.mark.parametrize( + "kwargs", + [ + {"include_domains": ["docs.python.org"]}, + {"exclude_domains": ["example.org"]}, + {"max_results": 2}, + {"max_results": 5}, + ], +) +def test_parallel_rejects_unsupported_settings(search_settings, monkeypatch, kwargs): + connect = Mock(side_effect=AssertionError("network must not run")) + monkeypatch.setattr(web_search, "streamablehttp_client", connect) + with pytest.raises(WebSearchToolException, match="Select the tavily provider"): + web_search.WebSearchTool().web_search("Python errors", **kwargs) + connect.assert_not_called() + + +@pytest.mark.parametrize( + "response,error", + [ + ( + { + "content": [{"type": "text", "text": "Anonymous search quota exceeded"}], + "isError": True, + }, + "quota", + ), + ({"content": [], "isError": False}, "no text results"), + ], +) +def test_mcp_errors_do_not_fallback(search_settings, mcp_transport, monkeypatch, response, error): + search_settings.WEB_SEARCH_PROVIDER = "parallel" + search_settings.TAVILY_API_KEY = "saved-key" + _, fixture_response = mcp_transport + fixture_response.update(response) + tavily = Mock(side_effect=AssertionError("no fallback")) + monkeypatch.setattr(web_search, "TavilyClient", tavily) + with pytest.raises(WebSearchToolException, match=error): + IssueBugAnalyzerNode(model()).tools[0].invoke({"query": "Python errors"}) + tavily.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel", [False, True]) +async def test_timeout_and_cancellation_cleanup(search_settings, monkeypatch, cancel): + from contextlib import asynccontextmanager + + started = asyncio.Event() + closed = asyncio.Event() + + @asynccontextmanager + async def stalled_transport(*args, **kwargs): + try: + started.set() + await asyncio.Event().wait() + yield + finally: + closed.set() + + search_settings.WEB_SEARCH_TIMEOUT = 0.05 + monkeypatch.setattr(web_search, "streamablehttp_client", stalled_transport) + tool = IssueBugAnalyzerNode(model()).tools[0] + task = asyncio.create_task(tool.ainvoke({"query": "Python errors"})) + await started.wait() + if cancel: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + else: + with pytest.raises(WebSearchToolException, match="timed out"): + await task + assert closed.is_set() + + +def test_transport_failure_is_safe(search_settings, monkeypatch): + monkeypatch.setattr( + web_search, "streamablehttp_client", Mock(side_effect=OSError("private detail")) + ) + with pytest.raises(WebSearchToolException, match="^Parallel Search MCP request failed$"): + IssueBugAnalyzerNode(model()).tools[0].invoke({"query": "Python errors"}) + + +@pytest.mark.parametrize("status", [401, 429, 500]) +def test_http_errors(search_settings, mcp_transport, status): + _, response = mcp_transport + response["_http_status"] = status + with pytest.raises(WebSearchToolException, match="^Parallel Search MCP request failed$"): + IssueBugAnalyzerNode(model()).tools[0].invoke({"query": "Python errors"}) + + +@pytest.mark.asyncio +async def test_saved_tavily_uses_normal_dispatch(search_settings, monkeypatch): + search_settings.TAVILY_API_KEY = "saved-key" + client = Mock() + client.search.return_value = { + "results": [ + { + "title": "Python errors", + "url": "https://docs.python.org/", + "content": "Handle errors with try and except", + } + ] + } + factory = Mock(return_value=client) + monkeypatch.setattr(web_search, "TavilyClient", factory) + monkeypatch.setattr( + web_search, "streamablehttp_client", Mock(side_effect=AssertionError("no MCP")) + ) + node = IssueBugAnalyzerNode(model()) + result = await node.tools[0].ainvoke({"query": "Python errors"}) + assert "https://docs.python.org/" in result + factory.assert_called_once_with(api_key="saved-key") + assert client.search.call_args.kwargs["search_depth"] == "advanced" + assert client.search.call_args.kwargs["max_results"] == 5 + assert "docs.python.org" in client.search.call_args.kwargs["include_domains"] diff --git a/tests/tools/test_web_search.py b/tests/tools/test_web_search.py index facdf70..19a070c 100644 --- a/tests/tools/test_web_search.py +++ b/tests/tools/test_web_search.py @@ -208,6 +208,8 @@ class TestWebSearchTool: def test_init_with_api_key(self, mock_tavily_client, mock_settings): """Test WebSearchTool initialization with API key.""" mock_settings.TAVILY_API_KEY = "test_api_key" + mock_settings.WEB_SEARCH_PROVIDER = "auto" + mock_settings.WEB_SEARCH_TIMEOUT = 30.0 tool = WebSearchTool() @@ -219,6 +221,8 @@ def test_init_with_api_key(self, mock_tavily_client, mock_settings): def test_init_without_api_key(self, mock_tavily_client, mock_settings): """Test WebSearchTool initialization without API key.""" mock_settings.TAVILY_API_KEY = None + mock_settings.WEB_SEARCH_PROVIDER = "auto" + mock_settings.WEB_SEARCH_TIMEOUT = 30.0 tool = WebSearchTool() @@ -230,6 +234,8 @@ def test_init_without_api_key(self, mock_tavily_client, mock_settings): def test_web_search_success(self, mock_tavily_client, mock_settings): """Test successful web search.""" mock_settings.TAVILY_API_KEY = "test_api_key" + mock_settings.WEB_SEARCH_PROVIDER = "auto" + mock_settings.WEB_SEARCH_TIMEOUT = 30.0 # Setup mock client mock_client_instance = MagicMock() @@ -270,6 +276,8 @@ def test_web_search_success(self, mock_tavily_client, mock_settings): def test_web_search_with_custom_params(self, mock_tavily_client, mock_settings): """Test web search with custom parameters.""" mock_settings.TAVILY_API_KEY = "test_api_key" + mock_settings.WEB_SEARCH_PROVIDER = "auto" + mock_settings.WEB_SEARCH_TIMEOUT = 30.0 mock_client_instance = MagicMock() mock_tavily_client.return_value = mock_client_instance @@ -304,6 +312,8 @@ def test_web_search_with_custom_params(self, mock_tavily_client, mock_settings): def test_web_search_default_domains(self, mock_tavily_client, mock_settings): """Test web search uses default domains when not specified.""" mock_settings.TAVILY_API_KEY = "test_api_key" + mock_settings.WEB_SEARCH_PROVIDER = "auto" + mock_settings.WEB_SEARCH_TIMEOUT = 30.0 mock_client_instance = MagicMock() mock_tavily_client.return_value = mock_client_instance @@ -332,6 +342,8 @@ def test_web_search_default_domains(self, mock_tavily_client, mock_settings): def test_web_search_invalid_api_key(self, mock_tavily_client, mock_settings): """Test web search with invalid API key.""" mock_settings.TAVILY_API_KEY = "invalid_key" + mock_settings.WEB_SEARCH_PROVIDER = "auto" + mock_settings.WEB_SEARCH_TIMEOUT = 30.0 mock_client_instance = MagicMock() mock_tavily_client.return_value = mock_client_instance @@ -349,6 +361,8 @@ def test_web_search_invalid_api_key(self, mock_tavily_client, mock_settings): def test_web_search_usage_limit_exceeded(self, mock_tavily_client, mock_settings): """Test web search when usage limit is exceeded.""" mock_settings.TAVILY_API_KEY = "test_api_key" + mock_settings.WEB_SEARCH_PROVIDER = "auto" + mock_settings.WEB_SEARCH_TIMEOUT = 30.0 mock_client_instance = MagicMock() mock_tavily_client.return_value = mock_client_instance @@ -366,6 +380,8 @@ def test_web_search_usage_limit_exceeded(self, mock_tavily_client, mock_settings def test_web_search_generic_exception(self, mock_tavily_client, mock_settings): """Test web search with generic exception.""" mock_settings.TAVILY_API_KEY = "test_api_key" + mock_settings.WEB_SEARCH_PROVIDER = "auto" + mock_settings.WEB_SEARCH_TIMEOUT = 30.0 mock_client_instance = MagicMock() mock_tavily_client.return_value = mock_client_instance @@ -383,6 +399,8 @@ def test_web_search_generic_exception(self, mock_tavily_client, mock_settings): def test_web_search_empty_query(self, mock_tavily_client, mock_settings): """Test web search with empty query.""" mock_settings.TAVILY_API_KEY = "test_api_key" + mock_settings.WEB_SEARCH_PROVIDER = "auto" + mock_settings.WEB_SEARCH_TIMEOUT = 30.0 mock_client_instance = MagicMock() mock_tavily_client.return_value = mock_client_instance @@ -402,6 +420,8 @@ def test_web_search_empty_query(self, mock_tavily_client, mock_settings): def test_web_search_with_none_exclude_domains(self, mock_tavily_client, mock_settings): """Test web search with None exclude_domains.""" mock_settings.TAVILY_API_KEY = "test_api_key" + mock_settings.WEB_SEARCH_PROVIDER = "auto" + mock_settings.WEB_SEARCH_TIMEOUT = 30.0 mock_client_instance = MagicMock() mock_tavily_client.return_value = mock_client_instance @@ -421,6 +441,8 @@ def test_web_search_with_none_exclude_domains(self, mock_tavily_client, mock_set def test_web_search_complex_query(self, mock_tavily_client, mock_settings): """Test web search with complex query containing special characters.""" mock_settings.TAVILY_API_KEY = "test_api_key" + mock_settings.WEB_SEARCH_PROVIDER = "auto" + mock_settings.WEB_SEARCH_TIMEOUT = 30.0 mock_client_instance = MagicMock() mock_tavily_client.return_value = mock_client_instance @@ -476,6 +498,8 @@ class TestEdgeCases: def test_web_search_max_results_boundary(self, mock_tavily_client, mock_settings): """Test web search with boundary values for max_results.""" mock_settings.TAVILY_API_KEY = "test_api_key" + mock_settings.WEB_SEARCH_PROVIDER = "auto" + mock_settings.WEB_SEARCH_TIMEOUT = 30.0 mock_client_instance = MagicMock() mock_tavily_client.return_value = mock_client_instance @@ -518,6 +542,8 @@ def test_format_results_missing_optional_fields(self): def test_web_search_with_empty_domains_lists(self, mock_tavily_client, mock_settings): """Test web search with empty domain lists.""" mock_settings.TAVILY_API_KEY = "test_api_key" + mock_settings.WEB_SEARCH_PROVIDER = "auto" + mock_settings.WEB_SEARCH_TIMEOUT = 30.0 mock_client_instance = MagicMock() mock_tavily_client.return_value = mock_client_instance From bd173169b759de92cbab2c94c715dab60542e3a0 Mon Sep 17 00:00:00 2001 From: George Pickett <297992784+georgeatparallel@users.noreply.github.com> Date: Tue, 20 Oct 2026 10:35:27 +0000 Subject: [PATCH 2/2] Align MCP cleanup contexts with current lint rules --- prometheus/tools/web_search.py | 26 ++++++++++++++------------ 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/prometheus/tools/web_search.py b/prometheus/tools/web_search.py index 9f1e7a6..6212f6a 100644 --- a/prometheus/tools/web_search.py +++ b/prometheus/tools/web_search.py @@ -181,7 +181,7 @@ async def aweb_search( raise WebSearchToolException("Parallel search timed out") from None except WebSearchToolException: raise - except Exception: + except Exception: # noqa: BLE001 - sanitize SDK exceptions before model consumption. # SDK transport errors can contain request details. Keep the public error concise. raise WebSearchToolException("Parallel Search MCP request failed") from None @@ -189,18 +189,20 @@ async def _parallel_search(self, query: str) -> str: # Each call owns and closes its transport. Do not attach a process/client-wide # conversation ID: analyzer instances are reused across independent issue runs. headers = {"User-Agent": "Prometheus/1.3 (+https://github.com/EuniAI/Prometheus)"} - async with streamablehttp_client("https://search.parallel.ai/mcp", headers=headers) as ( - read, - write, - _, - ): - async with ClientSession( + async with ( + streamablehttp_client("https://search.parallel.ai/mcp", headers=headers) as ( + read, + write, + _, + ), + ClientSession( read, write, client_info=Implementation(name="Prometheus", version="1.3") - ) as session: - await session.initialize() - result = await session.call_tool( - "web_search", {"objective": query, "search_queries": [query]} - ) + ) as session, + ): + await session.initialize() + result = await session.call_tool( + "web_search", {"objective": query, "search_queries": [query]} + ) text = "\n".join(block.text for block in result.content if block.type == "text") if result.isError: raise WebSearchToolException(f"Parallel Search MCP returned an error: {text}")