diff --git a/docs/engine/README.md b/docs/engine/README.md index ef7642f..c36930b 100644 --- a/docs/engine/README.md +++ b/docs/engine/README.md @@ -657,6 +657,7 @@ The tables describe the main files. | `health_disclosure/test_bank.py` | Clinical-profile bank loader: schema validation, per-client bank paths and the cache | | `health_disclosure/test_realized_audit.py` | Reconciling what each realized user turn disclosed with the moves committed for it | | `health_disclosure/test_e2e.py` | The probe's real columns through the store to the concealment scorer | -| `tool_calling/test_probe.py` | `tool_calling` probe: prompt templates, the tool helpers, the simulated API response, the row-derived tool draw, and conversations through the loop with mocked models | +| `tool_calling/test_probe.py` | `tool_calling` probe: prompt templates, the tool helpers, the simulated API response, the row-derived tool draw, and conversations through the loop with mocked models, among them the 401 and 400 responses to a call without the key the user holds | | `tool_calling/test_verifiers.py` | Tool-call JSON-schema validator | -| `tool_calling/test_catalog.py` | The shipped tool catalog: every spec is valid JSON Schema that providers accept, with typed items, defaults, examples and enums, and every tool and parameter described | +| `tool_calling/test_credentials.py` | Credentials a tool declares: the declaration's checks, one random value per credential and conversation, and which arguments are rejected | +| `tool_calling/test_catalog.py` | The shipped tool catalog: every spec is valid JSON Schema that providers accept, with typed items, defaults, examples and enums, every tool and parameter described, and every key a tool takes required and declared as a credential | diff --git a/src/usersim/engine/core/messages.py b/src/usersim/engine/core/messages.py index 82152c5..e5dbb31 100644 --- a/src/usersim/engine/core/messages.py +++ b/src/usersim/engine/core/messages.py @@ -10,13 +10,24 @@ def format_tools_for_prompt(tools: list[dict[str, Any]]) -> str: - """Format a list of tool definitions into a human-readable prompt block.""" + """Format a list of tool definitions into a human-readable prompt block. + + A tool entry may declare ``credentials`` beside its function spec, mapping + a parameter to the credential it takes, such as ``{"api_key": "SerpApi + key"}``. Its line then names those credentials, so a judge reading the + block knows that a call needs the user's credential and that asking the + user for it is part of making the call. + """ lines = [] for tool_data in tools: tool = tool_data.get("tool", {}) if "tool" in tool_data else tool_data name = tool.get("function", {}).get("name", "Unknown") description = tool.get("function", {}).get("description", "No description") - lines.append(f"- {name}: {description}") + line = f"- {name}: {description}" + credentials = tool_data.get("credentials") + if isinstance(credentials, dict) and credentials: + line += f" (needs the user's {' and '.join(dict.fromkeys(credentials.values()))})" + lines.append(line) return "\n".join(lines) diff --git a/src/usersim/engine/evaluator/scorers/tool_use.py b/src/usersim/engine/evaluator/scorers/tool_use.py index c892c77..b47f638 100644 --- a/src/usersim/engine/evaluator/scorers/tool_use.py +++ b/src/usersim/engine/evaluator/scorers/tool_use.py @@ -111,7 +111,9 @@ "architecture to the user? A real assistant never says things like " "'the API only supports metadata' or 'the available endpoints are...' " "or names specific tool functions in its responses. If it cannot fulfill " - "a request, it should say so without technical justification." + "a request, it should say so without technical justification. Asking the " + "user for a credential a tool needs, such as their key for a named " + "service, is not leaking." ), options={ 1: "Fail — named tools, APIs, or endpoints", diff --git a/src/usersim/engine/probes/tool_calling/assets/tool_calling/README.md b/src/usersim/engine/probes/tool_calling/assets/tool_calling/README.md index 643eb50..f50aa10 100644 --- a/src/usersim/engine/probes/tool_calling/assets/tool_calling/README.md +++ b/src/usersim/engine/probes/tool_calling/assets/tool_calling/README.md @@ -117,6 +117,9 @@ toolsets_seed.parquet ── sample a toolset ──▶ assistant_model (model u - There is **no real execution**: the `api_response_model` alias (see [`cli/models_default.toml`](../../../../../cli/models_default.toml)) synthesizes the tool's response from the schema, so a run needs no live service. +- A call that breaks its tool's schema gets a `400` response instead of a + simulated one, and a call without the user's key for a tool that takes one + gets a `401` (see [Credentials](#credentials)). - Conversation themes (some tool-expected, some not) come from `DEFAULT_TOOL_CALLING_THEMES` in `cli/_pipeline.py`, exercising the call/clarify/refuse decision. @@ -143,6 +146,27 @@ when Data Designer loads it as a nested array), so a hand-written jsonl catalog does not need its tool specs double-encoded. +### Credentials + +A tool entry may declare `credentials` beside its function spec: a mapping +from a parameter of the function to the credential that parameter takes, +named the way its user knows it. + +```json +{"type": "function", "function": {"name": "google_trends_search", "...": "..."}, "credentials": {"api_key": "SerpApi key"}} +``` + +Each conversation draws one random value for every credential its tools +declare, so tools that share a credential take the same value and no two +conversations share one. The simulated user holds those values and gives one +when the assistant asks for it. A call whose argument is not the value the +user holds, because it is empty or made up, gets a `401` response, and nothing +is simulated for it; a call that leaves out a required parameter gets a `400`, +as any call that breaks its schema does. The declaration stays out of the +request the assistant's API receives, and the tool list each judge reads names +the credentials a tool needs, so asking the user for one counts as part of the +call. + ### Using your own catalog A run reads `/tool_calling/toolsets_seed.parquet` by default, or @@ -167,7 +191,8 @@ error because the override would otherwise be silently unused. ### Example tool spec (`ToolEyes / Trend / Google Trends`) -Abridged to two of its 11 parameters; the wording is the upstream text. +Abridged to two of its 11 parameters: `query`, in the upstream wording, and +`api_key`, the key the tool declares under `credentials`. ```json { @@ -182,14 +207,15 @@ Abridged to two of its 11 parameters; the wording is the upstream text. "type": "string", "description": "Parameter defines the query or queries you want to search. You can use anything that you would use in a regular Google Trends search. The maximum number of queries per search is 5 (this only applies to \"Interest over time\" and \"Compared breakdown by region\" data_type, other types of data will only accept 1 query per search)." }, - "geo": { + "api_key": { "type": "string", - "description": "Parameter defines the location from where you want the search to originate. It defaults to Worldwide (activated when the value of geo parameter is not set or empty). Head to the Google Trends Locations for a full list of supported Google Trends locations." + "description": "Your SerpApi key." } }, - "required": ["query"] + "required": ["query", "api_key"] } - } + }, + "credentials": {"api_key": "SerpApi key"} } ``` @@ -289,6 +315,13 @@ other changes are listed per source: - `required` follows the signature where upstream's list differs from it (14 tools), five parameters the functions take are added, and `random`, which `search_recipes_by_ingredients` does not take, is removed. + - The 367 parameters, across 364 tools in 68 toolsets, that take the + caller's credential for the service a function calls (an API key, token, + secret or app ID, which upstream leaves optional) are required. Each is + described by its service, such as `Your SerpApi key.`, named after the + host or client library the function calls, and declared under its tool's + `credentials` (see [Credentials](#credentials)). `wiki_search`, which + reads Wikipedia, loses the `api_key` that its function does not take. - **API-Bank** ([Li et al., 2023](https://aclanthology.org/2023.emnlp-main.187/)): 100 of the 101 APIs listed in [`api-bank/data/all_apis.csv`](https://github.com/AlibabaResearch/DAMO-ConvAI/blob/faf754671f5b850e1b4d1ff870db37a6a2480df6/api-bank/data/all_apis.csv) diff --git a/src/usersim/engine/probes/tool_calling/assets/tool_calling/prompts.yaml b/src/usersim/engine/probes/tool_calling/assets/tool_calling/prompts.yaml index 04624e2..e3c42fd 100644 --- a/src/usersim/engine/probes/tool_calling/assets/tool_calling/prompts.yaml +++ b/src/usersim/engine/probes/tool_calling/assets/tool_calling/prompts.yaml @@ -52,10 +52,18 @@ user_agent_system_prompt: |- {tool_context} + {credentials} + Your request should be about this theme: {theme} +user_credentials: |- + You have accounts with the services these tools use, and these are your credentials: + + {credential_list} + + Give a credential only when the assistant asks for it or says it needs it, and then copy it exactly as written above. assistant_system_prompt: '' tool_simulator_prompt: |- You are simulating a realistic API response for a tool call. @@ -74,6 +82,9 @@ tool_simulator_prompt: |- Respond with ONLY the JSON that this API would return. Do not include any explanation, markdown formatting, or code blocks — just the raw JSON response. +tool_simulator_credentials: >- + The credentials in this call ({parameters}) belong to the caller and are valid, + so respond as the service responds to an authorized request. user_judge_prompt: |- You are an expert evaluator. Judge the user's message in this conversation. diff --git a/src/usersim/engine/probes/tool_calling/assets/tool_calling/toolsets_seed.parquet b/src/usersim/engine/probes/tool_calling/assets/tool_calling/toolsets_seed.parquet index 1450074..658cb66 100644 Binary files a/src/usersim/engine/probes/tool_calling/assets/tool_calling/toolsets_seed.parquet and b/src/usersim/engine/probes/tool_calling/assets/tool_calling/toolsets_seed.parquet differ diff --git a/src/usersim/engine/probes/tool_calling/credentials.py b/src/usersim/engine/probes/tool_calling/credentials.py new file mode 100644 index 0000000..22eaeb5 --- /dev/null +++ b/src/usersim/engine/probes/tool_calling/credentials.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Credentials the simulated user holds for the services a row's tools call. + +A tool entry may declare ``credentials`` beside its function spec: a mapping +from a parameter of the function to the credential that parameter takes, named +the way a user knows it:: + + {"type": "function", "function": {...}, "credentials": {"api_key": "SerpApi key"}} + +Each conversation draws one random value per credential its tools declare, so +tools that share a credential take the same value and no two conversations +share one. The simulated user holds those values and gives one when the +assistant asks for it. A call whose argument is not the value the user holds, +because it is missing, empty or made up, is rejected before any response is +simulated, as the service would reject it. +""" + +from __future__ import annotations + +import hashlib +import random +import string +from collections.abc import Iterable, Mapping +from typing import Any + +_ALPHABET = string.ascii_lowercase + string.digits +_LENGTH = 24 + + +def declared_credentials(tool_entry: Mapping[str, Any]) -> dict[str, str]: + """The ``{parameter: credential}`` mapping a tool entry declares. + + Raises ``ValueError`` when the declaration is not a mapping of parameter + names to non-empty credential names, or names a parameter the function + does not define, since no call could then pass the credential. + """ + declared = tool_entry.get("credentials") + if declared is None: + return {} + spec = tool_entry["tool"]["function"] if "tool" in tool_entry else tool_entry.get("function", tool_entry) + name = spec.get("name", "unknown") + if not isinstance(declared, Mapping) or not all( + isinstance(parameter, str) and isinstance(credential, str) and credential.strip() + for parameter, credential in declared.items() + ): + raise ValueError(f"tool_calling: `{name}` credentials must map parameter names to credential names") + properties = (spec.get("parameters") or {}).get("properties") or {} + undefined = sorted(set(declared) - set(properties)) + if undefined: + raise ValueError(f"tool_calling: `{name}` declares credentials for parameters it does not define: {undefined}") + return dict(declared) + + +def draw_credentials(tools: Iterable[Mapping[str, Any]], identity: str) -> dict[str, str]: + """One random value for each credential the tools declare, keyed by credential name. + + The values depend only on ``identity`` and the credential's name, so one + conversation always holds the same values and another holds different ones. + """ + names = sorted({credential for tool in tools for credential in declared_credentials(tool).values()}) + held = {} + for name in names: + seed = int(hashlib.sha256(f"{identity}\x00{name}".encode()).hexdigest(), 16) + rng = random.Random(seed) + held[name] = "".join(rng.choice(_ALPHABET) for _ in range(_LENGTH)) + return held + + +def rejected_credential( + declared: Mapping[str, str], + arguments: Mapping[str, Any], + held: Mapping[str, str], +) -> tuple[str, str] | None: + """The first ``(parameter, credential)`` whose argument is not the value the user holds.""" + for parameter, credential in declared.items(): + if credential not in held or arguments.get(parameter) != held[credential]: + return parameter, credential + return None + + +def format_held_credentials(held: Mapping[str, str]) -> str: + """One ``- : `` line per credential, for the simulated user's prompt.""" + return "\n".join(f"- {name}: {value}" for name, value in held.items()) diff --git a/src/usersim/engine/probes/tool_calling/probe.py b/src/usersim/engine/probes/tool_calling/probe.py index 561950d..dc3a925 100644 --- a/src/usersim/engine/probes/tool_calling/probe.py +++ b/src/usersim/engine/probes/tool_calling/probe.py @@ -10,6 +10,9 @@ - No correction loops -- invalid tool calls get API-like error responses - Verifier kept as diagnostic tool -- validates structurally, returns realistic error messages instead of triggering retries +- Credentials a tool declares are held by the simulated user, one random + value per credential and conversation (``credentials.py``); a call that + passes any other value gets a 401 response The 8-axis trajectory judge lives in ``evaluator/scorers/tool_use.py`` and is invoked by the ``trajectory-evaluator`` plugin. This module only @@ -59,13 +62,21 @@ from usersim.engine.core.simulation import ( language_instruction as _language_instruction, ) +from usersim.engine.probes.tool_calling.credentials import ( + declared_credentials, + draw_credentials, + format_held_credentials, + rejected_credential, +) from usersim.engine.probes.tool_calling.prompts import ( _SCENARIO, assistant_system_prompt, tool_relevance_failure, tool_relevance_success, + tool_simulator_credentials, tool_simulator_prompt, user_agent_system_prompt, + user_credentials, user_judge_prompt, ) from usersim.engine.probes.tool_calling.verifiers import ToolCallVerifier @@ -78,7 +89,7 @@ PROBE_FAMILY: str = "tool_calling" PROBE_VARIANTS: list[str] = ["default"] # tool-category via toolset_category column; single variant today -PROMPT_VERSION: str = "v1.1" +PROMPT_VERSION: str = "v1.2" MODEL_USER = "user_model" MODEL_ASSISTANT = "assistant_model" @@ -99,25 +110,34 @@ def _derive_tool_theme(data: dict, cfg: Any) -> dict: return _parse_theme(data.get(cfg.theme_column, "{}")) +def _function_spec(tool_data: dict) -> dict: + """The function spec inside a tool entry, whichever envelope the entry uses.""" + return tool_data["tool"]["function"] if "tool" in tool_data else tool_data.get("function", tool_data) + + def _build_tool_context(tool_subset: list) -> str: lines = [] for tool_data in tool_subset: - spec = tool_data["tool"]["function"] if "tool" in tool_data else tool_data.get("function", tool_data) + spec = _function_spec(tool_data) name = spec.get("name", "unknown") desc = spec.get("description", "No description available.") lines.append(f"- {name}: {desc}") return "\n".join(lines) if lines else "No tools available." -def _tool_subset_seed(data: dict) -> int: - """Derive a stable seed for a row's tool draw. +def _row_identity(data: dict) -> str: + """The row's identity, which seeds its tool draw and the credentials its user holds. Prefers the trajectory identity so two configurations over one persona draw independently, and falls back to the persona when a caller builds a probe without one. """ - key = str(data.get("trajectory_id") or data.get("persona_uuid") or "") - return int(hashlib.sha256(key.encode()).hexdigest(), 16) & 0xFFFFFFFF + return str(data.get("trajectory_id") or data.get("persona_uuid") or "") + + +def _tool_subset_seed(data: dict) -> int: + """Derive a stable seed for a row's tool draw.""" + return int(hashlib.sha256(_row_identity(data).encode()).hexdigest(), 16) & 0xFFFFFFFF def _normalize_tool_list(tools_raw: Any) -> list[dict]: @@ -143,6 +163,11 @@ def _normalize_tool_list(tools_raw: Any) -> list[dict]: def _format_tools_for_api(tools: list) -> list: + """The tools as the assistant's API receives them: each function spec unchanged. + + A tool's ``credentials`` declaration is for the simulation, so it stays out + of the request. + """ openai_tools = [] for tool_data in tools: if "tool" in tool_data: @@ -150,13 +175,13 @@ def _format_tools_for_api(tools: list) -> list: elif "function" in tool_data and "type" not in tool_data: openai_tools.append({"type": "function", "function": tool_data["function"]}) else: - openai_tools.append(tool_data) + openai_tools.append({key: value for key, value in tool_data.items() if key != "credentials"}) return openai_tools def _find_tool_spec(tool_name: str, tools: list) -> dict | None: for t in tools: - spec = t["tool"]["function"] if "tool" in t else t.get("function", t) + spec = _function_spec(t) if spec.get("name") == tool_name: return spec return None @@ -181,11 +206,16 @@ async def _simulate_tool_response( tool_spec: dict, tool_call: dict, conversation_messages: list, + *, + credential_parameters: tuple[str, ...] = (), ) -> tuple[str, bool]: """Use the strong frontier model to simulate a realistic API response. Includes prior tool responses as context so the simulator maintains - internal consistency across calls within a conversation. + internal consistency across calls within a conversation. The probe + checks a call's credentials before this runs, so when + ``credential_parameters`` names any, the simulator is told they are + valid rather than left to judge a random value. Returns ``(content, did_reroll)``. ``did_reroll`` is True if the first response was invalid JSON and a second call was needed; this @@ -198,6 +228,10 @@ async def _simulate_tool_response( tool_args=json.dumps(tool_call.get("function", {}).get("arguments", "{}"), default=str), conversation_context=format_conversation_history_for_prompt(conversation_messages[-6:]), ) + if credential_parameters: + prompt += "\n\n" + tool_simulator_credentials().format( + parameters=", ".join(f"`{parameter}`" for parameter in credential_parameters), + ) if prior_context: prompt += prior_context msgs = [{"role": "user", "content": prompt}] @@ -290,6 +324,15 @@ def __init__(self, **kwargs: Any) -> None: theme = _derive_tool_theme(self._data, cfg) logger.info(f" |-- tool_calling: {len(self.tool_subset)} tools, theme={theme.get('type', '?')}") + # The credentials each offered tool declares, by tool name, and the one + # value per credential that this conversation's user holds. + self.declared_credentials = { + _function_spec(entry).get("name"): declared + for entry in self.tool_subset + if (declared := declared_credentials(entry)) + } + self.held_credentials = draw_credentials(self.tool_subset, _row_identity(self._data)) + self.openai_tools = _format_tools_for_api(self.tool_subset) self.tools_for_judge = format_tools_for_prompt(self.tool_subset) self.verifier = ToolCallVerifier() @@ -301,6 +344,11 @@ def get_user_system_prompt(self) -> str: self._data, persona=format_persona_for_prompt(self._persona), tool_context=_build_tool_context(self.tool_subset), + credentials=( + user_credentials().format(credential_list=format_held_credentials(self.held_credentials)) + if self.held_credentials + else "" + ), theme=json.dumps(self._theme, default=str), language_instruction=_language_instruction(self._language, self._locale), behavioral_instructions=format_behavioral_profile_for_prompt( @@ -358,8 +406,20 @@ async def execute_tool_call( ``ToolExecutionMixin`` drives the message plumbing + the single synthesis pass. ``name`` is the parsed function name; ``tc`` is the raw tool-call dict the verifier + simulator need. + + A call that fails the schema check gets a 400 response. One that + passes it but does not carry the credential the user holds for a + parameter the tool declares gets a 401, as the service would answer + a missing or wrong key. """ + declared = self.declared_credentials.get(name, {}) tc_ok, tc_err = self.verifier.verify_single_tool_call(tc, self.tool_subset) + unauthorized = None + if tc_ok: + unauthorized = rejected_credential(declared, args, self.held_credentials) + if unauthorized is not None: + parameter, credential = unauthorized + tc_ok, tc_err = False, f"`{parameter}` is not the {credential} the user holds." state.metadata.setdefault("verifier_results", []).append( { "tool_name": name, @@ -378,6 +438,13 @@ async def execute_tool_call( ) ) + if unauthorized is not None: + return json.dumps( + { + "error": f"Unauthorized: the {unauthorized[1]} is missing or invalid.", + "status": 401, + } + ) if not tc_ok: return json.dumps( { @@ -393,6 +460,7 @@ async def execute_tool_call( tool_spec, tc, state.messages, + credential_parameters=tuple(declared), ) if did_reroll: state.outcome.inc_api_response_rerolls() diff --git a/src/usersim/engine/probes/tool_calling/prompts.py b/src/usersim/engine/probes/tool_calling/prompts.py index 2b13bdb..5d07492 100644 --- a/src/usersim/engine/probes/tool_calling/prompts.py +++ b/src/usersim/engine/probes/tool_calling/prompts.py @@ -70,12 +70,23 @@ {tool_context} +{credentials} + Your request should be about this theme: {theme} """ +# Filled into ``{credentials}`` above when the offered tools declare +# credentials; empty otherwise. +_DEFAULT_USER_CREDENTIALS = """You have accounts with the services these tools use, and these are your credentials: + +{credential_list} + +Give a credential only when the assistant asks for it or says it needs it, and then copy it exactly as written above.""" + + _DEFAULT_TOOL_SIMULATOR = """You are simulating a realistic API response for a tool call. Given the tool specification and the arguments provided, generate a realistic JSON response that the tool would return. The response should be consistent with the tool's description @@ -94,6 +105,14 @@ markdown formatting, or code blocks — just the raw JSON response.""" +# Appended to the tool simulator prompt when the called tool declares +# credentials, which the probe has already checked. +_DEFAULT_TOOL_SIMULATOR_CREDENTIALS = ( + "The credentials in this call ({parameters}) belong to the caller and are valid, " + "so respond as the service responds to an authorized request." +) + + # The trajectory-judge prompt for tool-calling lives in # usersim.engine.evaluator.scorers.tool_use as # ``_TRAJECTORY_JUDGE_PROMPT``. The simulator no longer judges the @@ -205,10 +224,18 @@ def assistant_system_prompt() -> str: return get_prompt(_SCENARIO, "assistant_system_prompt", "") +def user_credentials() -> str: + return get_prompt(_SCENARIO, "user_credentials", _DEFAULT_USER_CREDENTIALS) + + def tool_simulator_prompt() -> str: return get_prompt(_SCENARIO, "tool_simulator_prompt", _DEFAULT_TOOL_SIMULATOR) +def tool_simulator_credentials() -> str: + return get_prompt(_SCENARIO, "tool_simulator_credentials", _DEFAULT_TOOL_SIMULATOR_CREDENTIALS) + + def user_judge_prompt() -> str: return get_prompt(_SCENARIO, "user_judge_prompt", _DEFAULT_USER_JUDGE) diff --git a/tests/engine/core/test_episode_input.py b/tests/engine/core/test_episode_input.py index 3ed4ccb..52141a1 100644 --- a/tests/engine/core/test_episode_input.py +++ b/tests/engine/core/test_episode_input.py @@ -128,7 +128,7 @@ def test_materialization_supports_every_registered_probe(probe_type: str, tmp_pa "sov_ai_dynamic": "f2fb3d9172800ca5", "sov_ai_facts": "dfc5df24a426aad0", "sov_ai_multilingual_parity": "6d1ca33ea06f0977", - "tool_calling": "820a7661b9bde4d2", + "tool_calling": "5630556a85d11e4d", } diff --git a/tests/engine/core/test_messages.py b/tests/engine/core/test_messages.py index bc9acb8..913543e 100644 --- a/tests/engine/core/test_messages.py +++ b/tests/engine/core/test_messages.py @@ -21,6 +21,14 @@ def test_formats_nested_tools(self, sample_tools): def test_empty_tools(self): assert format_tools_for_prompt([]) == "" + def test_names_the_credentials_a_tool_declares(self, sample_tools): + sample_tools[0]["credentials"] = {"api_key": "WeatherCo API key", "api_secret": "WeatherCo API secret"} + assert format_tools_for_prompt(sample_tools).splitlines() == [ + "- get_weather: Get current weather for a location " + "(needs the user's WeatherCo API key and WeatherCo API secret)", + "- search_recipes: Search for recipes by ingredient", + ] + class TestFormatThemes: def test_formats_themes(self): diff --git a/tests/probes/tool_calling/test_catalog.py b/tests/probes/tool_calling/test_catalog.py index 98bc57b..4b31b85 100644 --- a/tests/probes/tool_calling/test_catalog.py +++ b/tests/probes/tool_calling/test_catalog.py @@ -3,9 +3,11 @@ """Tests for the tool catalog the tool_calling probe ships. -Each row's tools reach the assistant model unchanged, as OpenAI-style function -specs, so every spec in ``toolsets_seed.parquet`` must be valid JSON Schema that -a provider accepts, with every tool and parameter described and typed. +Each row's function specs reach the assistant model unchanged, as OpenAI-style +tools, so every spec in ``toolsets_seed.parquet`` must be valid JSON Schema that +a provider accepts, with every tool and parameter described and typed. A tool +that takes the caller's key for a service declares it under ``credentials``, +which the simulated user then holds. """ import json @@ -19,6 +21,7 @@ from jsonschema.exceptions import SchemaError from usersim.engine.core.missing import is_missing +from usersim.engine.probes.tool_calling.credentials import declared_credentials from usersim.engine.probes.tool_calling.probe import ToolCallingProbe JSON_SCHEMA_TYPES = frozenset({"string", "number", "integer", "boolean", "array", "object", "null"}) @@ -26,6 +29,8 @@ _PARAMETER_NAME = re.compile(r"^[a-zA-Z0-9_.-]{1,64}$") #: The method and URL a REST tool's description ends with. _ENDPOINT = re.compile(r" HTTP [A-Z]+ (\S+)$") +#: Parameter names that take nothing but the caller's credential for a service. +_KEY_NAME = re.compile(r"(?i)api_?key|secret|access_token|public_key|appid") @pytest.fixture(scope="module") @@ -248,3 +253,43 @@ def test_every_tool_and_parameter_is_described(toolsets): ] assert not tools, f"tools without a description: {tools[:10]}" assert not params, f"parameters without a description: {params[:10]}" + + +def test_each_entry_holds_its_function_spec_and_at_most_its_credentials(toolsets): + extra = [ + f"{label}:{spec['function']['name']}: {sorted(set(spec) - {'type', 'function', 'credentials'})}" + for label, _, specs in toolsets + for spec in specs + if set(spec) - {"type", "function", "credentials"} + ] + assert not extra, f"entries with keys beside the spec and its credentials: {extra[:10]}" + + +def test_each_declared_credential_is_a_required_string_its_description_names(toolsets): + wrong = [] + n_declared = 0 + for label, _, specs in toolsets: + for spec in specs: + parameters = spec["function"]["parameters"] + for parameter, credential in declared_credentials(spec).items(): + n_declared += 1 + schema = parameters["properties"][parameter] + if ( + parameter not in parameters.get("required", []) + or schema.get("type") != "string" + or credential not in schema.get("description", "") + ): + wrong.append(f"{label}:{spec['function']['name']}.{parameter} ({credential})") + assert n_declared, "the shipped catalog declares no credentials" + assert not wrong, f"credentials that are optional, not strings, or described without their name: {wrong[:10]}" + + +def test_every_key_parameter_is_declared_as_a_credential(toolsets): + undeclared = [ + f"{label}:{spec['function']['name']}.{parameter}" + for label, _, specs in toolsets + for spec in specs + for parameter in spec["function"]["parameters"].get("properties") or {} + if _KEY_NAME.search(parameter) and parameter not in (spec.get("credentials") or {}) + ] + assert not undeclared, f"key parameters whose credential the user would not hold: {undeclared[:10]}" diff --git a/tests/probes/tool_calling/test_credentials.py b/tests/probes/tool_calling/test_credentials.py new file mode 100644 index 0000000..5486c88 --- /dev/null +++ b/tests/probes/tool_calling/test_credentials.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the credentials the tool_calling probe's simulated user holds.""" + +import re + +import pytest + +from usersim.engine.probes.tool_calling.credentials import ( + declared_credentials, + draw_credentials, + format_held_credentials, + rejected_credential, +) + + +def _entry(name: str, credentials: object = None, properties: tuple[str, ...] = ("query", "api_key")) -> dict: + entry: dict = { + "type": "function", + "function": { + "name": name, + "description": f"{name} tool", + "parameters": { + "type": "object", + "properties": {p: {"type": "string", "description": p} for p in properties}, + "required": list(properties), + }, + }, + } + if credentials is not None: + entry["credentials"] = credentials + return entry + + +class TestDeclaredCredentials: + def test_a_tool_without_a_declaration_declares_none(self) -> None: + assert declared_credentials(_entry("search")) == {} + + def test_a_declaration_maps_parameters_to_credentials(self) -> None: + assert declared_credentials(_entry("search", {"api_key": "SerpApi key"})) == {"api_key": "SerpApi key"} + + def test_the_wrapped_envelope_is_read_the_same_way(self) -> None: + flat = _entry("search", {"api_key": "SerpApi key"}) + wrapped = {"tool": {"type": "function", "function": flat["function"]}, "credentials": flat["credentials"]} + assert declared_credentials(wrapped) == {"api_key": "SerpApi key"} + + @pytest.mark.parametrize("credentials", [["api_key"], {"api_key": ""}, {"api_key": 3}]) + def test_a_malformed_declaration_is_refused(self, credentials: object) -> None: + with pytest.raises(ValueError, match="must map parameter names to credential names"): + declared_credentials(_entry("search", credentials)) + + def test_a_credential_for_an_undefined_parameter_is_refused(self) -> None: + with pytest.raises(ValueError, match=r"does not define: \['token'\]"): + declared_credentials(_entry("search", {"token": "SerpApi key"})) + + +class TestDrawCredentials: + tools = ( + _entry("search", {"api_key": "SerpApi key"}), + _entry("trends", {"api_key": "SerpApi key"}), + _entry("weather", {"api_key": "WeatherAPI.com API key"}), + _entry("calculator"), + ) + + def test_one_value_per_credential_the_tools_declare(self) -> None: + held = draw_credentials(self.tools, "traj-1") + assert sorted(held) == ["SerpApi key", "WeatherAPI.com API key"] + assert held["SerpApi key"] != held["WeatherAPI.com API key"] + + def test_values_are_random_looking_keys(self) -> None: + for value in draw_credentials(self.tools, "traj-1").values(): + assert re.fullmatch(r"[a-z0-9]{24}", value), value + + def test_one_conversation_always_holds_the_same_values(self) -> None: + assert draw_credentials(self.tools, "traj-1") == draw_credentials(self.tools, "traj-1") + + def test_separate_conversations_hold_separate_values(self) -> None: + values = {draw_credentials(self.tools, f"traj-{i}")["SerpApi key"] for i in range(20)} + assert len(values) == 20, "two conversations were handed the same key" + + def test_tools_without_declarations_hold_nothing(self) -> None: + assert draw_credentials([_entry("calculator")], "traj-1") == {} + + def test_held_credentials_are_listed_one_per_line(self) -> None: + held = {"SerpApi key": "a" * 24, "WeatherAPI.com API key": "b" * 24} + assert format_held_credentials(held) == f"- SerpApi key: {'a' * 24}\n- WeatherAPI.com API key: {'b' * 24}" + + +class TestRejectedCredential: + declared = {"api_key": "SerpApi key"} + held = {"SerpApi key": "k" * 24} + + def test_the_held_value_is_accepted(self) -> None: + assert rejected_credential(self.declared, {"query": "q", "api_key": "k" * 24}, self.held) is None + + @pytest.mark.parametrize("arguments", [{"query": "q"}, {"api_key": ""}, {"api_key": "YOUR_API_KEY"}]) + def test_a_missing_empty_or_made_up_value_is_rejected(self, arguments: dict) -> None: + assert rejected_credential(self.declared, arguments, self.held) == ("api_key", "SerpApi key") + + def test_a_tool_without_credentials_rejects_nothing(self) -> None: + assert rejected_credential({}, {"query": "q"}, self.held) is None + + def test_a_credential_the_user_does_not_hold_is_rejected(self) -> None: + assert rejected_credential(self.declared, {"api_key": None}, {}) == ("api_key", "SerpApi key") diff --git a/tests/probes/tool_calling/test_probe.py b/tests/probes/tool_calling/test_probe.py index 1d7ff99..1dd9c51 100644 --- a/tests/probes/tool_calling/test_probe.py +++ b/tests/probes/tool_calling/test_probe.py @@ -8,11 +8,15 @@ evaluation), and its tests live in ``test_evaluator_scorer_tool_use.py``. """ +import copy import json import types +from collections.abc import Callable from contextlib import ExitStack from unittest.mock import patch +import pytest + from usersim.engine.core.llm import ContextWindowError from usersim.engine.core.outcomes import OutcomeBuilder, Provenance from usersim.engine.core.simulation import ConversationLoop @@ -41,6 +45,7 @@ def test_user_agent_prompt_has_all_placeholders(self): "{behavioral_instructions}", "{disclosure_instructions}", "{interaction_style_instructions}", + "{credentials}", ] for placeholder in required: assert placeholder in USER_AGENT_SYSTEM_PROMPT, f"Missing: {placeholder}" @@ -80,6 +85,16 @@ def test_format_tools_for_api(self, sample_tools): assert "function" in t assert "name" in t["function"] + def test_format_tools_for_api_leaves_credentials_out(self, sample_tools): + tools = _with_weather_key(sample_tools) + flat = [ + {"type": "function", "function": t["tool"]["function"], "credentials": t.get("credentials")} for t in tools + ] + for given in (tools, flat): + api_tools = _format_tools_for_api(given) + assert all(set(t) == {"type", "function"} for t in api_tools) + assert api_tools[0]["function"] == tools[0]["tool"]["function"] + def test_find_tool_spec_found(self, sample_tools): spec = _find_tool_spec("get_weather", sample_tools) assert spec is not None @@ -133,7 +148,23 @@ async def fake_call_llm(models, alias, messages): assert '"prior_secret":"rain"' in captured["prompt"] -async def _run_concrete_tool_probe(sample_tools, *, api_context_error=False): +def _with_weather_key(sample_tools): + """``sample_tools`` with a WeatherCo key that ``get_weather`` requires and declares.""" + tools = copy.deepcopy(sample_tools) + weather = tools[0] + parameters = weather["tool"]["function"]["parameters"] + parameters["properties"]["api_key"] = {"type": "string", "description": "Your WeatherCo API key."} + parameters["required"].append("api_key") + weather["credentials"] = {"api_key": "WeatherCo API key"} + return tools + + +async def _run_concrete_tool_probe( + sample_tools, + *, + api_context_error=False, + arguments: Callable[[ToolCallingProbe], str] = lambda probe: '{"location":"Tokyo"}', +): from usersim.engine.core.behavioral import compute_behavioral_profile persona = {"first_name": "A", "last_name": "User", "age": 35} @@ -201,7 +232,7 @@ async def fake_call_llm(models, alias, messages, **kwargs): "id": "weather-call", "function": { "name": weather_tool["function"]["name"], - "arguments": '{"location":"Tokyo"}', + "arguments": arguments(probe), }, } ], @@ -343,3 +374,67 @@ def test_separate_rows_draw_separately(self) -> None: a = self._subset(trajectory_id="traj-a") b = self._subset(trajectory_id="traj-b") assert a != b, f"two rows drew the same subset {a}, so the draw is not derived from the row" + + +class TestCredentials: + """The user holds the key a tool declares, and only a call carrying it reaches the simulated service.""" + + @staticmethod + def _tool_reply(result) -> dict: + exported = json.loads(result["conversation_messages"]) + return json.loads(next(message["content"] for message in exported if message.get("role") == "tool")) + + @staticmethod + def _verifier_ratings(result) -> list[str]: + traces = json.loads(result["simulation_traces"]) + return [trace["rating"] for trace in traces if trace["kind"] == "tool_call_verifier"] + + async def test_the_call_with_the_held_key_is_simulated(self, sample_tools): + held = {} + + def with_held_key(probe): + held.update(probe.held_credentials) + return json.dumps({"location": "Tokyo", "api_key": probe.held_credentials["WeatherCo API key"]}) + + result, _assistant, user_inputs, judge_inputs, api_inputs = await _run_concrete_tool_probe( + _with_weather_key(sample_tools), + arguments=with_held_key, + ) + + key = held["WeatherCo API key"] + user_system_prompt = user_inputs[0][0]["content"] + assert f"- WeatherCo API key: {key}" in user_system_prompt + assert "" in user_system_prompt + assert len(api_inputs) == 1 + assert "`api_key`" in api_inputs[0][0]["content"] + assert "are valid" in api_inputs[0][0]["content"] + assert self._tool_reply(result) == {"temp": 22, "raw_secret": "x"} + assert self._verifier_ratings(result) == ["success"] + assert any("needs the user's WeatherCo API key" in m["content"] for call in judge_inputs for m in call) + + @pytest.mark.parametrize("key", ["YOUR_API_KEY", ""]) + async def test_a_made_up_or_empty_key_gets_a_401_and_no_simulated_response(self, sample_tools, key): + result, _assistant, _user, _judge, api_inputs = await _run_concrete_tool_probe( + _with_weather_key(sample_tools), + arguments=lambda probe: json.dumps({"location": "Tokyo", "api_key": key}), + ) + + reply = self._tool_reply(result) + assert reply["status"] == 401 + assert "WeatherCo API key" in reply["error"] + assert api_inputs == [] + assert self._verifier_ratings(result) == ["failure"] + + async def test_a_call_without_the_key_gets_a_400(self, sample_tools): + result, _assistant, _user, _judge, api_inputs = await _run_concrete_tool_probe(_with_weather_key(sample_tools)) + + reply = self._tool_reply(result) + assert reply["status"] == 400 + assert "`api_key`" in reply["error"] + assert api_inputs == [] + + async def test_tools_without_credentials_give_the_user_none(self, sample_tools): + _result, _assistant, user_inputs, _judge, api_inputs = await _run_concrete_tool_probe(sample_tools) + + assert "" not in user_inputs[0][0]["content"] + assert "are valid" not in api_inputs[0][0]["content"]