diff --git a/ms_agent/tools/base.py b/ms_agent/tools/base.py index 18e66d175..66692e969 100644 --- a/ms_agent/tools/base.py +++ b/ms_agent/tools/base.py @@ -1,10 +1,20 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +import math from abc import abstractmethod from omegaconf import DictConfig -from typing import Any, Dict +from typing import Any, Dict, Optional from ms_agent.utils.constants import DEFAULT_OUTPUT_DIR +#: Return this from :attr:`ToolBase.max_output_chars` to declare that a tool +#: bounds its own output and must never be cut by the generic truncator. +SELF_MANAGED_OUTPUT = math.inf + +#: Where to keep text from when an oversized output IS cut generically. +TRUNCATE_KEEP_HEAD = 'head' +TRUNCATE_KEEP_TAIL = 'tail' +TRUNCATE_KEEP_BOTH = 'both' + class ToolBase: """The base class for all tools. @@ -19,6 +29,30 @@ def __init__(self, config): self.output_dir = getattr(self.config, 'output_dir', DEFAULT_OUTPUT_DIR) + @property + def max_output_chars(self) -> Optional[float]: + """Model-facing character budget for this tool's output. + + * ``None`` (default) — use the global ``MAX_TOOL_OUTPUT_LEN``. + * :data:`SELF_MANAGED_OUTPUT` — the tool guarantees its own bound + (paging, spilling to disk, …); never truncate it generically. + * a number — this tool's own budget, used instead of the global one. + + Override in a subclass to opt in. Declaring a budget is a promise about + SHAPE as much as size: a tool that returns structured data should keep + itself under budget so the generic cut never has to run. + """ + return None + + @property + def truncate_keep(self) -> str: + """Which end survives when this tool's output IS cut generically. + + ``'head'`` (a command's first output is the useful part), ``'tail'`` + (a long run whose verdict is last), or ``'both'`` (default). + """ + return TRUNCATE_KEEP_BOTH + def exclude_func(self, tool_config: DictConfig): if tool_config is not None: self.exclude_functions = getattr(tool_config, 'exclude', []) diff --git a/ms_agent/tools/search/tavily/fetcher.py b/ms_agent/tools/search/tavily/fetcher.py index d4ed4927a..58e41e63f 100644 --- a/ms_agent/tools/search/tavily/fetcher.py +++ b/ms_agent/tools/search/tavily/fetcher.py @@ -32,11 +32,9 @@ def __init__( include_favicon: bool = False, include_usage: bool = False, ): - key = api_key or os.getenv('TAVILY_API_KEY') - if not key: - raise ValueError( - 'TAVILY_API_KEY required for tavily_extract fetcher') - self._api_key = key + # Keyless is a supported mode here too (Tavily serves /extract without + # credentials under the same header as /search); see tavily/search.py. + self._api_key = api_key or os.getenv('TAVILY_API_KEY') or '' self._extract_depth = extract_depth self._format = format self._timeout = max(1.0, min(60.0, float(timeout))) @@ -52,7 +50,6 @@ def fetch(self, Extract one URL. Optional ``query`` enables chunk reranking (more relevant raw_content). """ body: Dict[str, Any] = { - 'api_key': self._api_key, 'urls': [url], 'extract_depth': self._extract_depth, 'format': self._format, @@ -64,10 +61,18 @@ def fetch(self, if query: body['query'] = query body['chunks_per_source'] = self._chunks_per_source + # Omitted when empty: any api_key in the body overrides the keyless + # header (see TavilySearchRequest.to_api_body). + if self._api_key: + body['api_key'] = self._api_key try: + from ms_agent.tools.search.tavily.search import KEYLESS_HEADER data = post_json( - TAVILY_EXTRACT_URL, body, timeout=self._timeout + 30.0) + TAVILY_EXTRACT_URL, + body, + timeout=self._timeout + 30.0, + headers=(dict(KEYLESS_HEADER) if not self._api_key else {})) except Exception as e: logger.warning(f'Tavily extract failed for {url[:80]}: {e}') return '', { diff --git a/ms_agent/tools/search/tavily/http.py b/ms_agent/tools/search/tavily/http.py index e2ad2cb85..39412d0f8 100644 --- a/ms_agent/tools/search/tavily/http.py +++ b/ms_agent/tools/search/tavily/http.py @@ -1,35 +1,124 @@ # Copyright (c) ModelScope Contributors. All rights reserved. """Minimal HTTP JSON client for Tavily REST API (stdlib only).""" import json -from typing import Any, Dict +from typing import Any, Dict, Optional from urllib.error import HTTPError, URLError from urllib.request import Request, urlopen +class TavilyHTTPError(RuntimeError): + """A Tavily call that failed, with the pieces a caller can act on. + + Plain ``RuntimeError`` forced every caller to re-parse the message to tell + "you are out of quota, ask the user for a key" from "that host is down". + Keyless mode makes that distinction routine rather than exceptional — the + free tier is a small hourly bucket — so the parts travel as fields: + ``status`` (HTTP code, None for transport failures), ``code`` (Tavily's own + machine-readable ``error.code``, e.g. ``hourly_cap_reached``) and + ``retry_after`` seconds when the response carried one. + """ + + def __init__(self, + message: str, + *, + status: Optional[int] = None, + code: str = '', + retry_after: Optional[int] = None, + detail: Any = None): + super().__init__(message) + self.status = status + self.code = code + self.retry_after = retry_after + self.detail = detail + + @property + def is_quota(self) -> bool: + """Out of quota — retryable later, and fixable now with an API key.""" + return self.status == 429 or self.code in ('hourly_cap_reached', + 'rate_limit_exceeded') + + @property + def is_auth(self) -> bool: + return self.status in (401, 403) + + +def _ssl_context(): + """A verifying TLS context that works on interpreters with no CA store. + + ``urlopen`` uses the interpreter's default store, which is empty in some + virtualenvs (``ssl.get_default_verify_paths().cafile is None`` — measured on + the WebUI backend's venv, where every Tavily call died with + CERTIFICATE_VERIFY_FAILED). certifi is already an indirect dependency there; + when it is missing we hand back None so urlopen behaves exactly as before. + Never disables verification. + """ + try: + import certifi + import ssl + return ssl.create_default_context(cafile=certifi.where()) + except Exception: + return None + + +def _parse_error_body(raw: str) -> Any: + try: + return json.loads(raw) if raw else {} + except json.JSONDecodeError: + return {'raw': raw} + + +def _dig_error(detail: Any) -> tuple: + """``(code, message, retry_after)`` out of Tavily's error envelope. + + Two shapes are in the wild: ``{"error": {"code", "message", + "retry_after_seconds"}}`` (keyless quota) and ``{"detail": {"error": ...}}`` + (auth). Anything else degrades to empty strings rather than raising while + already handling an error. + """ + code = message = '' + retry_after = None + node = detail + if isinstance(node, dict) and isinstance(node.get('detail'), dict): + node = node['detail'] + if isinstance(node, dict): + err = node.get('error') + if isinstance(err, dict): + code = str(err.get('code') or '') + message = str(err.get('message') or '') + ra = err.get('retry_after_seconds') + if isinstance(ra, (int, float)): + retry_after = int(ra) + elif isinstance(err, str): + message = err + return code, message, retry_after + + def post_json( url: str, body: Dict[str, Any], *, timeout: float = 120.0, + headers: Optional[Dict[str, str]] = None, ) -> Dict[str, Any]: """ POST JSON and parse JSON response. + ``headers`` is merged over the defaults — that is how keyless mode is + selected (``X-Tavily-Access-Mode: keyless``). + Raises: - RuntimeError: on HTTP errors or invalid JSON (includes Tavily error body). + TavilyHTTPError: on HTTP errors or invalid JSON (carries Tavily's own + error code / retry-after so callers can tell quota from outage). """ data = json.dumps(body, ensure_ascii=False).encode('utf-8') - req = Request( - url, - data=data, - method='POST', - headers={ - 'Content-Type': 'application/json', - 'Accept': 'application/json', - }, - ) + merged = { + 'Content-Type': 'application/json', + 'Accept': 'application/json', + } + merged.update(headers or {}) + req = Request(url, data=data, method='POST', headers=merged) try: - with urlopen(req, timeout=timeout) as resp: + with urlopen(req, timeout=timeout, context=_ssl_context()) as resp: raw = resp.read().decode('utf-8', errors='replace') if not raw.strip(): return {} @@ -40,10 +129,24 @@ def post_json( err_body = e.read().decode('utf-8', errors='replace') except Exception: pass - try: - detail = json.loads(err_body) if err_body else {} - except json.JSONDecodeError: - detail = {'raw': err_body} - raise RuntimeError(f'Tavily HTTP {e.code}: {detail}') from e + detail = _parse_error_body(err_body) + code, message, retry_after = _dig_error(detail) + if retry_after is None: + header_value = None + try: + header_value = e.headers.get('retry-after') + except Exception: + pass + if header_value: + try: + retry_after = int(float(header_value)) + except (TypeError, ValueError): + retry_after = None + raise TavilyHTTPError( + f'Tavily HTTP {e.code}: {message or detail}', + status=e.code, + code=code, + retry_after=retry_after, + detail=detail) from e except URLError as e: - raise RuntimeError(f'Tavily network error: {e}') from e + raise TavilyHTTPError(f'Tavily network error: {e}') from e diff --git a/ms_agent/tools/search/tavily/schema.py b/ms_agent/tools/search/tavily/schema.py index 112a41839..2d9d16451 100644 --- a/ms_agent/tools/search/tavily/schema.py +++ b/ms_agent/tools/search/tavily/schema.py @@ -33,7 +33,6 @@ class TavilySearchRequest: def to_api_body(self, api_key: str) -> Dict[str, Any]: n = max(0, min(20, int(self.max_results))) body: Dict[str, Any] = { - 'api_key': api_key, 'query': self.query, 'max_results': n, 'search_depth': self.search_depth, @@ -64,6 +63,13 @@ def to_api_body(self, api_key: str) -> Dict[str, Any]: body['exclude_domains'] = list(self.exclude_domains)[:150] if self.country: body['country'] = self.country + # Sent LAST and only when non-empty. Keyless mode (the + # X-Tavily-Access-Mode header) is overridden by any api_key present in + # the body: measured 2026-08-20, an empty string is tolerated but a + # non-empty one is validated and a bogus value 401s. Omitting the field + # entirely is the only shape that is unambiguous in both modes. + if api_key: + body['api_key'] = api_key return body diff --git a/ms_agent/tools/search/tavily/search.py b/ms_agent/tools/search/tavily/search.py index 38c3c69cb..49485ac28 100644 --- a/ms_agent/tools/search/tavily/search.py +++ b/ms_agent/tools/search/tavily/search.py @@ -15,6 +15,16 @@ TAVILY_SEARCH_URL = 'https://api.tavily.com/search' +#: Tavily serves /search and /extract without credentials when this header is +#: present (https://docs.tavily.com/documentation/keyless). Responses are +#: identical to keyed ones — same parameters, same result shape — but the quota +#: is a small sliding hourly bucket rather than the free tier's monthly credits. +#: Measured 2026-08-20: refill is roughly one request per 60-90s, and exhaustion +#: is a clean HTTP 429 (`error.code: hourly_cap_reached`) carrying Retry-After. +#: It exists so the framework works on first run with nothing configured; a key +#: always takes precedence when one is available. +KEYLESS_HEADER = {'X-Tavily-Access-Mode': 'keyless'} + class TavilySearch(SearchEngine): """ @@ -32,24 +42,36 @@ def __init__( api_key: Optional[str] = None, request_timeout: float = 120.0, ): - key = api_key or os.getenv('TAVILY_API_KEY') - if not key: - raise ValueError( - 'TAVILY_API_KEY must be set in environment or web_search.tavily_api_key' - ) - self._api_key = key + # No key is a supported mode, not an error: without one we fall back to + # Tavily's keyless tier so a fresh install can search out of the box. + # Constructing this used to raise, which WebSearchTool.connect() caught + # and turned into "engine unavailable" — the reason an unconfigured + # framework silently had no web search at all. + self._api_key = api_key or os.getenv('TAVILY_API_KEY') or '' self._request_timeout = float(request_timeout) + @property + def keyless(self) -> bool: + return not self._api_key + + def _headers(self) -> dict: + return dict(KEYLESS_HEADER) if self.keyless else {} + def search(self, search_request: TavilySearchRequest) -> TavilySearchResult: body = search_request.to_api_body(self._api_key) - try: - data = post_json( - TAVILY_SEARCH_URL, body, timeout=self._request_timeout) - except Exception as e: - raise RuntimeError(f'Tavily search failed: {e}') from e + # Deliberately unguarded: TavilyHTTPError carries the quota/auth fields + # the tool layer needs to tell the agent WHY a search failed. This used + # to be wrapped in a bare RuntimeError, which erased them. + data = post_json( + TAVILY_SEARCH_URL, + body, + timeout=self._request_timeout, + headers=self._headers()) safe_args = {k: v for k, v in body.items() if k != 'api_key'} - safe_args['api_key'] = '' + if self._api_key: + safe_args['api_key'] = '' + safe_args['access_mode'] = 'keyless' if self.keyless else 'api_key' return TavilySearchResult( query=search_request.query, arguments=safe_args, diff --git a/ms_agent/tools/search/websearch_tool.py b/ms_agent/tools/search/websearch_tool.py index aab578184..a0f2e3e4e 100644 --- a/ms_agent/tools/search/websearch_tool.py +++ b/ms_agent/tools/search/websearch_tool.py @@ -25,6 +25,17 @@ MAX_FETCH_CHARS = int(os.getenv('MAX_FETCH_CHARS', 100000)) +#: Model-facing budget for one web_search result payload. Matches the historical +#: global tool-output cap, so what reaches the context is the same size it always +#: was — the difference is that it now arrives as valid JSON (spill trims the +#: bodies) instead of being bisected by the generic truncator. +DEFAULT_SEARCH_OUTPUT_CHARS = int( + os.getenv('WEB_SEARCH_MAX_OUTPUT_CHARS', 20000)) + +#: Room reserved inside the budget for everything spill does NOT count: urls, +#: titles, scores, the spill digest, and JSON syntax/indentation. +SEARCH_JSON_OVERHEAD = 6000 + def default_per_url_fetch_timeout_s( fetch_timeout: float, @@ -617,13 +628,30 @@ def __init__(self, config, **kwargs): getattr(tool_cfg, 'summarization_timeout', 90.0) or 90.0) if tool_cfg else 90.0 + # This tool's model-facing budget, declared to the generic truncator via + # `max_output_chars` below. It is ONE knob: spill aims under it, and the + # final serialization is checked against it, so the JSON the model gets + # is always both bounded and parseable. + self._max_output_chars = int( + getattr(tool_cfg, 'max_output_chars', DEFAULT_SEARCH_OUTPUT_CHARS) + or DEFAULT_SEARCH_OUTPUT_CHARS) if tool_cfg else ( + DEFAULT_SEARCH_OUTPUT_CHARS) + # Large payload spill (write bodies to disk; keep JSON small) self._spill_enabled = bool( getattr(tool_cfg, 'spill_large_results', True)) if tool_cfg else True + # Derived from the budget, not an independent number. It used to default + # to 120000 while the generic truncator cut at 20000, so every payload + # between those two was left whole by spill and then bisected into + # invalid JSON — the mechanism meant to protect the payload never ran. + # The gap is the room the non-content fields need (urls, titles, scores, + # JSON syntax), which spill does not count. + _spill_default = max(2000, + self._max_output_chars - SEARCH_JSON_OVERHEAD) self._spill_max_inline_chars = int( - getattr(tool_cfg, 'spill_max_inline_chars', 120000) - or 120000) if tool_cfg else 120000 + getattr(tool_cfg, 'spill_max_inline_chars', _spill_default) + or _spill_default) if tool_cfg else _spill_default self._spill_subdir = str( getattr(tool_cfg, 'spill_subdir', '.ms_agent/web_search') or '.ms_agent/web_search') if tool_cfg else '.ms_agent/web_search' @@ -659,6 +687,65 @@ def __init__(self, config, **kwargs): } self._summary_usage_model: str = '' + @property + def max_output_chars(self) -> int: + """This tool's own budget (see ToolBase.max_output_chars). + + Declared rather than inherited because the generic cut is destructive + here: the payload is JSON, and a notice spliced into its middle makes it + unparseable. `_bounded_json` below keeps every response under this + number, so in practice the generic path never has to run. + """ + return self._max_output_chars + + def _bounded_json(self, response: Dict[str, Any]) -> str: + """Serialize a search response, guaranteed under budget AND valid JSON. + + Spill normally gets there first by moving bodies to disk. This is the + backstop for what spill cannot help with — a hundred results whose + previews alone overflow, or a run with no ``output_dir`` to spill into. + It sheds content in order of dispensability (previews, then whole rows) + and re-serializes each time, so the result is always parseable; the + alternative, letting the generic truncator cut it, is not. + """ + out = _json_dumps(response) + budget = self._max_output_chars + if len(out) <= budget: + return out + + rows = response.get('results') + if not isinstance(rows, list) or not rows: + return out # nothing to shed; caller's budget declaration handles it + + # 1) Drop the bulky per-row text, keeping the evidence (url/title). + trimmed = [] + for row in rows: + r = dict(row) if isinstance(row, dict) else {'value': row} + for field in ('content', 'summary', 'abstract', 'chunks'): + if r.get(field): + r[field] = '' + r['content_note'] = ( + 'Body omitted to stay within the tool output budget; ' + 'open the url, or read content_path if present.') + trimmed.append(r) + response = {**response, 'results': trimmed, 'body_omitted': True} + out = _json_dumps(response) + if len(out) <= budget: + return out + + # 2) Still over: keep fewer rows, and say how many were dropped. + keep = len(trimmed) + while keep > 1 and len(out) > budget: + keep = max(1, keep // 2) + response = { + **response, + 'results': trimmed[:keep], + 'count': keep, + 'results_omitted': len(trimmed) - keep, + } + out = _json_dumps(response) + return out + async def connect(self) -> None: """Initialize search engines and content fetcher.""" for engine_type in self._engine_types: @@ -977,6 +1064,52 @@ async def _bounded_fetch(url: str) -> Dict[str, Any]: tasks = [_bounded_fetch(url) for url in urls] return await asyncio.gather(*tasks) + @staticmethod + def _search_error_payload(engine_type: str, + exc: Exception) -> Dict[str, Any]: + """A failed search, described so the MODEL can act on it. + + The distinction that matters is "retry later / tell the user to add a + key" versus "this is broken" — a bare stringified exception leaves the + model to guess, and it usually guesses "no results". Tavily's keyless + tier makes the quota case routine, so it gets an explicit remedy. + """ + payload: Dict[str, Any] = { + 'engine': engine_type, + 'message': str(exc), + } + status = getattr(exc, 'status', None) + code = getattr(exc, 'code', '') or '' + retry_after = getattr(exc, 'retry_after', None) + if status is not None: + payload['http_status'] = status + if code: + payload['provider_code'] = code + if retry_after is not None: + payload['retry_after_seconds'] = retry_after + + if getattr(exc, 'is_quota', False): + payload['kind'] = 'quota_exceeded' + wait = (f' Retry in about {retry_after}s' + if retry_after is not None else ' Retry shortly') + payload['remedy'] = ( + f'The {engine_type} free/keyless quota is exhausted.{wait}, or ' + f'tell the user they can remove the limit by adding a ' + f'{engine_type} API key in Settings -> Search. Do NOT report ' + f'this as "no results found" — the search never ran.') + elif getattr(exc, 'is_auth', False): + payload['kind'] = 'auth_failed' + payload['remedy'] = ( + f'The configured {engine_type} API key was rejected. Tell the ' + f'user to check or replace it in Settings -> Search.') + else: + payload['kind'] = 'search_failed' + payload['remedy'] = ( + f'The {engine_type} search request did not complete. It may be ' + f'a transient network problem — retrying once is reasonable; ' + f'do not present this as an empty result set.') + return payload + def _do_search( self, engine_type: str, engine: SearchEngine, engine_cls: Type[SearchEngine], @@ -1005,8 +1138,17 @@ def _do_search( extra = result.extra_response_fields() return rows, extra except Exception as e: + # Reported, never swallowed. Returning ([], {}) here made every + # failure — expired key, rate limit, DNS outage — reach the model as + # "No search results found.", so it would tell the user their query + # matched nothing and move on. With the keyless Tavily tier that is + # actively harmful: running out of quota is an ordinary event whose + # fix (add an API key) only the user can apply, and they never heard + # about it. `extra` carries the diagnosis to _execute_search. logger.error(f'Search failed ({engine_type}): {e}') - return [], {} + return [], { + '__error__': self._search_error_payload(engine_type, e) + } async def _execute_search(self, engine_type: str, tool_args: Dict[str, Any]) -> str: @@ -1060,6 +1202,22 @@ async def _execute_search(self, engine_type: str, self._executor, self._do_search, engine_type, engine, engine_cls, tool_args) + # A failed search and a search that matched nothing are different facts + # and must not share a response shape (see _search_error_payload). + failure = (tavily_extra or {}).pop('__error__', None) + if failure: + out_error: Dict[str, Any] = { + 'status': 'error', + 'query': query, + 'engine': engine_type, + 'count': 0, + 'results': [], + **failure, + } + if tavily_extra: + out_error.update(tavily_extra) + return _json_dumps(out_error) + if not search_results: out_empty: Dict[str, Any] = { 'status': 'ok', @@ -1399,7 +1557,8 @@ async def _execute_search(self, engine_type: str, f'web_search spill failed (returning full inline JSON): {e}' ) - return _json_dumps(response) + # Bounded here, not by the generic truncator: this payload is JSON. + return self._bounded_json(response) async def fetch_page(self, url: str) -> str: """Fetch and parse a single web page.""" diff --git a/ms_agent/tools/tool_manager.py b/ms_agent/tools/tool_manager.py index 39a0d942a..df6b9a228 100644 --- a/ms_agent/tools/tool_manager.py +++ b/ms_agent/tools/tool_manager.py @@ -45,6 +45,40 @@ MAX_CONCURRENT_TOOLS = int(os.getenv('MAX_CONCURRENT_TOOLS', 20)) +def _truncate_tool_output(response: str, tool_ins: Any) -> str: + """Bound one tool result, honouring what the tool declared about itself. + + The global cap exists so a runaway result cannot eat the context window. + But cutting a payload in half and splicing a notice into the gap assumes + prose: done to JSON it lands inside a string literal and the whole result + stops parsing, so the model gets neither the data nor an error. Tools that + bound themselves therefore opt out (or set their own budget) via + ``ToolBase.max_output_chars``; everything else keeps the previous + behaviour exactly. + """ + budget = getattr(tool_ins, 'max_output_chars', None) + if budget is None: + budget = int(os.getenv('MAX_TOOL_OUTPUT_LEN', 20000)) + if budget == math.inf: # self-managed: the tool guarantees its own bound + return response + budget = int(budget) + if budget <= 0 or len(response) <= budget: + return response + + keep = str(getattr(tool_ins, 'truncate_keep', 'both') or 'both') + notice = ( + f'\n\n...[SYSTEM: Output truncated, {len(response)} chars total, ' + f'showing {{shown}}]...\n\n') + if keep == 'head': + return response[:budget] + notice.format(shown=f'first {budget} chars') + if keep == 'tail': + return notice.format(shown=f'last {budget} chars') + response[-budget:] + half = budget // 2 + return (response[:half] + + notice.format(shown=f'first and last {half} chars') + + response[-half:]) + + def _tool_on(config, name: str) -> bool: """Whether a builtin tool should load: its ``tools.`` key is present AND not explicitly disabled via ``tools..enabled: false``. Default @@ -653,14 +687,11 @@ async def single_call_tool(self, tool_info: ToolCall): tool_args=call_args), timeout=wait_sec) - # Truncate excessively long tool outputs to prevent context window explosion - max_len = int(os.getenv('MAX_TOOL_OUTPUT_LEN', 20000)) - if isinstance(response, str) and len(response) > max_len: - half = max_len // 2 - trunc_notice = ( - f'\n\n...[SYSTEM: Output truncated, {len(response)} chars total, ' - f'showing first and last {half} chars]...\n\n') - response = response[:half] + trunc_notice + response[-half:] + # Truncate excessively long tool outputs to prevent context + # window explosion — unless the tool declared its own budget or + # said it bounds itself (see ToolBase.max_output_chars). + if isinstance(response, str): + response = _truncate_tool_output(response, tool_ins) if (self.mcp_success_handler is not None and tool_ins is self.servers): diff --git a/requirements/framework.txt b/requirements/framework.txt index c808cb446..de60818d2 100644 --- a/requirements/framework.txt +++ b/requirements/framework.txt @@ -17,4 +17,5 @@ pytz pyyaml requests rich +socksio typing_extensions diff --git a/requirements/research.txt b/requirements/research.txt index 693301a94..67ce2d3fd 100644 --- a/requirements/research.txt +++ b/requirements/research.txt @@ -14,4 +14,3 @@ Pillow python-dotenv requests rich -socksio diff --git a/tests/search/test_tavily_keyless.py b/tests/search/test_tavily_keyless.py new file mode 100644 index 000000000..9c73aff37 --- /dev/null +++ b/tests/search/test_tavily_keyless.py @@ -0,0 +1,175 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Tavily's keyless tier, and telling the agent WHY a search failed. + +Keyless access (``X-Tavily-Access-Mode: keyless``) exists so the framework can +search on first run with nothing configured. Two facts measured against +api.tavily.com on 2026-08-20 shape these tests: + +* a non-empty ``api_key`` in the body OVERRIDES the keyless header — a bogus one + answers 401 — so the field must be absent, not empty, when going keyless; +* the quota is a small sliding hourly bucket (refill ~1 request/60-90s), so + running out is an ORDINARY event whose only fix (add an API key) belongs to + the user. It must therefore reach them, and the old code turned every failure + into "No search results found." +""" +import json +from io import BytesIO +from urllib.error import HTTPError + +import pytest +from ms_agent.tools.search.tavily import http as tavily_http +from ms_agent.tools.search.tavily.schema import TavilySearchRequest +from ms_agent.tools.search.tavily.search import KEYLESS_HEADER, TavilySearch +from ms_agent.tools.search.websearch_tool import WebSearchTool + +# The exact 429 envelope api.tavily.com returns once the keyless bucket is dry. +REAL_QUOTA_BODY = { + 'error': { + 'code': 'hourly_cap_reached', + 'message': ('You reached the hourly keyless Tavily limit. To continue ' + 'immediately, pay via x402 agentic payment or sign up at ' + 'https://tavily.com for a Tavily API key.'), + 'next_actions': [], + 'window': 'hour', + 'retry_after_seconds': 62, + } +} +# Auth failures nest one level deeper than quota ones. +REAL_AUTH_BODY = {'detail': {'error': 'Unauthorized: missing or invalid API key.'}} + + +class _Hdrs(dict): + + def get(self, key, default=None): + return dict.get(self, key.lower(), default) + + +def _raise_http(status, body, headers=None): + def boom(*args, **kwargs): + raise HTTPError('https://api.tavily.com/search', status, 'err', + _Hdrs(headers or {}), + BytesIO(json.dumps(body).encode())) + + return boom + + +@pytest.fixture() +def no_env_key(monkeypatch): + monkeypatch.delenv('TAVILY_API_KEY', raising=False) + + +# --------------------------------------------------------------- body shape + +def test_api_key_is_omitted_when_absent(): + """Present-but-empty would be tolerated; absent is the only unambiguous + shape, since ANY non-empty value makes Tavily validate it instead of + honouring the keyless header.""" + body = TavilySearchRequest(query='q', max_results=1).to_api_body('') + assert 'api_key' not in body + assert body['query'] == 'q' + + +def test_api_key_is_sent_when_present(): + body = TavilySearchRequest(query='q', max_results=1).to_api_body('tvly-K') + assert body['api_key'] == 'tvly-K' + + +# ------------------------------------------------------------ mode selection + +def test_missing_key_selects_keyless_instead_of_raising(no_env_key): + """Regression: this used to raise ValueError, which WebSearchTool.connect() + caught and logged — leaving an unconfigured install with no web search and + no explanation.""" + engine = TavilySearch() + assert engine.keyless is True + assert engine._headers() == KEYLESS_HEADER + + +def test_explicit_key_disables_keyless(no_env_key): + engine = TavilySearch(api_key='tvly-K') + assert engine.keyless is False + assert engine._headers() == {} + + +def test_env_key_disables_keyless(monkeypatch): + monkeypatch.setenv('TAVILY_API_KEY', 'tvly-FROM-ENV') + assert TavilySearch().keyless is False + + +def test_keyless_header_is_sent_and_key_field_absent(no_env_key, monkeypatch): + seen = {} + + def fake_post(url, body, *, timeout=120.0, headers=None): + seen['body'] = body + seen['headers'] = headers + return {'results': []} + + monkeypatch.setattr('ms_agent.tools.search.tavily.search.post_json', + fake_post) + TavilySearch().search(TavilySearchRequest(query='q', max_results=1)) + assert seen['headers'] == KEYLESS_HEADER + assert 'api_key' not in seen['body'] + + +# ------------------------------------------------------- error classification + +def test_quota_error_is_parsed_with_code_and_retry_after(monkeypatch): + monkeypatch.setattr( + tavily_http, 'urlopen', + _raise_http(429, REAL_QUOTA_BODY, {'retry-after': '62'})) + with pytest.raises(tavily_http.TavilyHTTPError) as ctx: + tavily_http.post_json('https://api.tavily.com/search', {'query': 'q'}) + err = ctx.value + assert err.status == 429 + assert err.code == 'hourly_cap_reached' + assert err.retry_after == 62 + assert err.is_quota and not err.is_auth + + +def test_auth_error_is_recognised_through_the_detail_envelope(monkeypatch): + monkeypatch.setattr(tavily_http, 'urlopen', + _raise_http(401, REAL_AUTH_BODY)) + with pytest.raises(tavily_http.TavilyHTTPError) as ctx: + tavily_http.post_json('https://api.tavily.com/search', {'query': 'q'}) + assert ctx.value.is_auth and not ctx.value.is_quota + + +def test_retry_after_falls_back_to_the_header(monkeypatch): + """Some responses carry the header but no retry_after_seconds field.""" + monkeypatch.setattr( + tavily_http, 'urlopen', + _raise_http(429, {'error': {'code': 'hourly_cap_reached'}}, + {'retry-after': '90'})) + with pytest.raises(tavily_http.TavilyHTTPError) as ctx: + tavily_http.post_json('https://api.tavily.com/search', {'query': 'q'}) + assert ctx.value.retry_after == 90 + + +# ------------------------------------------------ what the MODEL is told + +def _payload(status, code='', retry_after=None): + err = tavily_http.TavilyHTTPError( + 'boom', status=status, code=code, retry_after=retry_after) + return WebSearchTool._search_error_payload('tavily', err) + + +def test_quota_payload_names_the_remedy_and_forbids_no_results(): + p = _payload(429, 'hourly_cap_reached', 62) + assert p['kind'] == 'quota_exceeded' + assert p['retry_after_seconds'] == 62 + assert 'API key' in p['remedy'] + # The whole point: the model must not narrate this as an empty result set. + assert 'no results found' in p['remedy'].lower() + + +def test_auth_payload_points_at_the_key(): + p = _payload(401) + assert p['kind'] == 'auth_failed' + assert 'key' in p['remedy'].lower() + + +def test_transport_failure_is_neither_quota_nor_auth(): + err = tavily_http.TavilyHTTPError('Tavily network error: timed out') + p = WebSearchTool._search_error_payload('tavily', err) + assert p['kind'] == 'search_failed' + assert 'empty result' in p['remedy'] diff --git a/tests/tools/test_tool_budget.py b/tests/tools/test_tool_budget.py new file mode 100644 index 000000000..01a1c84f1 --- /dev/null +++ b/tests/tools/test_tool_budget.py @@ -0,0 +1,162 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Per-tool output budgets, and why the generic cut needed an opt-out. + +The generic truncator keeps the head and the tail of an oversized result and +splices a notice into the gap. That is right for prose and destructive for +structure: applied to JSON the notice lands inside a string literal, so the +payload stops parsing and the model receives neither the data nor an error. + +Measured on web_search (2026-08-25): an 84,429-char result was cut at 20,000 +and reached the model — and the UI — as invalid JSON, which is why search +results silently rendered with no sources. The protocol under test is the fix: +a tool declares its own budget (or declares that it bounds itself) and the +generic path leaves it alone. +""" +import json +import math + +import pytest +from ms_agent.tools.base import SELF_MANAGED_OUTPUT, ToolBase +from ms_agent.tools.tool_manager import _truncate_tool_output + + +class _Tool: + """Stands in for a tool instance; only the declared attrs are read.""" + + def __init__(self, budget=None, keep='both'): + if budget is not None: + self.max_output_chars = budget + self.truncate_keep = keep + + +# ----------------------------------------------------------- default behaviour + +def test_undeclared_tool_keeps_the_global_cap(monkeypatch): + """Tools that say nothing must behave exactly as before the protocol.""" + monkeypatch.setenv('MAX_TOOL_OUTPUT_LEN', '100') + out = _truncate_tool_output('x' * 500, _Tool()) + assert len(out) < 500 + assert 'Output truncated' in out + assert '500 chars total' in out + + +def test_short_output_is_untouched(monkeypatch): + monkeypatch.setenv('MAX_TOOL_OUTPUT_LEN', '100') + assert _truncate_tool_output('hello', _Tool()) == 'hello' + + +def test_base_class_default_is_the_global_cap(): + """ToolBase must not silently exempt everything — the cap is a safety net.""" + assert ToolBase.max_output_chars.fget(object()) is None + assert ToolBase.truncate_keep.fget(object()) == 'both' + + +# ------------------------------------------------------------------- opt-outs + +def test_self_managed_tool_is_never_cut(monkeypatch): + monkeypatch.setenv('MAX_TOOL_OUTPUT_LEN', '10') + payload = 'y' * 5000 + assert _truncate_tool_output(payload, _Tool(SELF_MANAGED_OUTPUT)) == payload + assert SELF_MANAGED_OUTPUT == math.inf + + +def test_per_tool_budget_overrides_the_global_one(monkeypatch): + monkeypatch.setenv('MAX_TOOL_OUTPUT_LEN', '10') # would shred it + out = _truncate_tool_output('z' * 900, _Tool(1000)) + assert out == 'z' * 900 # under ITS budget, so untouched + + +def test_budget_below_output_still_cuts(): + out = _truncate_tool_output('z' * 900, _Tool(100)) + assert 'Output truncated' in out + + +# ------------------------------------------------------------------ direction + +def test_keep_head_preserves_the_beginning(): + out = _truncate_tool_output('A' * 100 + 'B' * 100, _Tool(100, keep='head')) + assert out.startswith('A' * 100) + assert 'B' not in out.split('[SYSTEM')[0] + + +def test_keep_tail_preserves_the_end(): + out = _truncate_tool_output('A' * 100 + 'B' * 100, _Tool(100, keep='tail')) + assert out.endswith('B' * 100) + + +def test_keep_both_is_the_default_shape(): + out = _truncate_tool_output('A' * 100 + 'B' * 100, _Tool(100)) + assert out.startswith('A' * 50) and out.endswith('B' * 50) + + +# ------------------------------------------------- the regression, end to end + +def _search_payload(n_rows: int, body_chars: int) -> dict: + return { + 'status': 'ok', + 'query': '今日新闻', + 'engine': 'tavily', + 'count': n_rows, + 'results': [{ + 'url': f'https://example.com/{i}', + 'title': f'result {i}', + 'content': '正' * body_chars, + } for i in range(n_rows)], + } + + +def test_generic_cut_would_corrupt_json(): + """Pins the failure this protocol exists to prevent.""" + raw = json.dumps(_search_payload(10, 5000), ensure_ascii=False, indent=2) + cut = _truncate_tool_output(raw, _Tool(20000)) + with pytest.raises(json.JSONDecodeError): + json.loads(cut) + + +def test_web_search_declares_a_budget_and_stays_under_it(): + from ms_agent.tools.search.websearch_tool import WebSearchTool + + tool = WebSearchTool.__new__(WebSearchTool) + tool._max_output_chars = 20000 + + payload = _search_payload(10, 5000) # ~50k chars of bodies + out = tool._bounded_json(payload) + + assert len(out) <= tool.max_output_chars + parsed = json.loads(out) # THE point: still parseable + assert parsed['results'], 'evidence rows must survive' + assert parsed['results'][0]['url'].startswith('https://') + # And the generic path is now a no-op for it. + assert _truncate_tool_output(out, tool) == out + + +def test_bounded_json_sheds_rows_only_after_bodies(): + from ms_agent.tools.search.websearch_tool import WebSearchTool + + tool = WebSearchTool.__new__(WebSearchTool) + tool._max_output_chars = 20000 + # Bodies alone overflow; dropping them should be enough to fit every row. + out = json.loads(tool._bounded_json(_search_payload(8, 4000))) + assert len(out['results']) == 8, 'rows kept when shedding bodies suffices' + assert out.get('body_omitted') is True + + +def test_bounded_json_drops_rows_when_urls_alone_overflow(): + from ms_agent.tools.search.websearch_tool import WebSearchTool + + tool = WebSearchTool.__new__(WebSearchTool) + tool._max_output_chars = 1200 + out = json.loads(tool._bounded_json(_search_payload(200, 50))) + assert out['results_omitted'] > 0 + assert len(out['results']) < 200 + assert out['count'] == len(out['results']) + + +def test_small_payload_is_returned_verbatim(): + from ms_agent.tools.search.websearch_tool import WebSearchTool + + tool = WebSearchTool.__new__(WebSearchTool) + tool._max_output_chars = 20000 + payload = _search_payload(2, 10) + out = tool._bounded_json(payload) + assert json.loads(out) == payload # no shedding, no markers