diff --git a/docsgpt/agents/tool_executor.py b/docsgpt/agents/tool_executor.py index 94186b79..7bbde1f4 100644 --- a/docsgpt/agents/tool_executor.py +++ b/docsgpt/agents/tool_executor.py @@ -14,6 +14,7 @@ from docsgpt.agents.default_tools import ( synthesized_default_tools, ) from docsgpt import tracing +from docsgpt.agents.tool_pins import iter_parameters, llm_fills, resolve_arguments from docsgpt.agents.tools.tool_action_parser import ToolActionParser from docsgpt.agents.tools.tool_manager import ToolManager from docsgpt.guardrails.types import Stage as GuardrailStage, resolve_tool_result @@ -761,17 +762,22 @@ class ToolExecutor: ) return result - def _build_tool_parameters(self, action: Dict) -> Dict: + def _build_tool_parameters(self, action: Dict, hidden: Optional[set] = None) -> Dict: + """The JSON schema the model sees for ``action``. + + Parameters the model does not fill (fixed values) and ``hidden`` ones + (values the connection fixes) are left out, so the model is never + asked for them. + """ params = {"type": "object", "properties": {}, "required": []} - for param_type in ["query_params", "headers", "body", "parameters"]: - if param_type in action and action[param_type].get("properties"): - for k, v in action[param_type]["properties"].items(): - if v.get("filled_by_llm", True): - params["properties"][k] = { - key: value for key, value in v.items() if key not in ("filled_by_llm", "value", "required") - } - if v.get("required", False): - params["required"].append(k) + for _section, k, v in iter_parameters(action): + if not llm_fills(v) or (hidden and k in hidden): + continue + params["properties"][k] = { + key: value for key, value in v.items() if key not in ("filled_by_llm", "value", "required") + } + if v.get("required", False): + params["required"].append(k) return params def _guardrail_tool_result(self, result: Any, tool_name: str, action_name: str) -> Any: @@ -1393,23 +1399,10 @@ class ToolExecutor: tool_call_data["access"] = action_access(tool_data.get("name"), action_data) - query_params, headers, body, parameters = {}, {}, {}, {} - param_types = { - "query_params": query_params, - "headers": headers, - "body": body, - "parameters": parameters, - } - - for param_type, target_dict in param_types.items(): - if param_type in action_data and action_data[param_type].get("properties"): - for param, details in action_data[param_type]["properties"].items(): - if param not in call_args and "value" in details and details["value"]: - target_dict[param] = details["value"] - for param, value in call_args.items(): - for param_type, target_dict in param_types.items(): - if param_type in action_data and param in action_data[param_type].get("properties", {}): - target_dict[param] = value + # Fixed values win over whatever the model sent for the same key. + sections = resolve_arguments(action_data, call_args) + query_params, headers = sections["query_params"], sections["headers"] + body, parameters = sections["body"], sections["parameters"] # Load tool (with caching) from docsgpt.connectors.service import ConnectionUnavailable diff --git a/docsgpt/agents/tool_pins.py b/docsgpt/agents/tool_pins.py new file mode 100644 index 00000000..f473f175 --- /dev/null +++ b/docsgpt/agents/tool_pins.py @@ -0,0 +1,290 @@ +"""Fixed ("pinned") tool parameters. + +Every parameter of a stored tool action carries ``filled_by_llm`` and +``value``. A parameter with ``filled_by_llm`` false is hidden from the model; +when it also has a value, that value is *pinned*: it is sent on every call and +the model can neither see nor override it. An empty string means "no value": +the parameter is left out of the call (optional OpenAPI parameters are stored +that way), while ``0`` and ``false`` are real values. + +The run-time rules live in :func:`resolve_arguments`; the helpers below +validate and apply pin changes coming from the API. +""" + +from __future__ import annotations + +import copy +from typing import Any, Iterable, Mapping, Optional + +#: The places an action keeps parameter schemas: ``parameters`` for ordinary +#: tools and MCP, the other three for ``api_tool`` requests. +PARAM_SECTIONS = ("query_params", "headers", "body", "parameters") + +_NUMERIC = {"integer": int, "number": float} + + +def _properties(action: Mapping, section: str) -> dict: + block = action.get(section) + if not isinstance(block, dict): + return {} + properties = block.get("properties") + return properties if isinstance(properties, dict) else {} + + +def iter_parameters(action: Mapping) -> Iterable[tuple[str, str, dict]]: + """``(section, name, details)`` for every parameter schema of ``action``.""" + for section in PARAM_SECTIONS: + for name, details in _properties(action, section).items(): + if isinstance(details, dict): + yield section, name, details + + +def llm_fills(details: Mapping) -> bool: + """Whether the model fills this parameter (the default when unset).""" + return bool(details.get("filled_by_llm", True)) + + +def has_value(details: Mapping) -> bool: + """Whether a parameter carries a value: anything but a missing, None or empty string.""" + value = details.get("value") + return value is not None and value != "" + + +def is_pinned(details: Mapping) -> bool: + """Whether the parameter is hidden from the model and always sent with its value.""" + return not llm_fills(details) and has_value(details) + + +def pinned_names(action: Mapping) -> list[str]: + """Names of the parameters of ``action`` that carry a pinned value.""" + return [name for _section, name, details in iter_parameters(action) if is_pinned(details)] + + +def resolve_arguments( + action: Mapping, + llm_arguments: Mapping[str, Any], + connection_pins: Optional[Mapping[str, Any]] = None, +) -> dict[str, dict]: + """The arguments a call runs with, per parameter section. + + Only parameters the action declares are passed. For each one: + + * a value the connection fixes (``connection_pins``, e.g. Telegram's + default chat) always wins; + * a parameter hidden from the model takes its stored value, or is left + out when it has none; whatever the model sent for it is ignored, so a + model that names the key anyway (a mistake, or a prompt injection) + changes nothing; + * any other parameter takes the model's value, falling back to its + stored default. + + Args: + action: The stored action (``parameters`` / ``query_params`` / + ``headers`` / ``body`` schemas with ``filled_by_llm`` and ``value``). + llm_arguments: The arguments the model sent. + connection_pins: Parameter name to a value the connection fixes. + + Returns: + Section name to ``{parameter: value}`` for every section in + :data:`PARAM_SECTIONS` (empty dicts included). + """ + connection_pins = connection_pins or {} + resolved: dict[str, dict] = {section: {} for section in PARAM_SECTIONS} + for section, name, details in iter_parameters(action): + target = resolved[section] + if name in connection_pins: + target[name] = connection_pins[name] + elif not llm_fills(details): + if has_value(details): + target[name] = details["value"] + elif name in llm_arguments: + target[name] = llm_arguments[name] + elif has_value(details): + target[name] = details["value"] + return resolved + + +def coerce_value(details: Mapping, value: Any) -> Any: + """Check a value for a pin against the parameter's declared type. + + Strings typed into a form are converted for ``integer``, ``number`` and + ``boolean`` parameters. + + Raises: + ValueError: The value does not fit the type, or is empty. + """ + if value is None or value == "": + raise ValueError("A fixed value cannot be empty") + kind = details.get("type") + if isinstance(kind, list): + kind = next((k for k in kind if k != "null"), None) + if kind in _NUMERIC: + if isinstance(value, bool): + raise ValueError(f"Expected a {kind}") + try: + number = _NUMERIC[kind](value) + except (TypeError, ValueError): + raise ValueError(f"Expected a {kind}") from None + if kind == "integer" and isinstance(value, float) and not value.is_integer(): + raise ValueError("Expected an integer") + return number + if kind == "boolean": + if isinstance(value, bool): + return value + if isinstance(value, str) and value.strip().lower() in ("true", "false"): + return value.strip().lower() == "true" + raise ValueError("Expected true or false") + if kind == "array": + if not isinstance(value, list): + raise ValueError("Expected a list") + return value + if kind == "object": + if not isinstance(value, dict): + raise ValueError("Expected an object") + return value + if kind in (None, "string"): + if not isinstance(value, (str, int, float, bool)): + raise ValueError("Expected text") + return str(value) if kind == "string" else value + return value + + +def set_pins(action: Mapping, pins: Mapping[str, Any]) -> dict: + """Return ``action`` with ``pins`` applied. + + Args: + action: A stored action. + pins: Parameter name to the value to always use, or None to let the + model decide again. + + Raises: + ValueError: A parameter does not exist on the action, or a value does + not fit its type. + """ + updated = copy.deepcopy(dict(action)) + known = {name: details for _section, name, details in iter_parameters(updated)} + for name, value in pins.items(): + details = known.get(name) + if details is None: + raise ValueError(f"Unknown parameter: {name}") + if value is None: + details["filled_by_llm"] = True + details["value"] = "" + else: + details["value"] = coerce_value(details, value) + details["filled_by_llm"] = False + return updated + + +class PinChangeRefused(PermissionError): + """Someone other than the tool's owner tried to change a fixed value.""" + + +_ACTION_KEYS = ("active", "require_approval", "description") +_PARAM_KEYS = ("filled_by_llm", "value", "description") + + +def _pin_state(details: Mapping) -> tuple: + return (llm_fills(details), details.get("value") if has_value(details) else None) + + +def merge_submitted_actions(stored: list, submitted: Any, *, may_change_pins: bool) -> list: + """Validate a full list of actions sent by a client against the stored ones. + + The stored actions are the schema: a client can switch actions on and off, + change approval and descriptions and (the owner only) fixed values, but + cannot add actions or parameters, or change a parameter's type. + + Args: + stored: The tool's stored actions. + submitted: What the client sent (untrusted). + may_change_pins: False for a team editor: any change to + ``filled_by_llm`` or ``value`` is refused. + + Returns: + The actions to store: every stored action, updated from the + submission. + + Raises: + ValueError: The submission is malformed or names an unknown action or + parameter, or a fixed value does not fit its parameter. + PinChangeRefused: ``may_change_pins`` is False and a fixed value changed. + """ + if not isinstance(submitted, list): + raise ValueError("actions must be a list") + by_name = {a.get("name"): a for a in stored if isinstance(a, dict) and a.get("name")} + seen: set = set() + updates: dict[str, dict] = {} + for entry in submitted: + if not isinstance(entry, dict) or not isinstance(entry.get("name"), str): + raise ValueError("Each action needs a name") + name = entry["name"] + if name not in by_name: + raise ValueError(f"Unknown action: {name}") + if name in seen: + raise ValueError(f"Duplicate action: {name}") + seen.add(name) + updates[name] = entry + + merged = [] + for action in stored: + if not isinstance(action, dict) or action.get("name") not in updates: + merged.append(action) + continue + entry = updates[action["name"]] + result = copy.deepcopy(action) + for key in _ACTION_KEYS: + if key in entry: + value = entry[key] + if key == "description" and not isinstance(value, str): + raise ValueError("description must be text") + result[key] = bool(value) if key != "description" else value + for section in PARAM_SECTIONS: + sent = _properties(entry, section) + if not sent: + continue + properties = _properties(result, section) + for param, sent_details in sent.items(): + details = properties.get(param) + if details is None or not isinstance(details, dict): + raise ValueError(f"Unknown parameter {param} on {action['name']}") + if not isinstance(sent_details, dict): + raise ValueError(f"Parameter {param} must be an object") + before = _pin_state(details) + for key in _PARAM_KEYS: + if key in sent_details: + details[key] = sent_details[key] + details["filled_by_llm"] = llm_fills(details) + if _pin_state(details) == before: + continue + if not may_change_pins: + raise PinChangeRefused("Only the tool's owner can change fixed values") + if is_pinned(details): + details["value"] = coerce_value(details, details["value"]) + merged.append(result) + return merged + + +def carry_pins_between(previous: Optional[list], fresh: list) -> list: + """:func:`carry_pins` for whole action lists, matching actions by name.""" + old = {a.get("name"): a for a in previous or [] if isinstance(a, dict)} + return [carry_pins(old.get(a.get("name")), a) if isinstance(a, dict) else a for a in fresh] + + +def carry_pins(previous: Optional[Mapping], fresh: Mapping) -> dict: + """Copy fixed values from an action's old copy onto its re-discovered one. + + Parameters that still exist keep ``filled_by_llm`` and ``value``; removed + ones are gone with the new schema. + """ + result = copy.deepcopy(dict(fresh)) + if not previous: + return result + old = {(section, name): details for section, name, details in iter_parameters(previous)} + for section, name, details in iter_parameters(result): + before = old.get((section, name)) + if before is None: + continue + details["filled_by_llm"] = llm_fills(before) + details["value"] = before.get("value", "") + return result diff --git a/tests/agents/test_tool_pins.py b/tests/agents/test_tool_pins.py new file mode 100644 index 00000000..383262da --- /dev/null +++ b/tests/agents/test_tool_pins.py @@ -0,0 +1,240 @@ +"""Fixed ("pinned") tool parameters: the model never sees or overrides them.""" + +from unittest.mock import Mock + +import pytest + +from docsgpt.agents.tool_executor import ToolExecutor + + +def _action(**properties): + return { + "name": "telegram_send_message", + "description": "Send a message", + "active": True, + "parameters": {"type": "object", "properties": properties}, + } + + +def _llm(description="", **extra): + return {"type": "string", "description": description, "filled_by_llm": True, "value": "", **extra} + + +def _pinned(value, type_="string"): + return {"type": type_, "description": "", "filled_by_llm": False, "value": value} + + +def _run(executor, tools_dict, call_args, monkeypatch): + monkeypatch.setattr( + "docsgpt.agents.tool_executor.ToolActionParser", + lambda _cls, **kw: Mock(parse_args=Mock(return_value=("t1", "telegram_send_message", call_args))), + ) + call = Mock() + call.name = "telegram_send_message" + call.id = "c1" + call.arguments = "{}" + gen = executor.execute(tools_dict, call, "MockLLM") + while True: + try: + next(gen) + except StopIteration as stop: + return stop.value + + +def _tools(action): + return { + "t1": { + "id": "00000000-0000-0000-0000-000000000001", + "name": "telegram", + "config": {}, + "actions": [action], + } + } + + +@pytest.mark.unit +class TestPinnedValuesAtRunTime: + def test_llm_cannot_override_a_pinned_value(self, mock_tool_manager, monkeypatch): + """A model that sends the pinned key anyway (a mistake, or a prompt + injection naming another chat) still sends to the pinned chat.""" + tool = mock_tool_manager.load_tool.return_value + action = _action(text=_llm(), chat_id=_pinned("111")) + _run(ToolExecutor(user="u"), _tools(action), {"text": "hi", "chat_id": "666"}, monkeypatch) + tool.execute_action.assert_called_once_with("telegram_send_message", text="hi", chat_id="111") + + @pytest.mark.parametrize("value", [0, False]) + def test_falsy_pins_are_honoured(self, mock_tool_manager, monkeypatch, value): + tool = mock_tool_manager.load_tool.return_value + action = _action(text=_llm(), limit=_pinned(value, "integer")) + _run(ToolExecutor(user="u"), _tools(action), {"text": "hi", "limit": 50}, monkeypatch) + tool.execute_action.assert_called_once_with("telegram_send_message", text="hi", limit=value) + + def test_hidden_parameter_without_a_value_is_omitted_even_if_the_llm_sends_it( + self, mock_tool_manager, monkeypatch, + ): + """An empty fixed value means "leave it out" (an OpenAPI optional + parameter); the model was never shown it, so its value is dropped.""" + tool = mock_tool_manager.load_tool.return_value + action = _action(text=_llm(), chat_id=_pinned("")) + _run(ToolExecutor(user="u"), _tools(action), {"text": "hi", "chat_id": "666"}, monkeypatch) + tool.execute_action.assert_called_once_with("telegram_send_message", text="hi") + + def test_llm_filled_parameter_keeps_its_default_when_omitted(self, mock_tool_manager, monkeypatch): + tool = mock_tool_manager.load_tool.return_value + action = _action(text=_llm(), chat_id=_llm(value="42")) + _run(ToolExecutor(user="u"), _tools(action), {"text": "hi"}, monkeypatch) + tool.execute_action.assert_called_once_with("telegram_send_message", text="hi", chat_id="42") + + def test_llm_filled_parameter_takes_the_llm_value(self, mock_tool_manager, monkeypatch): + tool = mock_tool_manager.load_tool.return_value + action = _action(text=_llm(), chat_id=_llm(value="42")) + _run(ToolExecutor(user="u"), _tools(action), {"text": "hi", "chat_id": "7"}, monkeypatch) + tool.execute_action.assert_called_once_with("telegram_send_message", text="hi", chat_id="7") + + def test_unknown_llm_arguments_are_dropped(self, mock_tool_manager, monkeypatch): + tool = mock_tool_manager.load_tool.return_value + action = _action(text=_llm()) + _run(ToolExecutor(user="u"), _tools(action), {"text": "hi", "token": "x"}, monkeypatch) + tool.execute_action.assert_called_once_with("telegram_send_message", text="hi") + + def test_api_tool_pinned_header_is_not_overridden(self, mock_tool_manager, monkeypatch): + executor = ToolExecutor(user="u") + monkeypatch.setattr( + "docsgpt.agents.tool_executor.ToolActionParser", + lambda _cls, **kw: Mock(parse_args=Mock(return_value=("t1", "get_item", {"id": "1", "X-Tenant": "b"}))), + ) + tools_dict = { + "t1": { + "id": "00000000-0000-0000-0000-000000000001", + "name": "api_tool", + "config": { + "actions": { + "get_item": { + "name": "get_item", + "url": "https://api.example.com/items", + "method": "GET", + "active": True, + "headers": {"properties": {"X-Tenant": _pinned("a")}}, + "query_params": {"properties": {"id": _llm()}}, + "body": {"properties": {}}, + } + } + }, + } + } + call = Mock() + call.name = "get_item" + call.id = "c1" + gen = executor.execute(tools_dict, call, "MockLLM") + while True: + try: + next(gen) + except StopIteration: + break + _, kwargs = mock_tool_manager.load_tool.call_args + assert kwargs["tool_config"]["headers"] == {"X-Tenant": "a"} + assert kwargs["tool_config"]["query_params"] == {"id": "1"} + + +@pytest.mark.unit +class TestPinnedValuesInTheSchema: + def test_pinned_parameter_is_not_shown_to_the_llm(self): + executor = ToolExecutor(user="u") + action = _action(text=_llm(), chat_id=_pinned("111"), limit=_pinned(0, "integer")) + functions = executor.prepare_tools_for_llm(_tools(action)) + params = functions[0]["function"]["parameters"] + assert set(params["properties"]) == {"text"} + + def test_pinned_parameter_is_dropped_from_the_required_list(self): + executor = ToolExecutor(user="u") + action = _action(text=_llm(required=True), chat_id={**_pinned("111"), "required": True}) + params = executor.prepare_tools_for_llm(_tools(action))[0]["function"]["parameters"] + assert params["required"] == ["text"] + + +@pytest.mark.unit +class TestPinHelpers: + def test_set_pins_fixes_and_releases_parameters(self): + from docsgpt.agents.tool_pins import set_pins + + action = _action(text=_llm(), chat_id=_llm()) + pinned = set_pins(action, {"chat_id": "123"}) + assert pinned["parameters"]["properties"]["chat_id"] == { + "type": "string", "description": "", "filled_by_llm": False, "value": "123", + } + released = set_pins(pinned, {"chat_id": None}) + assert released["parameters"]["properties"]["chat_id"]["filled_by_llm"] is True + assert released["parameters"]["properties"]["chat_id"]["value"] == "" + # The input is not modified. + assert action["parameters"]["properties"]["chat_id"]["filled_by_llm"] is True + + def test_set_pins_rejects_unknown_parameters(self): + from docsgpt.agents.tool_pins import set_pins + + with pytest.raises(ValueError): + set_pins(_action(text=_llm()), {"token": "x"}) + + @pytest.mark.parametrize( + "type_, value, expected", + [ + ("integer", "5", 5), + ("integer", 0, 0), + ("number", "2.5", 2.5), + ("boolean", "false", False), + ("boolean", True, True), + ("string", 12, "12"), + ], + ) + def test_coerce_value(self, type_, value, expected): + from docsgpt.agents.tool_pins import coerce_value + + assert coerce_value({"type": type_}, value) == expected + + @pytest.mark.parametrize( + "type_, value", + [("integer", "five"), ("integer", 2.5), ("boolean", "maybe"), ("string", ""), ("string", None), + ("string", {"a": 1}), ("array", "x")], + ) + def test_coerce_value_rejects(self, type_, value): + from docsgpt.agents.tool_pins import coerce_value + + with pytest.raises(ValueError): + coerce_value({"type": type_}, value) + + def test_merge_keeps_schema_and_ignores_type_changes(self): + from docsgpt.agents.tool_pins import merge_submitted_actions + + stored = [_action(text=_llm(), chat_id=_llm())] + merged = merge_submitted_actions(stored, [{ + "name": "telegram_send_message", + "parameters": {"properties": {"chat_id": {"type": "object", "filled_by_llm": False, "value": "9"}}}, + }], may_change_pins=True) + chat_id = merged[0]["parameters"]["properties"]["chat_id"] + assert chat_id["type"] == "string" + assert chat_id["value"] == "9" and chat_id["filled_by_llm"] is False + + def test_merge_refuses_pin_changes_without_permission(self): + from docsgpt.agents.tool_pins import PinChangeRefused, merge_submitted_actions + + stored = [_action(text=_llm(), chat_id=_pinned("111"))] + with pytest.raises(PinChangeRefused): + merge_submitted_actions(stored, [{ + "name": "telegram_send_message", + "parameters": {"properties": {"chat_id": {"filled_by_llm": True}}}, + }], may_change_pins=False) + # Resending the stored values unchanged is fine. + merged = merge_submitted_actions(stored, [{ + "name": "telegram_send_message", + "active": False, + "parameters": {"properties": {"chat_id": {"filled_by_llm": False, "value": "111"}}}, + }], may_change_pins=False) + assert merged[0]["active"] is False + + def test_carry_pins(self): + from docsgpt.agents.tool_pins import carry_pins + + old = _action(q=_llm(), team=_pinned("ENG")) + fresh = _action(q={"type": "string", "filled_by_llm": True, "value": ""}, team={"type": "string"}) + carried = carry_pins(old, fresh) + assert carried["parameters"]["properties"]["team"]["value"] == "ENG" + assert carried["parameters"]["properties"]["team"]["filled_by_llm"] is False