From d56fdafe65cbbf929a949439e60513ea1d88f1b3 Mon Sep 17 00:00:00 2001 From: Juan Denis <13461850+jhd3197@users.noreply.github.com> Date: Wed, 5 Aug 2026 16:47:17 -0400 Subject: [PATCH 1/8] fix(tools): give the model the whole tool docstring, not just line one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit tool_from_function derived a tool's description from the first line of the docstring only, so everything after the summary — the prose that actually says what a tool covers, when to use it, and what it will not do — never reached the model. The description is now built from the docstring's summary plus its extended description, with Returns:/Yields: appended as a single line. Args: stays out since it is already encoded in the parameter schema, and Raises:/Examples:/ Notes: are dropped to keep descriptions dense. Descriptions are capped at MAX_TOOL_DESCRIPTION_CHARS (1024, OpenAI's limit) and truncated on a paragraph, sentence, or word boundary; pass max_description_chars=None to opt out. ToolRegistry.register() threads the cap through. _parse_docstring_params is rewritten on the new section splitter, which also fixes Args: not terminating at multi-word headers such as See Also. Two renderers assumed single-line descriptions and now indent continuation lines: ToolDefinition.to_prompt_format() and the XML tool-grammar block. Co-Authored-By: Claude Opus 5 (1M context) --- prompture/agents/tool_grammars.py | 5 +- prompture/agents/tools_schema.py | 163 +++++++++++++++++++++++++----- tests/test_tool_use.py | 100 +++++++++++++++++- 3 files changed, 241 insertions(+), 27 deletions(-) diff --git a/prompture/agents/tool_grammars.py b/prompture/agents/tool_grammars.py index dbcccb9f..899223a1 100644 --- a/prompture/agents/tool_grammars.py +++ b/prompture/agents/tool_grammars.py @@ -103,7 +103,10 @@ def _xml_render_system_prompt(tools: list[dict[str, Any]]) -> str: desc = fn.get("description", "") or "(no description)" params = fn.get("parameters", {}) or {} params_json = json.dumps(params, separators=(",", ":")) - lines.append(f"- **{name}** — {desc}") + desc_lines = desc.strip().split("\n") + lines.append(f"- **{name}** — {desc_lines[0]}") + # Indent continuation lines so a multi-line description stays inside the bullet. + lines.extend(f" {line.strip()}" if line.strip() else "" for line in desc_lines[1:]) lines.append(f" Parameters JSON Schema: `{params_json}`") lines.extend( [ diff --git a/prompture/agents/tools_schema.py b/prompture/agents/tools_schema.py index fd1e5f38..5f430301 100644 --- a/prompture/agents/tools_schema.py +++ b/prompture/agents/tools_schema.py @@ -30,6 +30,11 @@ def get_weather(city: str, units: str = "celsius") -> str: logger = logging.getLogger("prompture.tools_schema") +#: Default cap on a tool description's length. OpenAI rejects ``function`` +#: descriptions longer than 1024 characters; other providers are more lenient. +#: Pass ``max_description_chars=None`` to opt out of truncation. +MAX_TOOL_DESCRIPTION_CHARS = 1024 + # Mapping from Python types to JSON Schema types _TYPE_MAP: dict[type, str] = { str: "string", @@ -152,7 +157,11 @@ def security_metadata(self) -> dict[str, Any] | None: def to_prompt_format(self) -> str: """Plain-text description suitable for prompt-based tool calling.""" - lines = [f"Tool: {self.name}", f" Description: {self.description}", " Parameters:"] + desc_lines = self.description.strip().split("\n") + lines = [f"Tool: {self.name}", f" Description: {desc_lines[0]}"] + # Keep continuation lines under the Description label so the block stays readable. + lines.extend(f" {line.strip()}" if line.strip() else "" for line in desc_lines[1:]) + lines.append(" Parameters:") props = self.parameters.get("properties", {}) required = set(self.parameters.get("required", [])) if not props: @@ -169,13 +178,122 @@ def to_prompt_format(self) -> str: return "\n".join(lines) +#: Google-style section headers recognised when splitting a docstring. +_ARGS_HEADERS = ("args", "arguments", "parameters") +_RETURNS_HEADERS = ("returns", "return") +_YIELDS_HEADERS = ("yields", "yield") +_SECTION_HEADERS = frozenset( + _ARGS_HEADERS + + _RETURNS_HEADERS + + _YIELDS_HEADERS + + ( + "raises", + "raise", + "exceptions", + "examples", + "example", + "notes", + "note", + "warning", + "warnings", + "attributes", + "see also", + "references", + "todo", + ) +) + + +def _section_key(line: str) -> str | None: + """Return the normalised section name if *line* is a docstring section header.""" + stripped = line.strip() + if not stripped.endswith(":"): + return None + key = stripped[:-1].strip().lower() + return key if key in _SECTION_HEADERS else None + + +def _split_docstring(docstring: str | None) -> tuple[str, dict[str, list[str]]]: + """Split a Google-style docstring into its lead text and named sections. + + Returns ``(lead, sections)`` where *lead* is everything before the first + recognised section header (summary line plus any extended description) and + *sections* maps a lower-cased section name to its still-indented lines. + """ + if not docstring: + return "", {} + + lead: list[str] = [] + sections: dict[str, list[str]] = {} + current: list[str] | None = None + + for line in docstring.split("\n"): + key = _section_key(line) + if key is not None: + current = sections.setdefault(key, []) + continue + if current is None: + lead.append(line) + else: + current.append(line) + + return "\n".join(lead).strip(), sections + + +def _first_section(sections: dict[str, list[str]], *keys: str) -> list[str]: + """Return the first present section among *keys* (empty list if none).""" + for key in keys: + if key in sections: + return sections[key] + return [] + + +def _collapse(lines: list[str]) -> str: + """Flatten a docstring section into a single space-joined paragraph.""" + return " ".join(stripped for line in lines if (stripped := line.strip())) + + +def _docstring_description(docstring: str | None) -> str: + """Build the tool description the model sees from *docstring*. + + Keeps the summary line **and** the extended description that follows it — + that prose is usually the only place a tool's scope, caveats, and intended + use are written down. ``Args:`` is dropped (it is already encoded in the + parameter schema), while ``Returns:``/``Yields:`` are appended as a single + line so the model knows what it gets back. + """ + lead, sections = _split_docstring(docstring) + parts: list[str] = [lead] if lead else [] + + for label, keys in (("Returns", _RETURNS_HEADERS), ("Yields", _YIELDS_HEADERS)): + body = _collapse(_first_section(sections, *keys)) + if body: + parts.append(f"{label}: {body}") + + return "\n\n".join(parts).strip() + + +def _truncate_description(text: str, limit: int | None) -> str: + """Trim *text* to *limit* characters on a paragraph, sentence, or word boundary.""" + if limit is None or len(text) <= limit: + return text + head = text[: limit - 1] + for sep in ("\n\n", ". ", " "): + idx = head.rfind(sep) + if idx > limit // 2: + head = head[:idx] + break + return head.rstrip(" \n.,;:") + "…" + + def _parse_docstring_params(docstring: str | None) -> dict[str, str]: """Extract parameter descriptions from a Google-style docstring ``Args:`` section.""" - if not docstring: + _, sections = _split_docstring(docstring) + lines = _first_section(sections, *_ARGS_HEADERS) + if not lines: return {} - lines = docstring.split("\n") + params: dict[str, str] = {} - in_args = False current_param: str | None = None current_desc_parts: list[str] = [] args_indent: int | None = None @@ -183,23 +301,7 @@ def _parse_docstring_params(docstring: str | None) -> dict[str, str]: for line in lines: stripped = line.strip() - # Detect start of Args section - if stripped in ("Args:", "Arguments:", "Parameters:"): - in_args = True - args_indent = None - continue - - if not in_args: - continue - - # Detect end of Args section (next section header like Returns:, Raises:, etc.) - if stripped and not stripped.startswith("-") and stripped.endswith(":") and " " not in stripped: - # Save last param - if current_param is not None: - params[current_param] = " ".join(current_desc_parts).strip() - break - - # Empty line inside Args might end the section or just be spacing + # Blank lines inside Args are just spacing if not stripped: continue @@ -233,18 +335,28 @@ def _parse_docstring_params(docstring: str | None) -> dict[str, str]: def tool_from_function( - fn: Callable[..., Any], *, name: str | None = None, description: str | None = None + fn: Callable[..., Any], + *, + name: str | None = None, + description: str | None = None, + max_description_chars: int | None = MAX_TOOL_DESCRIPTION_CHARS, ) -> ToolDefinition: """Build a :class:`ToolDefinition` by inspecting *fn*'s signature and docstring. Parameters: fn: The callable to wrap. name: Override the tool name (defaults to ``fn.__name__``). - description: Override the description (defaults to the first line of the docstring). + description: Override the description. By default the docstring's + summary *and* extended description are used, with ``Returns:`` + appended; ``Args:`` is omitted because it already lives in the + parameter schema. + max_description_chars: Truncate the description to this many characters + (see :data:`MAX_TOOL_DESCRIPTION_CHARS`). ``None`` disables it. """ tool_name = name or fn.__name__ raw_doc = inspect.getdoc(fn) or "" - tool_desc = description or raw_doc.split("\n")[0] or f"Call {tool_name}" + tool_desc = description or _docstring_description(raw_doc) or f"Call {tool_name}" + tool_desc = _truncate_description(tool_desc, max_description_chars) param_docs = _parse_docstring_params(raw_doc) sig = inspect.signature(fn) @@ -320,9 +432,10 @@ def register( *, name: str | None = None, description: str | None = None, + max_description_chars: int | None = MAX_TOOL_DESCRIPTION_CHARS, ) -> ToolDefinition: """Register *fn* as a tool and return the :class:`ToolDefinition`.""" - td = tool_from_function(fn, name=name, description=description) + td = tool_from_function(fn, name=name, description=description, max_description_chars=max_description_chars) self._tools[td.name] = td return td diff --git a/tests/test_tool_use.py b/tests/test_tool_use.py index 72af21fe..35302ccb 100644 --- a/tests/test_tool_use.py +++ b/tests/test_tool_use.py @@ -7,7 +7,7 @@ import pytest from prompture.agents.conversation import Conversation -from prompture.agents.tools_schema import ToolRegistry, tool_from_function +from prompture.agents.tools_schema import MAX_TOOL_DESCRIPTION_CHARS, ToolRegistry, tool_from_function from prompture.drivers.base import Driver # --------------------------------------------------------------------------- @@ -68,6 +68,104 @@ def f(x: int) -> str: assert "input_schema" in fmt +class TestDocstringDescription: + def test_extended_description_is_kept(self): + def search(query: str) -> str: + """Search the internal wiki. + + Covers engineering runbooks and onboarding docs only. + Does not reach public web pages. + + Args: + query: Full-text search query. + + Returns: + The top matching page as markdown. + """ + return query + + td = tool_from_function(search) + assert td.description.startswith("Search the internal wiki.") + assert "engineering runbooks" in td.description + assert "Does not reach public web pages." in td.description + # Args live in the parameter schema, not the description + assert "Full-text search query" not in td.description + assert td.parameters["properties"]["query"]["description"] == "Full-text search query." + # Returns is appended so the model knows what comes back + assert "Returns: The top matching page as markdown." in td.description + + def test_summary_only_docstring_unchanged(self): + def ping() -> str: + """Check liveness.""" + return "ok" + + assert tool_from_function(ping).description == "Check liveness." + + def test_no_docstring_falls_back(self): + def bare(x: int) -> int: + return x + + assert tool_from_function(bare).description == "Call bare" + + def test_sections_after_args_are_dropped(self): + def f(x: int) -> int: + """Do a thing. + + Args: + x: A number. + + Raises: + ValueError: If x is negative. + + Examples: + >>> f(1) + 1 + """ + return x + + td = tool_from_function(f) + assert td.description == "Do a thing." + assert td.parameters["properties"]["x"]["description"] == "A number." + + def test_description_truncated_to_limit(self): + def f() -> None: + """Summary line. + + {body} + """ + + f.__doc__ = "Summary line.\n\n" + ("word " * 500) + td = tool_from_function(f, max_description_chars=200) + assert len(td.description) <= 200 + assert td.description.startswith("Summary line.") + assert td.description.endswith("…") + + def test_truncation_can_be_disabled(self): + def f() -> None: + pass + + f.__doc__ = "Summary line.\n\n" + ("word " * 500) + td = tool_from_function(f, max_description_chars=None) + assert len(td.description) > MAX_TOOL_DESCRIPTION_CHARS + + def test_prompt_format_indents_multiline_description(self): + def f(x: int) -> int: + """Do a thing. + + With a second paragraph. + + Args: + x: A number. + """ + return x + + text = tool_from_function(f).to_prompt_format() + assert " Description: Do a thing." in text + assert " With a second paragraph." in text + # Parameters block still renders after the description + assert text.index(" Parameters:") > text.index("With a second paragraph.") + + # --------------------------------------------------------------------------- # ToolRegistry tests # --------------------------------------------------------------------------- From c24c8c934b795ea2ef81c57df001ac09f465afbc Mon Sep 17 00:00:00 2001 From: Juan Denis <13461850+jhd3197@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:16:36 -0400 Subject: [PATCH 2/8] fix(tukuy): resolve PEP 563 annotations when building skill schemas MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit tukuy derives a skill's JSON Schema from the raw `__annotations__`, so a skill defined in a module using `from __future__ import annotations` hands it the *string* `"int"` rather than `int` — and every such parameter silently degraded to `{"type": "string"}`. That style is used throughout this codebase, so it was the common case rather than an edge case. Previously this only misinformed the model. Now that tool arguments are validated against the advertised schema it rejects correct calls outright (`double(x=7)` -> "7 is not of type 'string'"). Re-derive the parameter fragments via `typing.get_type_hints`, which resolves the strings, and merge them over tukuy's output. Only properties tukuy already emitted are touched, so parameters it deliberately hides (`context`, config params) stay hidden, and non-shape keys it attached (description, default) are preserved. Also reduce `integrations/tukuy_bridge.py` to a re-export shim. It was a byte-identical copy of the canonical module, which meant two *distinct* `current_tool_call_id` ContextVars — a tool executed via one module's var was invisible to code reading the other's, and this fix would have applied to only one copy. Co-Authored-By: Claude Opus 5 (1M context) --- prompture/extraction/tukuy_bridge.py | 80 +++- prompture/integrations/tukuy_bridge.py | 506 ++----------------------- tests/test_tukuy_bridge.py | 24 ++ 3 files changed, 137 insertions(+), 473 deletions(-) diff --git a/prompture/extraction/tukuy_bridge.py b/prompture/extraction/tukuy_bridge.py index 26101428..6d95bae0 100644 --- a/prompture/extraction/tukuy_bridge.py +++ b/prompture/extraction/tukuy_bridge.py @@ -13,9 +13,9 @@ import inspect import logging from collections.abc import Callable -from typing import Any +from typing import Any, get_type_hints -from ..agents.tools_schema import ToolDefinition, ToolRegistry +from ..agents.tools_schema import ToolDefinition, ToolRegistry, _python_type_to_json_schema logger = logging.getLogger("prompture.tukuy_bridge") @@ -30,6 +30,80 @@ # Skill → ToolDefinition # ------------------------------------------------------------------ +# Keys that describe a parameter's *shape*. When we recover a better +# annotation than tukuy managed, these are the ones we replace; anything +# else tukuy attached (description, default, examples) is preserved. +_SCHEMA_SHAPE_KEYS = frozenset( + { + "type", + "items", + "properties", + "required", + "additionalProperties", + "enum", + "anyOf", + "$ref", + "$defs", + "format", + } +) + + +def _resolved_param_schemas(skill_obj: Any) -> dict[str, dict[str, Any]]: + """JSON Schema fragments for *skill_obj*'s params, from resolved hints. + + tukuy derives its schema from the raw ``__annotations__``, so a skill + defined in a module using ``from __future__ import annotations`` (PEP + 563) hands it the *string* ``"int"`` rather than :class:`int` — and every + such parameter silently degrades to ``{"type": "string"}``. That used to + merely misinform the model; now that arguments are validated against the + schema it rejects correct calls outright. + + :func:`typing.get_type_hints` resolves those strings, so re-derive the + fragments here. Unannotated parameters are absent from the hints and so + are left to tukuy. Returns ``{}`` when the skill exposes no + introspectable function. + """ + fn = getattr(skill_obj, "fn", None) + if fn is None or not callable(fn): + return {} + try: + hints = get_type_hints(fn) + except Exception: # unresolvable forward refs, exotic annotations, … + logger.debug("Could not resolve type hints for tukuy skill %r", skill_obj, exc_info=True) + return {} + return {name: _python_type_to_json_schema(hint) for name, hint in hints.items() if name != "return"} + + +def _repair_parameter_types(skill_obj: Any, parameters: dict[str, Any]) -> dict[str, Any]: + """Return *parameters* with PEP 563-degraded property types corrected. + + Only properties tukuy already emitted are touched, so parameters it + deliberately hides (``context``, config params) stay hidden. + """ + properties = parameters.get("properties") + if not isinstance(properties, dict): + return parameters + resolved = _resolved_param_schemas(skill_obj) + if not resolved: + return parameters + + repaired: dict[str, Any] = {} + for name, prop in properties.items(): + better = resolved.get(name) + if better is None or not isinstance(prop, dict): + repaired[name] = prop + continue + merged = dict(better) + for key, value in prop.items(): + if key not in _SCHEMA_SHAPE_KEYS: + merged[key] = value + repaired[name] = merged + + out = dict(parameters) + out["properties"] = repaired + return out + def skill_to_tool_definition(skill_or_fn: Any, *, config: dict[str, Any] | None = None) -> ToolDefinition: """Convert a tukuy :class:`Skill` or ``@skill``-decorated function to a :class:`ToolDefinition`. @@ -56,7 +130,7 @@ def skill_to_tool_definition(skill_or_fn: Any, *, config: dict[str, Any] | None skill_obj = _normalize(skill_or_fn) desc = skill_obj.descriptor - parameters = _wrap_as_parameters(skill_obj) + parameters = _repair_parameter_types(skill_obj, _wrap_as_parameters(skill_obj)) # Build wrapper that calls invoke() and unwraps SkillResult def _wrapper(**kwargs: Any) -> Any: diff --git a/prompture/integrations/tukuy_bridge.py b/prompture/integrations/tukuy_bridge.py index 26101428..0269371d 100644 --- a/prompture/integrations/tukuy_bridge.py +++ b/prompture/integrations/tukuy_bridge.py @@ -1,476 +1,42 @@ -"""Bridge between tukuy's skill/chain/safety system and Prompture's agent/tool architecture. +"""Deprecated alias for :mod:`prompture.extraction.tukuy_bridge`. -Converts tukuy ``@skill``-decorated functions into Prompture :class:`ToolDefinition` -objects, wraps tukuy :class:`Chain` as pipeline steps, and applies -:class:`SafetyPolicy` to gate tool execution. +.. deprecated:: + This module moved to ``prompture.extraction.tukuy_bridge``. Import from + there instead. -All tukuy imports are lazy to avoid import-time errors if tukuy is not installed. +Everything is re-exported from the canonical module rather than duplicated. +Two copies of the source meant two *distinct* ``current_tool_call_id`` +ContextVars, so a tool executed via one module's var was invisible to code +reading the other's — and any fix applied to one copy silently skipped the +other. """ from __future__ import annotations -import contextvars -import inspect -import logging -from collections.abc import Callable -from typing import Any - -from ..agents.tools_schema import ToolDefinition, ToolRegistry - -logger = logging.getLogger("prompture.tukuy_bridge") - -# ContextVar that holds the tool_call_id for the currently executing tool. -# Set by AsyncConversation.ask_with_tool_events() before each tool execution -# so that the tukuy bridge can associate streaming deltas with the correct -# tool call. -current_tool_call_id: contextvars.ContextVar[str | None] = contextvars.ContextVar("current_tool_call_id", default=None) - - -# ------------------------------------------------------------------ -# Skill → ToolDefinition -# ------------------------------------------------------------------ - - -def skill_to_tool_definition(skill_or_fn: Any, *, config: dict[str, Any] | None = None) -> ToolDefinition: - """Convert a tukuy :class:`Skill` or ``@skill``-decorated function to a :class:`ToolDefinition`. - - Uses tukuy's ``bridges._normalize()`` and ``bridges._wrap_as_parameters()`` - to extract name, description, and JSON Schema parameters. The returned - tool's ``function`` calls ``skill.invoke(**kwargs)`` and returns - ``result.value`` on success or ``f"Error: {result.error}"`` on failure. - - Args: - skill_or_fn: A tukuy ``Skill`` instance or a ``@skill``-decorated function. - config: Optional configuration dict injected as a - :class:`SkillContext` into ``invoke()`` when no context is - already present in kwargs. - - Returns: - A :class:`ToolDefinition` wrapping the skill. - - Raises: - TypeError: If *skill_or_fn* is not a recognised tukuy skill type. - """ - from tukuy.bridges import _normalize, _wrap_as_parameters - - skill_obj = _normalize(skill_or_fn) - desc = skill_obj.descriptor - - parameters = _wrap_as_parameters(skill_obj) - - # Build wrapper that calls invoke() and unwraps SkillResult - def _wrapper(**kwargs: Any) -> Any: - invoke_kwargs = dict(kwargs) - if config is not None and "context" not in invoke_kwargs: - from tukuy import SkillContext - - ctx = SkillContext(config=config) - invoke_kwargs["context"] = ctx - result = skill_obj.invoke(**invoke_kwargs) - if result.success: - return result.value - return f"Error: {result.error}" - - # For async skill functions (or Instructions, which are always async), - # attach a dedicated async wrapper that uses ainvoke(). - from tukuy.instruction import Instruction as _TukuyInstruction - - _needs_async = ( - isinstance(skill_obj, _TukuyInstruction) - or getattr(desc, "is_async", False) - or inspect.iscoroutinefunction(getattr(skill_obj, "fn", None)) - ) - if _needs_async: - _is_instruction = isinstance(skill_obj, _TukuyInstruction) - - async def _async_wrapper(**kwargs: Any) -> Any: - invoke_kwargs = dict(kwargs) - if config is not None and "context" not in invoke_kwargs: - from tukuy import SkillContext - - ctx = SkillContext(config=config) - invoke_kwargs["context"] = ctx - - # For instructions: wire up on_delta streaming if a sender is - # available in the config and ainvoke() accepts the parameter. - if _is_instruction and config is not None: - delta_sender = config.get("on_instruction_delta") - if delta_sender is not None: - tool_id = current_tool_call_id.get() - if tool_id is not None: - - async def _on_delta(text: str) -> None: - await delta_sender(tool_id, text) - - invoke_kwargs["on_delta"] = _on_delta - - result = await skill_obj.ainvoke(**invoke_kwargs) - if result.success: - return result.value - return f"Error: {result.error}" - - _wrapper._async_fn = _async_wrapper # type: ignore[attr-defined] - - # Attach the original skill for reverse-bridge detection - _wrapper.__skill__ = skill_obj # type: ignore[attr-defined] - - return ToolDefinition( - name=desc.name, - description=desc.description or f"Call {desc.name}", - parameters=parameters, - function=_wrapper, - ) - - -def skills_to_registry(skills: list[Any]) -> ToolRegistry: - """Batch-convert tukuy skills to a :class:`ToolRegistry`. - - Args: - skills: List of tukuy ``Skill`` instances or ``@skill``-decorated functions. - - Returns: - A populated :class:`ToolRegistry`. - """ - registry = ToolRegistry() - for s in skills: - td = skill_to_tool_definition(s) - registry.add(td) - return registry - - -# ------------------------------------------------------------------ -# ToolDefinition → Skill (reverse bridge) -# ------------------------------------------------------------------ - - -def tool_definition_to_skill(td: ToolDefinition) -> Any: - """Convert a :class:`ToolDefinition` back to a tukuy :class:`Skill`. - - Builds a :class:`SkillDescriptor` from the tool's name, description, - and parameters, then wraps the tool's function as a ``Skill``. - - Args: - td: The Prompture tool definition to convert. - - Returns: - A tukuy ``Skill`` instance. - """ - from tukuy import Skill, SkillDescriptor - - descriptor = SkillDescriptor( - name=td.name, - description=td.description, - input_schema=td.parameters, - ) - return Skill(descriptor=descriptor, fn=td.function) - - -# ------------------------------------------------------------------ -# Registry ↔ skill dict -# ------------------------------------------------------------------ - - -def registry_to_skill_dict(registry: ToolRegistry) -> dict[str, Any]: - """Convert a :class:`ToolRegistry` to a dict compatible with tukuy's ``dispatch_openai()``/``dispatch_anthropic()``. - - For tools that originated from tukuy skills (detected via ``__skill__`` - attribute on the wrapper function), the original decorated function is used. - For native Prompture tools, a reverse bridge via :func:`tool_definition_to_skill` - is applied. - - Args: - registry: The Prompture tool registry. - - Returns: - A ``{name: skill_or_fn}`` dict usable with tukuy dispatch functions. - """ - result: dict[str, Any] = {} - for td in registry.definitions: - # Check if the tool's function came from a tukuy skill - skill_obj = getattr(td.function, "__skill__", None) - if skill_obj is not None: - result[td.name] = skill_obj - else: - result[td.name] = tool_definition_to_skill(td) - return result - - -# ------------------------------------------------------------------ -# TukuyChainStep — pipeline adapter -# ------------------------------------------------------------------ - - -class TukuyChainStep: - """Adapter that makes a tukuy :class:`Chain` usable as a :class:`SkillPipeline` step. - - Args: - chain: A tukuy ``Chain`` instance. - name: Display name for this step (default ``"tukuy_chain"``). - """ - - def __init__(self, chain: Any, *, name: str = "tukuy_chain") -> None: - self.chain = chain - self.name = name - - def run(self, input_text: str) -> str: - """Execute the chain synchronously. - - Args: - input_text: The input string to transform. - - Returns: - The chain's output as a string. - """ - result = self.chain.run(input_text) - return str(result) - - async def arun(self, input_text: str) -> str: - """Execute the chain asynchronously. - - Args: - input_text: The input string to transform. - - Returns: - The chain's output as a string. - """ - result = await self.chain.arun(input_text) - return str(result) - - -# ------------------------------------------------------------------ -# Safety policy gating -# ------------------------------------------------------------------ - - -def apply_safety_policy(registry: ToolRegistry, policy: Any) -> ToolRegistry: - """Return a new :class:`ToolRegistry` where tukuy-backed tools are gated by *policy*. - - For each tool whose function has a ``__skill__`` attribute (i.e. it was - created from a tukuy skill), the wrapper is replaced with one that calls - ``policy.validate()`` before execution. If validation fails, the tool - returns an error string. - - Non-tukuy tools pass through unchanged. - - Args: - registry: The source tool registry. - policy: A tukuy ``SafetyPolicy`` instance. - - Returns: - A new :class:`ToolRegistry` with safety-gated tools. - """ - new_registry = ToolRegistry() - - for td in registry.definitions: - skill_obj = getattr(td.function, "__skill__", None) - if skill_obj is not None: - # Gate this tool with the safety policy - original_fn = td.function - - def _make_gated(fn: Callable[..., Any], skill: Any) -> Callable[..., Any]: - def _gated_wrapper(**kwargs: Any) -> Any: - violations = policy.validate(skill.descriptor) - if violations: - msgs = "; ".join(v.message for v in violations) - return f"Error: Safety policy violation: {msgs}" - return fn(**kwargs) - - # Preserve the __skill__ attribute - _gated_wrapper.__skill__ = skill # type: ignore[attr-defined] - return _gated_wrapper - - gated = _make_gated(original_fn, skill_obj) - new_td = ToolDefinition( - name=td.name, - description=td.description, - parameters=td.parameters, - function=gated, - ) - new_registry.add(new_td) - else: - # Non-tukuy tool: pass through - new_registry.add(td) - - return new_registry - - -# ------------------------------------------------------------------ -# Security context gating -# ------------------------------------------------------------------ - - -def apply_security_context(registry: ToolRegistry, security_context: Any) -> ToolRegistry: - """Return a new :class:`ToolRegistry` where tukuy-backed tools run inside *security_context*. - - For each tool whose function has a ``__skill__`` attribute (i.e. it was - created from a tukuy skill), the wrapper is replaced with one that calls - ``set_security_context()`` before execution and ``reset_security_context()`` - after (even on error). - - Non-tukuy tools pass through unchanged. - - Args: - registry: The source tool registry. - security_context: A tukuy ``SecurityContext`` instance. - - Returns: - A new :class:`ToolRegistry` with security-scoped tools. - """ - from tukuy.safety import reset_security_context, set_security_context - - new_registry = ToolRegistry() - - for td in registry.definitions: - skill_obj = getattr(td.function, "__skill__", None) - if skill_obj is not None: - original_fn = td.function - - def _make_scoped(fn: Callable[..., Any], skill: Any) -> Callable[..., Any]: - def _scoped_wrapper(**kwargs: Any) -> Any: - token = set_security_context(security_context) - try: - return fn(**kwargs) - finally: - reset_security_context(token) - - _scoped_wrapper.__skill__ = skill # type: ignore[attr-defined] - return _scoped_wrapper - - scoped = _make_scoped(original_fn, skill_obj) - new_td = ToolDefinition( - name=td.name, - description=td.description, - parameters=td.parameters, - function=scoped, - ) - new_registry.add(new_td) - else: - new_registry.add(td) - - return new_registry - - -# ------------------------------------------------------------------ -# Availability filtering -# ------------------------------------------------------------------ - - -def filter_available_skills( - registry: ToolRegistry, - *, - policy: Any | None = None, -) -> ToolRegistry: - """Return a new :class:`ToolRegistry` containing only available tukuy skills. - - Uses tukuy's :func:`get_available_skills` with a virtual plugin wrapper - to filter a :class:`ToolRegistry`. Non-tukuy tools always pass through. - - Args: - registry: The source tool registry. - policy: Optional tukuy ``SafetyPolicy``. When ``None``, all - skills are considered available. - - Returns: - A new :class:`ToolRegistry` with only available tools. - """ - from tukuy import get_available_skills - from tukuy.plugins.base import TransformerPlugin - - new_registry = ToolRegistry() - - # Collect tukuy skills for availability check - tukuy_tools: list[tuple[ToolDefinition, Any]] = [] - for td in registry.definitions: - skill_obj = getattr(td.function, "__skill__", None) - if skill_obj is not None: - tukuy_tools.append((td, skill_obj)) - else: - # Non-tukuy tools pass through unchanged - new_registry.add(td) - - if not tukuy_tools: - return new_registry - - # Build a virtual plugin wrapping the registry's tukuy skills - skill_dict = {s.descriptor.name: s for _, s in tukuy_tools} - - class _VirtualPlugin(TransformerPlugin): # type: ignore[misc] - @property - def skills(self) -> dict[str, Any]: - return skill_dict - - @property - def transformers(self) -> dict[str, Any]: - return {} - - virtual = _VirtualPlugin("_prompture_filter") - available = get_available_skills([virtual], policy=policy) - available_names = {a.skill.descriptor.name for a in available if a.available} - - for td, skill_obj in tukuy_tools: - if skill_obj.descriptor.name in available_names: - new_registry.add(td) - - return new_registry - - -def discover_and_register_plugins( - plugins: list[Any], - *, - config: dict[str, Any] | None = None, -) -> ToolRegistry: - """Discover available plugins and register their skills into a new :class:`ToolRegistry`. - - Takes a list of tukuy ``TransformerPlugin`` instances, runs - :func:`discover_plugins`, and registers all skills from available - plugins. - - Args: - plugins: List of tukuy ``TransformerPlugin`` instances. - config: Optional configuration dict forwarded to - :func:`skill_to_tool_definition`. - - Returns: - A :class:`ToolRegistry` populated with skills from available plugins. - """ - from tukuy import discover_plugins - - registry = ToolRegistry() - - if not plugins: - return registry - - results = discover_plugins(plugins) - for result in results: - if result.available: - for skill_obj in result.plugin.skills.values(): - td = skill_to_tool_definition(skill_obj, config=config) - registry.add(td) - - return registry - - -# ------------------------------------------------------------------ -# Transform chain convenience -# ------------------------------------------------------------------ - - -def make_transform_chain(*transforms: str) -> Callable[[str], str]: - """Return a callable that applies tukuy transforms to a string. - - Convenience wrapper around tukuy's :class:`Chain`. - - Args: - *transforms: Transform names (e.g. ``"strip"``, ``"lowercase"``). - - Returns: - A callable ``(str) -> str`` that applies the transforms in order. - """ - from tukuy import Chain - - chain = Chain(list(transforms)) - - def _apply(value: str) -> str: - return str(chain.run(value)) - - return _apply +from ..extraction.tukuy_bridge import ( + TukuyChainStep, + apply_safety_policy, + apply_security_context, + current_tool_call_id, + discover_and_register_plugins, + filter_available_skills, + make_transform_chain, + registry_to_skill_dict, + skill_to_tool_definition, + skills_to_registry, + tool_definition_to_skill, +) + +__all__ = [ + "TukuyChainStep", + "apply_safety_policy", + "apply_security_context", + "current_tool_call_id", + "discover_and_register_plugins", + "filter_available_skills", + "make_transform_chain", + "registry_to_skill_dict", + "skill_to_tool_definition", + "skills_to_registry", + "tool_definition_to_skill", +] diff --git a/tests/test_tukuy_bridge.py b/tests/test_tukuy_bridge.py index bc2736e2..cf096882 100644 --- a/tests/test_tukuy_bridge.py +++ b/tests/test_tukuy_bridge.py @@ -67,6 +67,30 @@ def test_from_skill_instance(self): td = skill_to_tool_definition(skill_obj) assert td.name == "double" + def test_param_types_survive_pep563_annotations(self): + """``x: int`` stays an integer even under ``from __future__ import annotations``. + + tukuy builds its schema from raw ``__annotations__``, so with PEP 563 + in effect it sees the *string* ``"int"`` and degrades every parameter + to ``{"type": "string"}``. Since arguments are now validated against + the schema, that made correct calls fail outright — the bridge + re-derives the types from resolved hints. This module uses + ``from __future__ import annotations``, so it is the real case. + """ + td = skill_to_tool_definition(double) + assert td.parameters["properties"]["x"]["type"] == "integer" + + def test_pep563_skill_executes_with_native_types(self): + """The end-to-end path an agent takes: schema, validation, execution.""" + reg = skills_to_registry([double]) + assert reg.execute("double", {"x": 21}) == 42 + + def test_str_param_still_typed_as_string(self): + """Repairing types must not mangle parameters tukuy already got right.""" + td = skill_to_tool_definition(greet) + assert td.parameters["properties"]["name"]["type"] == "string" + assert td.function(name="Ada") == "Hello, Ada!" + def test_wrapper_execution_success(self): td = skill_to_tool_definition(double) result = td.function(x=5) From 1457e36c28128903105c11f8e62091eb853fe45e Mon Sep 17 00:00:00 2001 From: Juan Denis <13461850+jhd3197@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:17:03 -0400 Subject: [PATCH 3/8] feat(drivers): forward tool_choice, normalize stop reasons, flag bad args MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three gaps in the driver layer surfaced by a tool-calling audit: `tool_choice` was silently dropped by the OpenAI, Claude and Google drivers (only minor raw-HTTP drivers forwarded it), and `parallel_tool_calls` was exposed nowhere. Both now travel through a shared `_translate_tool_choice` that maps the normalized form onto each provider's wire format, warning and ignoring rather than sending something the API will reject. `tool_choice="required"` on a final round is the standard cure for an agent that never stops, so it needs to actually reach the API. `stop_reason` was passed through raw, so callers had to know each provider's vocabulary (`tool_calls` vs `tool_use` vs numeric Gemini enums) and truncation was invisible to the agent loop. Responses now carry a normalized reason (`end_turn`, `tool_use`, `max_tokens`, `content_filter`, `error`) with the provider's original value preserved in `meta["raw_stop_reason"]` / `usage["raw_stop_reason"]`. Unknown strings pass through unchanged so no information is lost. Malformed or truncated tool arguments were logged and then executed as `{}` — a `write_file` call truncated at `max_tokens` ran with empty arguments instead of the model being told to retry. Both the buffered and streaming paths now attach `arguments_error` / `truncated` to the tool call so the conversation layer can feed the problem back to the model. Also accept single-quoted attributes in the prompted-tool XML grammar and give `ToolGrammar` an `open_prefix` so the streaming parser can hold back a partial delimiter split across chunks. Co-Authored-By: Claude Opus 5 (1M context) --- prompture/agents/live_events.py | 17 +- prompture/agents/tool_grammars.py | 16 +- prompture/drivers/_openai_compat_stream.py | 102 ++++++-- prompture/drivers/_prompted_tool_stream.py | 38 +-- prompture/drivers/async_base.py | 11 +- prompture/drivers/async_claude_driver.py | 49 +++- prompture/drivers/async_google_driver.py | 48 ++-- prompture/drivers/async_openai_driver.py | 8 +- prompture/drivers/base.py | 278 ++++++++++++++++++++- prompture/drivers/claude_driver.py | 50 +++- prompture/drivers/google_driver.py | 60 +++-- prompture/drivers/groq_driver.py | 16 +- prompture/drivers/openai_driver.py | 16 +- tests/test_live_events.py | 5 +- 14 files changed, 582 insertions(+), 132 deletions(-) diff --git a/prompture/agents/live_events.py b/prompture/agents/live_events.py index bca2fa92..b896d76b 100644 --- a/prompture/agents/live_events.py +++ b/prompture/agents/live_events.py @@ -76,11 +76,20 @@ class ToolInputDelta: @dataclass(frozen=True) class ToolUseStop: - """The tool's input is complete. ``input`` is the parsed dict.""" + """The tool's input is complete. ``input`` is the parsed dict. + + ``truncated`` is set when the provider stopped early (``length`` / + ``max_tokens`` finish reason) and the streamed arguments failed to + parse — the conversation layer treats these like malformed-argument + calls (no execution, error fed back to the model). ``raw_stop_reason`` + preserves the provider's original finish reason for diagnostics. + """ id: str name: str input: dict[str, Any] + truncated: bool = False + raw_stop_reason: str | None = None event_type: Literal["tool_use_stop"] = field(default="tool_use_stop", init=False) @@ -97,8 +106,10 @@ class ToolResult: @dataclass(frozen=True) class MessageStop: - """End of one assistant turn. ``stop_reason`` matches provider semantics - (``end_turn``, ``tool_use``, ``max_tokens``, ``stop``, ``length`` …).""" + """End of one assistant turn. ``stop_reason`` uses the shared vocabulary + (``end_turn``, ``tool_use``, ``max_tokens``, ``content_filter``, + ``error``); drivers preserve the provider's raw value in + ``usage["raw_stop_reason"]``.""" stop_reason: str usage: dict[str, Any] = field(default_factory=dict) diff --git a/prompture/agents/tool_grammars.py b/prompture/agents/tool_grammars.py index 899223a1..0046c063 100644 --- a/prompture/agents/tool_grammars.py +++ b/prompture/agents/tool_grammars.py @@ -66,20 +66,29 @@ class ToolGrammar: """Given the OpenAI-style ``tools`` list, return the instruction block appended to the system prompt so the model emits this grammar.""" + open_prefix: str = " or +# Attribute values may be double- or single-quoted. _XML_OPEN_RE = re.compile( - r"[^\"]+)\"" - r"(?:\s+id\s*=\s*\"(?P[^\"]+)\")?\s*>" + r"[^\"]+)\"|'(?P[^']+)')" + r"(?:\s+id\s*=\s*(?:\"(?P[^\"]+)\"|'(?P[^']+)'))?\s*>" ) def _xml_parse_open(match: re.Match[str]) -> tuple[str, str | None]: - return match.group("name"), match.group("id") + name = match.group("name_dq") or match.group("name_sq") + tag_id = match.group("id_dq") or match.group("id_sq") + return name, tag_id def _xml_render_system_prompt(tools: list[dict[str, Any]]) -> str: @@ -138,6 +147,7 @@ def _xml_render_system_prompt(tools: list[dict[str, Any]]) -> str: close_marker="", parse_open_tag=_xml_parse_open, render_system_prompt=_xml_render_system_prompt, + open_prefix="{json}``. diff --git a/prompture/drivers/_openai_compat_stream.py b/prompture/drivers/_openai_compat_stream.py index 6856c75a..e648ac51 100644 --- a/prompture/drivers/_openai_compat_stream.py +++ b/prompture/drivers/_openai_compat_stream.py @@ -83,7 +83,7 @@ def _process_chunk( ``state["cached_prompt_tokens"]``, ``state["finish_reason"]``, ``state["tool_calls"]`` (dict indexed by tool-call slot). """ - from ..agents.live_events import TextDelta, ToolInputDelta, ToolUseStart + from ..agents.live_events import TextDelta, ThinkingDelta, ToolInputDelta, ToolUseStart if getattr(chunk, "usage", None): state["prompt_tokens"] = chunk.usage.prompt_tokens or 0 @@ -101,6 +101,11 @@ def _process_chunk( state["finish_reason"] = fr return + # Reasoning/thinking deltas (Grok, DeepSeek, Moonshot reasoning models) + reasoning = getattr(delta, "reasoning_content", None) + if reasoning: + yield ThinkingDelta(text=reasoning) + content = getattr(delta, "content", None) or "" if content: yield TextDelta(text=content) @@ -137,27 +142,55 @@ def _process_chunk( def _finalize_tool_use_stops(state: dict[str, Any]) -> Iterator[Any]: - """After the stream ends, emit ``ToolUseStop`` per accumulated tool.""" + """After the stream ends, emit ``ToolUseStop`` per accumulated tool. + + Guarantees a ``ToolUseStart`` always precedes its ``ToolUseStop`` + (synthesizing one here when the provider's chunks never completed the + start conditions) and generates a ``call_`` fallback id when the + provider omitted tool-call ids. When the finish reason was + ``length``/``max_tokens`` and the assembled arguments fail to parse, + the stop is flagged ``truncated=True`` (contract C3). + """ + import uuid + from ..agents.live_events import ToolUseStart, ToolUseStop + finish_reason = state.get("finish_reason") for idx in sorted(state["tool_calls"].keys()): bucket = state["tool_calls"][idx] - if not bucket["start_emitted"] and bucket["id"] and bucket["name"]: + if not bucket["id"]: + bucket["id"] = f"call_{uuid.uuid4().hex[:24]}" + if not bucket["start_emitted"]: yield ToolUseStart(id=bucket["id"], name=bucket["name"]) bucket["start_emitted"] = True args_str = "".join(bucket["args_fragments"]) + parse_failed = False try: parsed = json.loads(args_str) if args_str else {} if not isinstance(parsed, dict): + parse_failed = True + logger.warning( + "Streamed tool input for %s parsed to non-object JSON: %r", + bucket["name"], + args_str[:200], + ) parsed = {} except json.JSONDecodeError: + parse_failed = True logger.warning( "Failed to parse streamed tool input for %s: %r", bucket["name"], args_str[:200], ) parsed = {} - yield ToolUseStop(id=bucket["id"], name=bucket["name"], input=parsed) + truncated = parse_failed and finish_reason in ("length", "max_tokens") + yield ToolUseStop( + id=bucket["id"], + name=bucket["name"], + input=parsed, + truncated=truncated, + raw_stop_reason=finish_reason if parse_failed else None, + ) def _build_meta(state: dict[str, Any], model: str, cost: float) -> dict[str, Any]: @@ -171,6 +204,22 @@ def _build_meta(state: dict[str, Any], model: str, cost: float) -> dict[str, Any } +def _build_message_stop(state: dict[str, Any], model: str, cost: float) -> Any: + """Build the terminal ``MessageStop`` with a normalized ``stop_reason`` + (shared vocabulary, contract M3); the provider's raw finish reason is + preserved in ``usage["raw_stop_reason"]``.""" + from ..agents.live_events import MessageStop + from .base import _normalize_stop_reason + + meta = _build_meta(state, model, cost) + raw_finish_reason = state.get("finish_reason") + meta["raw_stop_reason"] = raw_finish_reason + return MessageStop( + stop_reason=_normalize_stop_reason(raw_finish_reason, tool_calls_present=bool(state["tool_calls"])), + usage=meta, + ) + + def _fresh_state() -> dict[str, Any]: return { "prompt_tokens": 0, @@ -201,8 +250,6 @@ def iter_openai_compat_live_events( ``stream_options={"include_usage": True}`` to populate usage in the final chunk; without it ``cost_fn`` will receive zeros. """ - from ..agents.live_events import MessageStop - state = _fresh_state() for chunk in stream: yield from _process_chunk(chunk, state) @@ -212,7 +259,7 @@ def iter_openai_compat_live_events( state["completion_tokens"], state["cached_prompt_tokens"], ) - yield MessageStop(stop_reason=state["finish_reason"], usage=_build_meta(state, model, cost)) + yield _build_message_stop(state, model, cost) async def aiter_openai_compat_live_events( @@ -227,8 +274,6 @@ async def aiter_openai_compat_live_events( """ import inspect - from ..agents.live_events import MessageStop - state = _fresh_state() async for chunk in stream: for ev in _process_chunk(chunk, state): @@ -241,7 +286,7 @@ async def aiter_openai_compat_live_events( state["cached_prompt_tokens"], ) cost = await raw_cost if inspect.isawaitable(raw_cost) else raw_cost - yield MessageStop(stop_reason=state["finish_reason"], usage=_build_meta(state, model, cost)) + yield _build_message_stop(state, model, cost) # ---------------------------------------------------------------------- @@ -268,6 +313,7 @@ def stream_openai_compat_tool_call( The driver must satisfy the contract described in this module's docstring. *provider* is the pricing-table key (e.g. ``"groq"``). """ + from .base import _apply_openai_tool_options from .openai_driver import _build_openai_base_kwargs model = options.get("model", driver.model) @@ -276,6 +322,16 @@ def stream_openai_compat_tool_call( supports_temperature = model_config["supports_temperature"] opts = {"temperature": default_temperature, "max_tokens": default_max_tokens, **options} + + # Only first-party OpenAI gets prompt_cache_key — third-party + # OpenAI-compatible endpoints reject unknown fields (see + # _build_openai_base_kwargs). Mirrors the buffered path. + prompt_cache_key = None + if provider == "openai": + from .openai_driver import _openai_prompt_cache_key + + prompt_cache_key = _openai_prompt_cache_key(messages, opts, tools) + kwargs = _build_openai_base_kwargs( model, messages, @@ -288,7 +344,9 @@ def stream_openai_compat_tool_call( "stream": True, "stream_options": {"include_usage": True}, }, + prompt_cache_key=prompt_cache_key, ) + _apply_openai_tool_options(kwargs, options) stream = driver.client.chat.completions.create(**kwargs) @@ -309,6 +367,7 @@ async def astream_openai_compat_tool_call( default_temperature: float = 1.0, ) -> AsyncIterator[Any]: """Async sibling of :func:`stream_openai_compat_tool_call`.""" + from .base import _apply_openai_tool_options from .openai_driver import _build_openai_base_kwargs model = options.get("model", driver.model) @@ -317,6 +376,13 @@ async def astream_openai_compat_tool_call( supports_temperature = model_config["supports_temperature"] opts = {"temperature": default_temperature, "max_tokens": default_max_tokens, **options} + + prompt_cache_key = None + if provider == "openai": + from .openai_driver import _openai_prompt_cache_key + + prompt_cache_key = _openai_prompt_cache_key(messages, opts, tools) + kwargs = _build_openai_base_kwargs( model, messages, @@ -329,7 +395,9 @@ async def astream_openai_compat_tool_call( "stream": True, "stream_options": {"include_usage": True}, }, + prompt_cache_key=prompt_cache_key, ) + _apply_openai_tool_options(kwargs, options) stream = await driver.client.chat.completions.create(**kwargs) @@ -394,8 +462,10 @@ def _build_raw_http_payload( payload[tokens_param] = opts.get("max_tokens", default_max_tokens) if supports_temperature and "temperature" in opts: payload["temperature"] = opts["temperature"] - if "tool_choice" in options: - payload["tool_choice"] = options["tool_choice"] + + from .base import _apply_openai_tool_options + + _apply_openai_tool_options(payload, options) return model, payload @@ -429,8 +499,6 @@ def stream_raw_http_compat_tool_call( """ import requests - from ..agents.live_events import MessageStop - model, payload = _build_raw_http_payload( driver, messages, @@ -476,7 +544,7 @@ def stream_raw_http_compat_tool_call( state["completion_tokens"], cached_tokens=state["cached_prompt_tokens"], ) - yield MessageStop(stop_reason=state["finish_reason"], usage=_build_meta(state, model, cost)) + yield _build_message_stop(state, model, cost) async def astream_raw_http_compat_tool_call( @@ -496,8 +564,6 @@ async def astream_raw_http_compat_tool_call( ``httpx.AsyncClient.stream``.""" import httpx - from ..agents.live_events import MessageStop - model, payload = _build_raw_http_payload( driver, messages, @@ -547,7 +613,7 @@ async def astream_raw_http_compat_tool_call( state["completion_tokens"], cached_tokens=state["cached_prompt_tokens"], ) - yield MessageStop(stop_reason=state["finish_reason"], usage=_build_meta(state, model, cost)) + yield _build_message_stop(state, model, cost) __all__ = [ diff --git a/prompture/drivers/_prompted_tool_stream.py b/prompture/drivers/_prompted_tool_stream.py index 17da9b67..65b99292 100644 --- a/prompture/drivers/_prompted_tool_stream.py +++ b/prompture/drivers/_prompted_tool_stream.py @@ -60,18 +60,13 @@ logger = logging.getLogger(__name__) -# The opening-delimiter prefix the parser scans for. When in narration -# mode, the parser holds back the tail of un-emitted text that could -# still grow into this prefix. -_OPEN_PREFIX = " None: self._state = _ParserState.NARRATION self._narration_buf = "" self._current_tool: _ToolInProgress | None = None - self._open_holdback = len(_OPEN_PREFIX) - 1 + self._open_holdback = max(0, len(grammar.open_prefix) - 1) self._close_holdback = max(0, len(grammar.close_marker) - 1) def feed(self, chunk: str) -> Iterator[LiveEvent]: @@ -178,8 +173,16 @@ def _drain_narration(self) -> Iterator[LiveEvent]: match = self.grammar.open_regex.search(self._narration_buf) if match is None: # No opening tag yet. Emit everything except a tail that - # could still grow into `` bool: opening tag with more characters appended. Cheap O(holdback) check that avoids unnecessary text holdback - when the buffer's tail clearly can't be the start of - `` dict[str, Any]: """Parse the accumulated args JSON, with graceful fallback.""" diff --git a/prompture/drivers/async_base.py b/prompture/drivers/async_base.py index 60af0422..0396e64e 100644 --- a/prompture/drivers/async_base.py +++ b/prompture/drivers/async_base.py @@ -132,9 +132,10 @@ async def generate_messages_with_tools_stream( yield TextDelta(text=text) import json as _json + import uuid as _uuid for tc in tool_calls: - tc_id = tc.get("id", "") or "" + tc_id = tc.get("id") or f"call_{_uuid.uuid4().hex[:24]}" tc_name = tc.get("name", "") or "" tc_args = tc.get("arguments", {}) or {} yield ToolUseStart(id=tc_id, name=tc_name) @@ -144,7 +145,13 @@ async def generate_messages_with_tools_stream( fragment = "{}" if fragment and fragment != "{}": yield ToolInputDelta(id=tc_id, fragment=fragment) - yield ToolUseStop(id=tc_id, name=tc_name, input=tc_args) + yield ToolUseStop( + id=tc_id, + name=tc_name, + input=tc_args, + truncated=bool(tc.get("arguments_error")) and stop_reason in ("length", "max_tokens"), + raw_stop_reason=meta.get("raw_stop_reason") if tc.get("arguments_error") else None, + ) yield MessageStop(stop_reason=stop_reason, usage=meta) diff --git a/prompture/drivers/async_claude_driver.py b/prompture/drivers/async_claude_driver.py index 9d2904e8..2929a213 100644 --- a/prompture/drivers/async_claude_driver.py +++ b/prompture/drivers/async_claude_driver.py @@ -30,6 +30,7 @@ cache_write_multiplier as _cache_write_multiplier, ) from .async_base import AsyncDriver +from .base import _normalize_stop_reason, _translate_tool_choice from .claude_driver import ( ClaudeDriver, _build_anthropic_json_mode_tool_def, @@ -206,7 +207,7 @@ async def generate_messages_with_tools( 'anthropic package not installed. Install it with: pip install "prompture[anthropic]"' ) - opts = {**{"temperature": 0.0, "max_tokens": 512}, **options} + opts = {**{"temperature": 0.0, "max_tokens": 4096}, **options} model = options.get("model", self.model) self._validate_model_capabilities("claude", model, using_tool_use=True) @@ -235,6 +236,9 @@ async def generate_messages_with_tools( kwargs["temperature"] = opts["temperature"] if wrapped_system is not None: kwargs["system"] = wrapped_system + tool_choice = _translate_tool_choice(options.get("tool_choice"), "anthropic") + if tool_choice is not None: + kwargs["tool_choice"] = tool_choice resp = await client.messages.create(**kwargs) @@ -249,6 +253,7 @@ async def generate_messages_with_tools( cache_write_multiplier=_cache_write_multiplier(opts.get("cache_ttl", "5m")), ) meta = _build_anthropic_meta(resp, model, total_cost) + meta["raw_stop_reason"] = resp.stop_reason text, tool_calls_out = _extract_anthropic_text_and_tool_calls(resp.content) reasoning_content = ClaudeDriver._extract_thinking(resp.content) @@ -257,7 +262,7 @@ async def generate_messages_with_tools( "text": text, "meta": meta, "tool_calls": tool_calls_out, - "stop_reason": resp.stop_reason, + "stop_reason": _normalize_stop_reason(resp.stop_reason, tool_calls_present=bool(tool_calls_out)), } if reasoning_content is not None: result["reasoning_content"] = reasoning_content @@ -411,10 +416,17 @@ async def generate_messages_with_tools_stream( kwargs["temperature"] = opts["temperature"] if wrapped_system is not None: kwargs["system"] = wrapped_system + tool_choice = _translate_tool_choice(options.get("tool_choice"), "anthropic") + if tool_choice is not None: + kwargs["tool_choice"] = tool_choice block_kinds: dict[int, str] = {} tool_block_info: dict[int, dict[str, Any]] = {} tool_input_buffers: dict[int, list[str]] = {} + # Tool blocks whose input JSON failed to parse are withheld until + # the message-level stop_reason arrives (message_delta comes AFTER + # content_block_stop) so truncation can be flagged accurately. + pending_failed_stops: list[dict[str, Any]] = [] base_input = 0 cache_read = 0 @@ -464,22 +476,28 @@ async def generate_messages_with_tools_stream( if block_kinds.get(idx) == "tool_use": info = tool_block_info.get(idx, {}) buf = "".join(tool_input_buffers.get(idx, [])) + parse_failed = False try: parsed = json.loads(buf) if buf else {} if not isinstance(parsed, dict): + parse_failed = True parsed = {} except json.JSONDecodeError: + parse_failed = True + parsed = {} + if parse_failed: logger.warning( "Failed to parse streamed tool input for %s: %r", info.get("name", "?"), buf[:200], ) - parsed = {} - yield ToolUseStop( - id=info.get("id", ""), - name=info.get("name", ""), - input=parsed, - ) + pending_failed_stops.append({"info": info}) + else: + yield ToolUseStop( + id=info.get("id", ""), + name=info.get("name", ""), + input=parsed, + ) elif ev_type == "message_delta": usage = getattr(event, "usage", None) if usage is not None: @@ -488,6 +506,18 @@ async def generate_messages_with_tools_stream( if sr: stop_reason = sr + # Flush tool stops whose input failed to parse, now that the final + # stop_reason is known (contract C3: truncated iff max_tokens). + for pending in pending_failed_stops: + info = pending["info"] + yield ToolUseStop( + id=info.get("id", ""), + name=info.get("name", ""), + input={}, + truncated=stop_reason == "max_tokens", + raw_stop_reason=stop_reason, + ) + prompt_tokens = base_input + cache_read + cache_create total_cost = self._calculate_cost( "claude", @@ -506,5 +536,6 @@ async def generate_messages_with_tools_stream( "cache_creation_tokens": cache_create, "cost": round(total_cost, 6), "model_name": model, + "raw_stop_reason": stop_reason, } - yield MessageStop(stop_reason=stop_reason, usage=meta) + yield MessageStop(stop_reason=_normalize_stop_reason(stop_reason), usage=meta) diff --git a/prompture/drivers/async_google_driver.py b/prompture/drivers/async_google_driver.py index 50728ef0..5f4e4941 100644 --- a/prompture/drivers/async_google_driver.py +++ b/prompture/drivers/async_google_driver.py @@ -10,13 +10,16 @@ try: from google import genai + from google.genai import errors as genai_errors from google.genai import types except ImportError: genai = None # type: ignore[assignment] + genai_errors = None # type: ignore[assignment] types = None # type: ignore[assignment] from ..infra.cost_mixin import CostMixin from .async_base import AsyncDriver +from .base import DriverHTTPError, _normalize_stop_reason, _translate_tool_choice from .google_driver import GoogleDriver logger = logging.getLogger(__name__) @@ -232,11 +235,13 @@ async def _do_generate( return {"text": response.text, "meta": meta} - except Exception as e: + except (genai_errors.APIError, ValueError) as e: logger.error(f"Google API request failed: {e}") - from ..exceptions import DriverError - - raise DriverError(f"Google API request failed: {e}") from e + raise DriverHTTPError( + f"Google API request failed: {e}", + status_code=getattr(e, "code", None), + provider="google", + ) from e # ------------------------------------------------------------------ # Tool use @@ -288,6 +293,12 @@ async def generate_messages_with_tools( config_dict["tools"] = [types.Tool(function_declarations=function_declarations)] + tool_choice = _translate_tool_choice(options.get("tool_choice"), "google") + if tool_choice is not None: + config_dict["tool_config"] = types.ToolConfig( + function_calling_config=types.FunctionCallingConfig(**tool_choice) + ) + try: config = types.GenerateContentConfig(**config_dict) response = await self._client.aio.models.generate_content( @@ -305,7 +316,7 @@ async def generate_messages_with_tools( text = "" tool_calls_out: list[dict[str, Any]] = [] - stop_reason = "stop" + raw_stop_reason: Any = None for candidate in response.candidates or []: if candidate.content is None or candidate.content.parts is None: @@ -325,11 +336,10 @@ async def generate_messages_with_tools( finish_reason = getattr(candidate, "finish_reason", None) if finish_reason is not None: - reason_map = {1: "stop", 2: "max_tokens", 3: "safety", 4: "recitation", 5: "other"} - stop_reason = reason_map.get(finish_reason, "stop") + raw_stop_reason = finish_reason - if tool_calls_out: - stop_reason = "tool_use" + stop_reason = _normalize_stop_reason(raw_stop_reason, tool_calls_present=bool(tool_calls_out)) + meta["raw_stop_reason"] = str(raw_stop_reason) if raw_stop_reason is not None else None return { "text": text, @@ -338,11 +348,13 @@ async def generate_messages_with_tools( "stop_reason": stop_reason, } - except Exception as e: + except (genai_errors.APIError, ValueError) as e: logger.error(f"Google API tool call request failed: {e}") - from ..exceptions import DriverError - - raise DriverError(f"Google API tool call request failed: {e}") from e + raise DriverHTTPError( + f"Google API tool call request failed: {e}", + status_code=getattr(e, "code", None), + provider="google", + ) from e # ------------------------------------------------------------------ # Streaming @@ -384,8 +396,10 @@ async def generate_messages_stream( }, } - except Exception as e: + except (genai_errors.APIError, ValueError) as e: logger.error(f"Google API streaming request failed: {e}") - from ..exceptions import DriverError - - raise DriverError(f"Google API streaming request failed: {e}") from e + raise DriverHTTPError( + f"Google API streaming request failed: {e}", + status_code=getattr(e, "code", None), + provider="google", + ) from e diff --git a/prompture/drivers/async_openai_driver.py b/prompture/drivers/async_openai_driver.py index 8830a43c..2939aa92 100644 --- a/prompture/drivers/async_openai_driver.py +++ b/prompture/drivers/async_openai_driver.py @@ -14,6 +14,7 @@ from ..infra.cost_mixin import CostMixin from .async_base import AsyncDriver +from .base import _apply_openai_tool_options, _normalize_stop_reason from .openai_driver import ( OpenAIDriver, _build_openai_base_kwargs, @@ -161,6 +162,7 @@ async def generate_messages_with_tools( extra={"tools": tools}, prompt_cache_key=_openai_prompt_cache_key(messages, opts, tools), ) + _apply_openai_tool_options(kwargs, options) resp = await self.client.chat.completions.create(**kwargs) @@ -175,8 +177,10 @@ async def generate_messages_with_tools( choice = resp.choices[0] text = choice.message.content or "" - stop_reason = choice.finish_reason - tool_calls_out = _extract_openai_tool_calls(choice.message, stop_reason) + raw_stop_reason = choice.finish_reason + tool_calls_out = _extract_openai_tool_calls(choice.message, raw_stop_reason) + stop_reason = _normalize_stop_reason(raw_stop_reason, tool_calls_present=bool(tool_calls_out)) + meta["raw_stop_reason"] = raw_stop_reason return { "text": text, diff --git a/prompture/drivers/base.py b/prompture/drivers/base.py index ceba1cab..661948d6 100644 --- a/prompture/drivers/base.py +++ b/prompture/drivers/base.py @@ -16,31 +16,235 @@ import contextlib +from ..exceptions import DriverError as _DriverError from ..infra.callbacks import DriverCallbacks logger = logging.getLogger("prompture.driver") +# ------------------------------------------------------------------ +# Shared driver error with HTTP context +# ------------------------------------------------------------------ + + +class DriverHTTPError(_DriverError): + """A driver request failed with HTTP/provider context attached. + + Subclasses :class:`prompture.exceptions.DriverError` so existing + ``except DriverError`` callers keep working, while adding structured + fields for retry logic and observability: + + - ``status_code`` — HTTP status when the provider returned one. + - ``provider`` — pricing-table/provider key (``"grok"``, ``"ollama"`` …). + - ``retryable`` — heuristic: timeouts/5xx/429 are worth retrying, + 4xx auth/validation errors are not. + """ + + def __init__( + self, + message: str, + *, + status_code: int | None = None, + provider: str | None = None, + retryable: bool | None = None, + ) -> None: + super().__init__(message) + self.status_code = status_code + self.provider = provider + if retryable is None: + retryable = status_code is None or status_code in (408, 409, 425, 429) or status_code >= 500 + self.retryable = retryable + + +# ------------------------------------------------------------------ +# Shared stop-reason normalization (driver boundary) +# ------------------------------------------------------------------ + +#: Canonical stop-reason vocabulary shared by all drivers. +STOP_REASONS: frozenset[str] = frozenset({"end_turn", "tool_use", "max_tokens", "content_filter", "error"}) + +_STOP_REASON_MAP: dict[str, str] = { + # OpenAI-compatible finish_reason values + "stop": "end_turn", + "length": "max_tokens", + "tool_calls": "tool_use", + "function_call": "tool_use", + "content_filter": "content_filter", + # Anthropic (already canonical, listed for completeness) + "end_turn": "end_turn", + "tool_use": "tool_use", + "max_tokens": "max_tokens", + "stop_sequence": "end_turn", + "refusal": "content_filter", + # Cohere v2 finish_reason values + "complete": "end_turn", + "tool_call": "tool_use", + "error": "error", + "error_limit": "error", + # Ollama done_reason values ("stop"/"length" covered above) + "load": "end_turn", + "unload": "end_turn", + # Google Gemini FinishReason enum values (string form) + "safety": "content_filter", + "recitation": "content_filter", + "blocklist": "content_filter", + "prohibited_content": "content_filter", + "image_safety": "content_filter", + "malformed_function_call": "error", + "finish_reason_unspecified": "end_turn", + "other": "end_turn", +} + +#: Google Gemini numeric FinishReason enum values. +_GOOGLE_STOP_REASON_MAP: dict[int, str] = { + 1: "end_turn", # STOP + 2: "max_tokens", # MAX_TOKENS + 3: "content_filter", # SAFETY + 4: "content_filter", # RECITATION + 5: "end_turn", # OTHER +} + + +def _normalize_stop_reason(raw: Any, *, tool_calls_present: bool = False) -> str: + """Normalize a provider stop/finish reason to the shared vocabulary. + + Returns one of ``end_turn``, ``tool_use``, ``max_tokens``, + ``content_filter``, ``error``. Unknown strings pass through unchanged + so no information is lost. When the response contains tool calls but + the provider reported a plain end-of-turn (Ollama does this), the + reason is upgraded to ``tool_use``. + """ + import enum + + if isinstance(raw, enum.Enum): + raw = raw.value + if isinstance(raw, bool): + normalized = "end_turn" + elif isinstance(raw, int): + normalized = _GOOGLE_STOP_REASON_MAP.get(raw, "end_turn") + elif isinstance(raw, str): + normalized = _STOP_REASON_MAP.get(raw.lower(), raw) + else: + normalized = "end_turn" + if tool_calls_present and normalized == "end_turn": + return "tool_use" + return normalized + + +# ------------------------------------------------------------------ +# Shared tool_choice translation +# ------------------------------------------------------------------ + + +def _translate_tool_choice(tool_choice: Any, api: str) -> Any: + """Translate a normalized ``tool_choice`` option into wire format. + + Accepted normalized input: ``"auto"``, ``"none"``, ``"required"``, or + ``{"name": ""}`` to force a specific tool. Provider-shaped + values (already carrying ``type``/``mode`` keys) pass through + unchanged so advanced users keep an escape hatch. + + *api* is ``"openai"`` (OpenAI-compatible chat completions), + ``"anthropic"``, or ``"google"`` (returns a ``function_calling_config`` + dict the caller wraps in ``types.ToolConfig``). + + Returns ``None`` when *tool_choice* is ``None`` or unusable (a + warning is logged in the latter case). + """ + if tool_choice is None: + return None + + if api == "anthropic": + if isinstance(tool_choice, str): + mapping = {"auto": "auto", "none": "none", "required": "any"} + t = mapping.get(tool_choice) + if t is None: + logger.warning("Unsupported tool_choice %r for anthropic; ignoring", tool_choice) + return None + return {"type": t} + if isinstance(tool_choice, dict): + if "type" in tool_choice: + return tool_choice + if "name" in tool_choice: + return {"type": "tool", "name": tool_choice["name"]} + logger.warning("Unsupported tool_choice %r for anthropic; ignoring", tool_choice) + return None + + if api == "google": + if isinstance(tool_choice, str): + mapping = {"auto": "AUTO", "none": "NONE", "required": "ANY"} + m = mapping.get(tool_choice) + if m is None: + logger.warning("Unsupported tool_choice %r for google; ignoring", tool_choice) + return None + return {"mode": m} + if isinstance(tool_choice, dict): + if "mode" in tool_choice: + return tool_choice + if "name" in tool_choice: + return {"mode": "ANY", "allowed_function_names": [tool_choice["name"]]} + logger.warning("Unsupported tool_choice %r for google; ignoring", tool_choice) + return None + + # OpenAI-compatible wire format + if isinstance(tool_choice, str): + if tool_choice in ("auto", "none", "required"): + return tool_choice + logger.warning("Unsupported tool_choice %r for openai-compatible API; ignoring", tool_choice) + return None + if isinstance(tool_choice, dict): + if "type" in tool_choice: + return tool_choice + if "name" in tool_choice: + return {"type": "function", "function": {"name": tool_choice["name"]}} + logger.warning("Unsupported tool_choice %r for openai-compatible API; ignoring", tool_choice) + return None + + +def _apply_openai_tool_options(kwargs: dict[str, Any], options: dict[str, Any]) -> None: + """Pass ``tool_choice`` / ``parallel_tool_calls`` through to an + OpenAI-compatible request payload, translating ``tool_choice`` from + the normalized form (see :func:`_translate_tool_choice`).""" + if "tool_choice" in options: + translated = _translate_tool_choice(options["tool_choice"], "openai") + if translated is not None: + kwargs["tool_choice"] = translated + if "parallel_tool_calls" in options: + kwargs["parallel_tool_calls"] = options["parallel_tool_calls"] + + # ------------------------------------------------------------------ # Shared tool-argument parser for OpenAI-compatible drivers # ------------------------------------------------------------------ -def _parse_tool_arguments(raw_args: Any, tool_name: str, stop_reason: str | None = None) -> dict[str, Any]: - """Parse tool call arguments, handling both string and dict formats. +def _parse_tool_arguments_with_error( + raw_args: Any, tool_name: str, stop_reason: str | None = None +) -> tuple[dict[str, Any], str | None]: + """Parse tool call arguments, returning ``(arguments, error)``. - Some providers return ``arguments`` as a JSON string, others as an - already-parsed dict. Calling ``json.loads()`` on a dict raises - ``TypeError`` which previously caused a silent fallback to ``{}``. + Same resilience contract as :func:`_parse_tool_arguments` (never + raises, falls back to ``{}``) but also returns a human-readable + error message when parsing failed or truncation is detected, so + callers can attach ``tc["arguments_error"]`` and the conversation + layer can ask the model to retry instead of executing garbage. """ if isinstance(raw_args, dict): - return raw_args + return raw_args, None if isinstance(raw_args, str): try: parsed = json.loads(raw_args) - return parsed if isinstance(parsed, dict) else {} + if isinstance(parsed, dict): + return parsed, None + msg = f"Tool arguments for {tool_name} parsed to non-object JSON: {raw_args[:200]!r}" + logger.warning("Tool arguments for %s parsed to non-object JSON: %r", tool_name, raw_args[:200]) + return {}, msg except json.JSONDecodeError: - if stop_reason == "length": + if stop_reason in ("length", "max_tokens"): + msg = ( + f"Tool arguments for {tool_name} were truncated due to the max_tokens limit. " + "Increase max_tokens in options to allow longer tool outputs." + ) logger.warning( "Tool arguments for %s were truncated due to max_tokens limit. " "Increase max_tokens in options to allow longer tool outputs. " @@ -49,21 +253,62 @@ def _parse_tool_arguments(raw_args: Any, tool_name: str, stop_reason: str | None raw_args[:200] if raw_args else raw_args, ) else: + msg = f"Failed to parse tool arguments for {tool_name} as JSON: {raw_args[:200]!r}" logger.warning( "Failed to parse tool arguments for %s: %r", tool_name, raw_args[:200] if raw_args else raw_args, ) - return {} + return {}, msg if raw_args is None: - return {} + return {}, None + msg = f"Unexpected argument type {type(raw_args).__name__} for tool {tool_name}" logger.warning( "Unexpected argument type %s for tool %s: %r", type(raw_args).__name__, tool_name, raw_args, ) - return {} + return {}, msg + + +def _parse_tool_arguments(raw_args: Any, tool_name: str, stop_reason: str | None = None) -> dict[str, Any]: + """Parse tool call arguments, handling both string and dict formats. + + Some providers return ``arguments`` as a JSON string, others as an + already-parsed dict. Calling ``json.loads()`` on a dict raises + ``TypeError`` which previously caused a silent fallback to ``{}``. + """ + args, _error = _parse_tool_arguments_with_error(raw_args, tool_name, stop_reason) + return args + + +def _tool_call_dict( + tool_id: Any, + name: str, + raw_args: Any, + stop_reason: str | None = None, + *, + generate_id: bool = True, +) -> dict[str, Any]: + """Build a normalized tool-call dict for ``generate_messages_with_tools``. + + Generates a ``call_`` fallback id when the provider omits one + (OpenAI-compat streaming and some raw-HTTP providers do) and attaches + ``arguments_error`` when argument parsing failed or truncation is + detected (contract C1). + """ + import uuid as _uuid + + args, args_error = _parse_tool_arguments_with_error(raw_args, name, stop_reason) + tc: dict[str, Any] = { + "id": tool_id or (f"call_{_uuid.uuid4().hex[:24]}" if generate_id else ""), + "name": name, + "arguments": args, + } + if args_error: + tc["arguments_error"] = args_error + return tc # ------------------------------------------------------------------ @@ -278,9 +523,10 @@ def generate_messages_with_tools_stream( yield TextDelta(text=text) import json as _json + import uuid as _uuid for tc in tool_calls: - tc_id = tc.get("id", "") or "" + tc_id = tc.get("id") or f"call_{_uuid.uuid4().hex[:24]}" tc_name = tc.get("name", "") or "" tc_args = tc.get("arguments", {}) or {} yield ToolUseStart(id=tc_id, name=tc_name) @@ -290,7 +536,13 @@ def generate_messages_with_tools_stream( fragment = "{}" if fragment and fragment != "{}": yield ToolInputDelta(id=tc_id, fragment=fragment) - yield ToolUseStop(id=tc_id, name=tc_name, input=tc_args) + yield ToolUseStop( + id=tc_id, + name=tc_name, + input=tc_args, + truncated=bool(tc.get("arguments_error")) and stop_reason in ("length", "max_tokens"), + raw_stop_reason=meta.get("raw_stop_reason") if tc.get("arguments_error") else None, + ) yield MessageStop(stop_reason=stop_reason, usage=meta) diff --git a/prompture/drivers/claude_driver.py b/prompture/drivers/claude_driver.py index aaf3ceb9..812f7a2b 100644 --- a/prompture/drivers/claude_driver.py +++ b/prompture/drivers/claude_driver.py @@ -35,7 +35,7 @@ from ._prompt_cache import ( cache_write_multiplier as _cache_write_multiplier, ) -from .base import Driver +from .base import Driver, _normalize_stop_reason, _translate_tool_choice logger = logging.getLogger(__name__) @@ -470,7 +470,7 @@ def generate_messages_with_tools( 'anthropic package not installed. Install it with: pip install "prompture[anthropic]"' ) - opts = {**{"temperature": 0.0, "max_tokens": 512}, **options} + opts = {**{"temperature": 0.0, "max_tokens": 4096}, **options} model = options.get("model", self.model) self._validate_model_capabilities("claude", model, using_tool_use=True) @@ -499,6 +499,9 @@ def generate_messages_with_tools( kwargs["temperature"] = opts["temperature"] if wrapped_system is not None: kwargs["system"] = wrapped_system + tool_choice = _translate_tool_choice(options.get("tool_choice"), "anthropic") + if tool_choice is not None: + kwargs["tool_choice"] = tool_choice resp = client.messages.create(**kwargs) @@ -513,6 +516,7 @@ def generate_messages_with_tools( cache_write_multiplier=_cache_write_multiplier(opts.get("cache_ttl", "5m")), ) meta = _build_anthropic_meta(resp, model, total_cost) + meta["raw_stop_reason"] = resp.stop_reason text, tool_calls_out = _extract_anthropic_text_and_tool_calls(resp.content) reasoning_content = self._extract_thinking(resp.content) @@ -521,7 +525,7 @@ def generate_messages_with_tools( "text": text, "meta": meta, "tool_calls": tool_calls_out, - "stop_reason": resp.stop_reason, + "stop_reason": _normalize_stop_reason(resp.stop_reason, tool_calls_present=bool(tool_calls_out)), } if reasoning_content is not None: result["reasoning_content"] = reasoning_content @@ -684,10 +688,17 @@ def generate_messages_with_tools_stream( kwargs["temperature"] = opts["temperature"] if wrapped_system is not None: kwargs["system"] = wrapped_system + tool_choice = _translate_tool_choice(options.get("tool_choice"), "anthropic") + if tool_choice is not None: + kwargs["tool_choice"] = tool_choice block_kinds: dict[int, str] = {} tool_block_info: dict[int, dict[str, Any]] = {} tool_input_buffers: dict[int, list[str]] = {} + # Tool blocks whose input JSON failed to parse are withheld until + # the message-level stop_reason arrives (message_delta comes AFTER + # content_block_stop) so truncation can be flagged accurately. + pending_failed_stops: list[dict[str, Any]] = [] base_input = 0 cache_read = 0 @@ -737,22 +748,28 @@ def generate_messages_with_tools_stream( if block_kinds.get(idx) == "tool_use": info = tool_block_info.get(idx, {}) buf = "".join(tool_input_buffers.get(idx, [])) + parse_failed = False try: parsed = json.loads(buf) if buf else {} if not isinstance(parsed, dict): + parse_failed = True parsed = {} except json.JSONDecodeError: + parse_failed = True + parsed = {} + if parse_failed: logger.warning( "Failed to parse streamed tool input for %s: %r", info.get("name", "?"), buf[:200], ) - parsed = {} - yield ToolUseStop( - id=info.get("id", ""), - name=info.get("name", ""), - input=parsed, - ) + pending_failed_stops.append({"info": info}) + else: + yield ToolUseStop( + id=info.get("id", ""), + name=info.get("name", ""), + input=parsed, + ) elif ev_type == "message_delta": usage = getattr(event, "usage", None) if usage is not None: @@ -761,6 +778,18 @@ def generate_messages_with_tools_stream( if sr: stop_reason = sr + # Flush tool stops whose input failed to parse, now that the final + # stop_reason is known (contract C3: truncated iff max_tokens). + for pending in pending_failed_stops: + info = pending["info"] + yield ToolUseStop( + id=info.get("id", ""), + name=info.get("name", ""), + input={}, + truncated=stop_reason == "max_tokens", + raw_stop_reason=stop_reason, + ) + prompt_tokens = base_input + cache_read + cache_create total_cost = self._calculate_cost( "claude", @@ -779,5 +808,6 @@ def generate_messages_with_tools_stream( "cache_creation_tokens": cache_create, "cost": round(total_cost, 6), "model_name": model, + "raw_stop_reason": stop_reason, } - yield MessageStop(stop_reason=stop_reason, usage=meta) + yield MessageStop(stop_reason=_normalize_stop_reason(stop_reason), usage=meta) diff --git a/prompture/drivers/google_driver.py b/prompture/drivers/google_driver.py index d9e36ee7..028dc257 100644 --- a/prompture/drivers/google_driver.py +++ b/prompture/drivers/google_driver.py @@ -7,13 +7,16 @@ try: from google import genai + from google.genai import errors as genai_errors from google.genai import types except ImportError: genai = None # type: ignore[assignment] + genai_errors = None # type: ignore[assignment] types = None # type: ignore[assignment] +from ..exceptions import DriverError from ..infra.cost_mixin import CostMixin -from .base import Driver +from .base import Driver, DriverHTTPError, _normalize_stop_reason, _translate_tool_choice logger = logging.getLogger(__name__) @@ -256,7 +259,14 @@ def _build_generation_args( elif role == "tool": # Tool result → user with function_response part tc_id = msg.get("tool_call_id", "") - name = tool_call_names.get(tc_id, "unknown_tool") + name = tool_call_names.get(tc_id) + if name is None: + raise DriverError( + f"Tool result message references unknown tool_call_id {tc_id!r}: " + "no matching assistant tool call exists in the conversation " + "history, and Gemini rejects function_response parts with a " + "fabricated name." + ) result_content = content if isinstance(result_content, str): try: @@ -337,11 +347,13 @@ def _do_generate(self, messages: list[dict[str, str]], options: dict[str, Any] | return {"text": response.text, "meta": meta} - except Exception as e: + except (genai_errors.APIError, ValueError) as e: logger.error(f"Google API request failed: {e}") - from ..exceptions import DriverError - - raise DriverError(f"Google API request failed: {e}") from e + raise DriverHTTPError( + f"Google API request failed: {e}", + status_code=getattr(e, "code", None), + provider="google", + ) from e # ------------------------------------------------------------------ # Tool use @@ -394,6 +406,12 @@ def generate_messages_with_tools( config_dict["tools"] = [types.Tool(function_declarations=function_declarations)] + tool_choice = _translate_tool_choice(options.get("tool_choice"), "google") + if tool_choice is not None: + config_dict["tool_config"] = types.ToolConfig( + function_calling_config=types.FunctionCallingConfig(**tool_choice) + ) + try: config = types.GenerateContentConfig(**config_dict) response = self._client.models.generate_content( @@ -411,7 +429,7 @@ def generate_messages_with_tools( text = "" tool_calls_out: list[dict[str, Any]] = [] - stop_reason = "stop" + raw_stop_reason: Any = None for candidate in response.candidates or []: if candidate.content is None or candidate.content.parts is None: @@ -431,12 +449,10 @@ def generate_messages_with_tools( finish_reason = getattr(candidate, "finish_reason", None) if finish_reason is not None: - # Map Gemini finish reasons to standard stop reasons - reason_map = {1: "stop", 2: "max_tokens", 3: "safety", 4: "recitation", 5: "other"} - stop_reason = reason_map.get(finish_reason, "stop") + raw_stop_reason = finish_reason - if tool_calls_out: - stop_reason = "tool_use" + stop_reason = _normalize_stop_reason(raw_stop_reason, tool_calls_present=bool(tool_calls_out)) + meta["raw_stop_reason"] = str(raw_stop_reason) if raw_stop_reason is not None else None return { "text": text, @@ -445,11 +461,13 @@ def generate_messages_with_tools( "stop_reason": stop_reason, } - except Exception as e: + except (genai_errors.APIError, ValueError) as e: logger.error(f"Google API tool call request failed: {e}") - from ..exceptions import DriverError - - raise DriverError(f"Google API tool call request failed: {e}") from e + raise DriverHTTPError( + f"Google API tool call request failed: {e}", + status_code=getattr(e, "code", None), + provider="google", + ) from e # ------------------------------------------------------------------ # Streaming @@ -491,8 +509,10 @@ def generate_messages_stream( }, } - except Exception as e: + except (genai_errors.APIError, ValueError) as e: logger.error(f"Google API streaming request failed: {e}") - from ..exceptions import DriverError - - raise DriverError(f"Google API streaming request failed: {e}") from e + raise DriverHTTPError( + f"Google API streaming request failed: {e}", + status_code=getattr(e, "code", None), + provider="google", + ) from e diff --git a/prompture/drivers/groq_driver.py b/prompture/drivers/groq_driver.py index 8380f978..05578c3f 100644 --- a/prompture/drivers/groq_driver.py +++ b/prompture/drivers/groq_driver.py @@ -12,7 +12,7 @@ groq = None # type: ignore[assignment] from ..infra.cost_mixin import CostMixin -from .base import Driver, _parse_tool_arguments +from .base import Driver, _apply_openai_tool_options, _normalize_stop_reason, _tool_call_dict logger = logging.getLogger(__name__) @@ -188,6 +188,8 @@ def generate_messages_with_tools( if supports_temperature and "temperature" in opts: kwargs["temperature"] = opts["temperature"] + _apply_openai_tool_options(kwargs, options) + resp = self.client.chat.completions.create(**kwargs) from .openai_driver import _extract_openai_cached_tokens @@ -217,20 +219,18 @@ def generate_messages_with_tools( choice = resp.choices[0] text = choice.message.content or "" - stop_reason = choice.finish_reason + raw_stop_reason = choice.finish_reason tool_calls_out: list[dict[str, Any]] = [] if choice.message.tool_calls: for tc in choice.message.tool_calls: - args = _parse_tool_arguments(tc.function.arguments, tc.function.name, stop_reason) tool_calls_out.append( - { - "id": tc.id, - "name": tc.function.name, - "arguments": args, - } + _tool_call_dict(getattr(tc, "id", None), tc.function.name, tc.function.arguments, raw_stop_reason) ) + stop_reason = _normalize_stop_reason(raw_stop_reason, tool_calls_present=bool(tool_calls_out)) + meta["raw_stop_reason"] = raw_stop_reason + result: dict[str, Any] = { "text": text, "meta": meta, diff --git a/prompture/drivers/openai_driver.py b/prompture/drivers/openai_driver.py index 5bb2959c..92a045a5 100644 --- a/prompture/drivers/openai_driver.py +++ b/prompture/drivers/openai_driver.py @@ -14,7 +14,7 @@ from ..infra.cost_mixin import CostMixin, prepare_strict_schema from ._prompt_cache import derive_prompt_cache_key -from .base import Driver, _parse_tool_arguments +from .base import Driver, _apply_openai_tool_options, _normalize_stop_reason, _tool_call_dict logger = logging.getLogger(__name__) @@ -127,13 +127,8 @@ def _extract_openai_tool_calls(message: Any, stop_reason: str | None) -> list[di tool_calls_out: list[dict[str, Any]] = [] if message.tool_calls: for tc in message.tool_calls: - args = _parse_tool_arguments(tc.function.arguments, tc.function.name, stop_reason) tool_calls_out.append( - { - "id": tc.id, - "name": tc.function.name, - "arguments": args, - } + _tool_call_dict(getattr(tc, "id", None), tc.function.name, tc.function.arguments, stop_reason) ) return tool_calls_out @@ -307,6 +302,7 @@ def generate_messages_with_tools( extra={"tools": tools}, prompt_cache_key=_openai_prompt_cache_key(messages, opts, tools), ) + _apply_openai_tool_options(kwargs, options) resp = self.client.chat.completions.create(**kwargs) @@ -321,8 +317,10 @@ def generate_messages_with_tools( choice = resp.choices[0] text = choice.message.content or "" - stop_reason = choice.finish_reason - tool_calls_out = _extract_openai_tool_calls(choice.message, stop_reason) + raw_stop_reason = choice.finish_reason + tool_calls_out = _extract_openai_tool_calls(choice.message, raw_stop_reason) + stop_reason = _normalize_stop_reason(raw_stop_reason, tool_calls_present=bool(tool_calls_out)) + meta["raw_stop_reason"] = raw_stop_reason return { "text": text, diff --git a/tests/test_live_events.py b/tests/test_live_events.py index bcfd9369..62b8689e 100644 --- a/tests/test_live_events.py +++ b/tests/test_live_events.py @@ -809,7 +809,10 @@ def fake_post(url, headers, json, stream, timeout): assert tool_stop.input == {"city": "Tokyo"} message_stop = events[-1] - assert message_stop.stop_reason == "tool_calls" + # stop_reason uses the shared vocabulary, so OpenAI's "tool_calls" + # surfaces as "tool_use"; the provider's raw value is preserved. + assert message_stop.stop_reason == "tool_use" + assert message_stop.usage["raw_stop_reason"] == "tool_calls" assert message_stop.usage["prompt_tokens"] == 12 assert message_stop.usage["completion_tokens"] == 9 # Cost = 0.0001 * 12 + 0.0005 * 9 = 0.0012 + 0.0045 = 0.0057 From ade646ec1972b2f23609c74f14e5c431224c63f2 Mon Sep 17 00:00:00 2001 From: Juan Denis <13461850+jhd3197@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:17:22 -0400 Subject: [PATCH 4/8] fix(tools): generate honest tool schemas and validate arguments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_python_type_to_json_schema` degraded most non-trivial annotations to `{"type": "string"}`: `int | None`, `Literal`, `Enum`, nested pydantic models, dataclasses, `TypedDict`, `dict[str, X]` and container element types. The schema handed to the model therefore disagreed with what the tool actually accepted, and coercion existed but was wired only to extraction, never to tool arguments. Now emitted faithfully: unions as `anyOf` (with `{"type": "null"}` for NoneType), `Literal`/`Enum` as `enum` (with a shared `type` when all values agree), containers with `items`/`prefixItems`/`additionalProperties`, date-ish scalars with a `format`, and nested structured types expanded via `pydantic.TypeAdapter`. Arguments are coerced towards the annotation and validated against the schema, returning an LLM-friendly message naming the offending argument so the model can self-correct instead of a bare TypeError escaping the tool. A *missing* annotation now yields the empty (unconstrained) schema rather than a string. An unannotated parameter means "type unknown", and since arguments are now validated, claiming `string` rejected every correct non-string call — `lambda x: x + 1` could no longer be called with an int. Adds direct coverage for the type mapping, which previously had none: 38 cases across unions, enumerations, containers, formats and nested models, plus end-to-end checks that a tool's advertised schema matches what it accepts and survives conversion to OpenAI wire format. Co-Authored-By: Claude Opus 5 (1M context) --- prompture/agents/tools_schema.py | 512 ++++++++++++++++++++++++++++--- tests/test_tools_schema.py | 309 ++++++++++++++++++- 2 files changed, 772 insertions(+), 49 deletions(-) diff --git a/prompture/agents/tools_schema.py b/prompture/agents/tools_schema.py index 5f430301..4ddddf20 100644 --- a/prompture/agents/tools_schema.py +++ b/prompture/agents/tools_schema.py @@ -24,9 +24,23 @@ def get_weather(city: str, units: str = "celsius") -> str: import inspect import json import logging -from collections.abc import Callable -from dataclasses import dataclass, field -from typing import Any, get_type_hints +import re +import uuid +from collections.abc import Callable, Mapping +from dataclasses import asdict, dataclass, field, is_dataclass +from datetime import date, datetime, time +from enum import Enum +from typing import Any, Literal, get_args, get_origin, get_type_hints, is_typeddict + +from pydantic import BaseModel, TypeAdapter + +from ..extraction.tools import _is_union_origin, convert_value +from ..infra.cost_mixin import prepare_strict_schema + +try: # same defensive pattern as prompture/extraction/validator.py + import jsonschema +except Exception: # pragma: no cover - jsonschema is a hard dependency + jsonschema = None logger = logging.getLogger("prompture.tools_schema") @@ -42,41 +56,189 @@ def get_weather(city: str, units: str = "celsius") -> str: float: "number", bool: "boolean", list: "array", + tuple: "array", dict: "object", } +def _json_type_name(value: Any) -> str | None: + """JSON Schema type name for a literal Python value (``None`` if unmappable).""" + if isinstance(value, bool): + return "boolean" + if isinstance(value, int): + return "integer" + if isinstance(value, float): + return "number" + if isinstance(value, str): + return "string" + if isinstance(value, list): + return "array" + if isinstance(value, dict): + return "object" + return None + + +def _enum_schema(values: list[Any]) -> dict[str, Any]: + """Build an ``enum`` schema, adding ``type`` when all values share one.""" + schema: dict[str, Any] = {"enum": values} + names = [_json_type_name(v) for v in values] + if values and all(n is not None and n == names[0] for n in names): + schema["type"] = names[0] + return schema + + +def _is_structured_type(annotation: Any) -> bool: + """True for nested structured types: pydantic models, dataclasses, TypedDicts.""" + if not isinstance(annotation, type): + return False + if is_typeddict(annotation): + return True + if is_dataclass(annotation): + return True + return issubclass(annotation, BaseModel) + + def _python_type_to_json_schema(annotation: Any) -> dict[str, Any]: - """Convert a Python type annotation to a JSON Schema snippet.""" - if annotation is inspect.Parameter.empty or annotation is None: - return {"type": "string"} + """Convert a Python type annotation to a JSON Schema snippet. + + Handles ``typing.Union`` and PEP-604 unions (``int | None``) as ``anyOf`` + (with ``{"type": "null"}`` for ``NoneType``), ``Literal`` and ``Enum`` as + ``enum``, ``datetime``/``date``/``time``/``UUID`` as formatted strings, + ``list``/``tuple``/``dict`` containers (``dict[str, X]`` gets + ``additionalProperties``), nested pydantic models, dataclasses and + ``TypedDict`` via :class:`pydantic.TypeAdapter`, and ``Any`` as the empty + schema. Unknown types fall back to ``{"type": "string"}``. + + A *missing* annotation yields the empty (unconstrained) schema rather than + a string: an unannotated parameter means "type unknown", and since + arguments are validated against this schema, claiming ``string`` would + reject every correct non-string call (e.g. ``lambda x: x + 1``). + """ + if annotation is inspect.Parameter.empty: + return {} + if annotation is Any or annotation is None: + return {} - # Handle Optional[X] (Union[X, None]) - origin = getattr(annotation, "__origin__", None) - args = getattr(annotation, "__args__", ()) + origin = get_origin(annotation) + args = get_args(annotation) - if origin is type(None): - return {"type": "string"} + # Union / Optional — both typing.Union[X, ...] and PEP-604 X | Y. + if _is_union_origin(origin): + return { + "anyOf": [{"type": "null"} if a is type(None) else _python_type_to_json_schema(a) for a in args], + } - # Union types (Optional) - if origin is not None and hasattr(origin, "__name__") and origin.__name__ == "Union": - non_none = [a for a in args if a is not type(None)] - if len(non_none) == 1: - return _python_type_to_json_schema(non_none[0]) + # Literal[...] values. + if origin is Literal: + return _enum_schema(list(args)) # list[X] - if origin is list and args: - return {"type": "array", "items": _python_type_to_json_schema(args[0])} + if origin is list: + return { + "type": "array", + "items": _python_type_to_json_schema(args[0]) if args else {}, + } + + # tuple[X, ...] / tuple[A, B] + if origin is tuple: + if args and args[-1] is Ellipsis: + return {"type": "array", "items": _python_type_to_json_schema(args[0])} + if args: + return { + "type": "array", + "prefixItems": [_python_type_to_json_schema(a) for a in args], + "items": False, + } + return {"type": "array"} # dict[str, X] if origin is dict: - return {"type": "object"} + schema: dict[str, Any] = {"type": "object"} + if len(args) > 1 and args[1] is not Any: + schema["additionalProperties"] = _python_type_to_json_schema(args[1]) + return schema + + # Enum subclasses. + if isinstance(annotation, type) and issubclass(annotation, Enum): + return _enum_schema([member.value for member in annotation]) + + # Date/time-ish scalars. + if annotation is datetime: + return {"type": "string", "format": "date-time"} + if annotation is date: + return {"type": "string", "format": "date"} + if annotation is time: + return {"type": "string", "format": "time"} + if annotation is uuid.UUID: + return {"type": "string", "format": "uuid"} + + # Nested structured types: pydantic BaseModel, dataclass, TypedDict. + if _is_structured_type(annotation): + try: + return TypeAdapter(annotation).json_schema() + except Exception as exc: # pragma: no cover - defensive + logger.debug("TypeAdapter schema generation failed for %r: %s", annotation, exc) # Simple types json_type = _TYPE_MAP.get(annotation, "string") return {"type": json_type} +def _make_nullable(schema: dict[str, Any]) -> dict[str, Any]: + """Return a copy of *schema* that also accepts ``null``.""" + schema = dict(schema) + if isinstance(schema.get("type"), str): + schema["type"] = [schema["type"], "null"] + elif isinstance(schema.get("type"), list): + if "null" not in schema["type"]: + schema["type"] = [*schema["type"], "null"] + elif isinstance(schema.get("anyOf"), list): + if {"type": "null"} not in schema["anyOf"]: + schema["anyOf"] = [*schema["anyOf"], {"type": "null"}] + else: + schema = {"anyOf": [schema, {"type": "null"}]} + return schema + + +def _strict_parameters(parameters: dict[str, Any]) -> dict[str, Any]: + """Normalise a tool-parameters schema for OpenAI strict tool use. + + Applies the same normalization as + :func:`prompture.infra.cost_mixin.prepare_strict_schema` + (``additionalProperties: false``, every property in ``required``) and + additionally makes originally-optional parameters (those with a default, + i.e. not in the source ``required`` list) nullable, since strict mode + forces them into ``required``. + """ + originally_required = set(parameters.get("required", []) or []) + strict_schema = prepare_strict_schema(parameters) + properties = strict_schema.get("properties") + if isinstance(properties, dict): + for key, prop in properties.items(): + if key not in originally_required and isinstance(prop, dict): + properties[key] = _make_nullable(prop) + return strict_schema + + +#: Valid tool names (OpenAI / Anthropic / Google all accept this shape). +_TOOL_NAME_RE = re.compile(r"^[a-zA-Z0-9_-]{1,64}$") + + +def _validate_tool_name(name: str) -> None: + """Raise ``ValueError`` if *name* is not a valid tool name. + + Provider APIs reject tool names outside ``^[a-zA-Z0-9_-]{1,64}$``; + validating at registration surfaces the problem at build time instead of + as a mid-conversation API error (e.g. unchecked tukuy skill names flowing + in via ``extraction/tukuy_bridge.py``). + """ + if not _TOOL_NAME_RE.match(name): + raise ValueError( + f"Invalid tool name {name!r}: must match ^[a-zA-Z0-9_-]{{1,64}}$ " + "(1-64 characters; letters, digits, underscore, hyphen)." + ) + + @dataclass class ToolDefinition: """Describes a single callable tool the LLM can invoke. @@ -97,19 +259,33 @@ class ToolDefinition: # Serialisation helpers # ------------------------------------------------------------------ - def to_openai_format(self) -> dict[str, Any]: - """Serialise to OpenAI ``tools`` array element format.""" - return { - "type": "function", - "function": { - "name": self.name, - "description": self.description, - "parameters": self.parameters, - }, + def to_openai_format(self, strict: bool = False) -> dict[str, Any]: + """Serialise to OpenAI ``tools`` array element format. + + With ``strict=True`` the parameters schema is normalised for OpenAI + strict tool use — ``additionalProperties: false`` on every object, + every property key listed in ``required``, and parameters that were + optional (i.e. have a default) made nullable — and ``"strict": true`` + is set on the function object. + """ + parameters = _strict_parameters(self.parameters) if strict else self.parameters + function: dict[str, Any] = { + "name": self.name, + "description": self.description, + "parameters": parameters, } + if strict: + function["strict"] = True + return {"type": "function", "function": function} def to_anthropic_format(self) -> dict[str, Any]: - """Serialise to Anthropic ``tools`` array element format.""" + """Serialise to Anthropic ``tools`` array element format. + + .. note:: + Maintenance-only: every built-in driver serialises tools from the + OpenAI shape (:meth:`to_openai_format`), so this converter is kept + for external consumers and is a candidate for future deprecation. + """ return { "name": self.name, "description": self.description, @@ -253,6 +429,21 @@ def _collapse(lines: list[str]) -> str: return " ".join(stripped for line in lines if (stripped := line.strip())) +def _strip_numpy_sections(text: str) -> str: + """Truncate *text* at the first NumPy-style section header (``Parameters``, + ``Returns``, … — a word-only line followed by a dash underline).""" + lines = text.split("\n") + for i in range(len(lines) - 1): + if re.match(r"^[A-Za-z][A-Za-z _]*$", lines[i].strip()) and _is_numpy_underline(lines[i + 1]): + return "\n".join(lines[:i]) + return text + + +def _strip_rest_fields(text: str) -> str: + """Drop reST field lines (``:param x:``, ``:return:``, …) from *text*.""" + return "\n".join(line for line in text.split("\n") if not _REST_FIELD_RE.match(line.strip())) + + def _docstring_description(docstring: str | None) -> str: """Build the tool description the model sees from *docstring*. @@ -260,9 +451,11 @@ def _docstring_description(docstring: str | None) -> str: that prose is usually the only place a tool's scope, caveats, and intended use are written down. ``Args:`` is dropped (it is already encoded in the parameter schema), while ``Returns:``/``Yields:`` are appended as a single - line so the model knows what it gets back. + line so the model knows what it gets back. NumPy-style sections and reST + field lines are stripped from the lead as well. """ lead, sections = _split_docstring(docstring) + lead = _strip_rest_fields(_strip_numpy_sections(lead)).strip() parts: list[str] = [lead] if lead else [] for label, keys in (("Returns", _RETURNS_HEADERS), ("Yields", _YIELDS_HEADERS)): @@ -286,7 +479,7 @@ def _truncate_description(text: str, limit: int | None) -> str: return head.rstrip(" \n.,;:") + "…" -def _parse_docstring_params(docstring: str | None) -> dict[str, str]: +def _parse_google_params(docstring: str | None) -> dict[str, str]: """Extract parameter descriptions from a Google-style docstring ``Args:`` section.""" _, sections = _split_docstring(docstring) lines = _first_section(sections, *_ARGS_HEADERS) @@ -334,6 +527,152 @@ def _parse_docstring_params(docstring: str | None) -> dict[str, str]: return params +def _is_numpy_underline(line: str) -> bool: + """True for a NumPy-style section underline (``----------``).""" + stripped = line.strip() + return len(stripped) >= 3 and set(stripped) == {"-"} + + +#: NumPy-style parameter entry header: ``name : type`` (``*``/``**`` allowed). +_NUMPY_PARAM_RE = re.compile(r"^\*{0,2}(\w+)\s*:") + +#: reST-style parameter field: ``:param name: desc`` or ``:param type name: desc``. +_REST_PARAM_RE = re.compile(r"^:param\s+(?:[\w.\[\], ]+\s+)?(\w+)\s*:\s*(.*)$") + +#: reST field lines that should not leak into the tool description. +_REST_FIELD_RE = re.compile(r"^:(param|type|return|rtype|raises?|yield|yields)\b") + + +def _parse_numpy_params(docstring: str | None) -> dict[str, str]: + """Extract parameter descriptions from a NumPy-style ``Parameters`` section. + + Expects the canonical shape:: + + Parameters + ---------- + x : int + Description of x. + """ + if not docstring: + return {} + lines = docstring.split("\n") + + # Locate the "Parameters" header (word-only line followed by a dash underline). + start = None + for i in range(len(lines) - 1): + if lines[i].strip().lower() == "parameters" and _is_numpy_underline(lines[i + 1]): + start = i + 2 + break + if start is None: + return {} + + params: dict[str, str] = {} + current_param: str | None = None + current_desc_parts: list[str] = [] + base_indent: int | None = None + + def _flush() -> None: + if current_param is not None: + params[current_param] = " ".join(current_desc_parts).strip() + + for j in range(start, len(lines)): + line = lines[j] + stripped = line.strip() + indent = len(line) - len(line.lstrip()) + + # The section ends at the next underline-headed section (e.g. "Returns"). + if ( + j + 1 < len(lines) + and stripped + and re.match(r"^[A-Za-z][A-Za-z _]*$", stripped) + and _is_numpy_underline(lines[j + 1]) + ): + break + + match = _NUMPY_PARAM_RE.match(stripped) + if match and (base_indent is None or indent <= base_indent): + if base_indent is None: + base_indent = indent + if indent == base_indent: + _flush() + current_param = match.group(1) + current_desc_parts = [] + continue + if current_param is not None and stripped: + current_desc_parts.append(stripped) + + _flush() + return params + + +def _parse_rest_params(docstring: str | None) -> dict[str, str]: + """Extract parameter descriptions from reST-style ``:param x:`` fields.""" + if not docstring: + return {} + params: dict[str, str] = {} + for line in docstring.split("\n"): + match = _REST_PARAM_RE.match(line.strip()) + if match: + params[match.group(1)] = match.group(2).strip() + return params + + +def _parse_docstring_params(docstring: str | None) -> dict[str, str]: + """Extract parameter descriptions from a docstring. + + Supports Google-style (``Args:``), NumPy-style (``Parameters`` followed by + a dash underline) and reST-style (``:param x:``) docstrings, tried in that + order; the first style that yields any descriptions wins. + """ + params = _parse_google_params(docstring) + if params: + return params + params = _parse_numpy_params(docstring) + if params: + return params + return _parse_rest_params(docstring) + + +def _coerce_argument(tool_name: str, param: str, value: Any, annotation: Any) -> Any: + """Coerce one tool argument to *annotation* via ``extraction.tools.convert_value``. + + Raises ``ValueError`` with an LLM-friendly message (naming the argument, + the expected type and the received value) so the model can self-correct. + """ + if get_origin(annotation) is Literal: + allowed = get_args(annotation) + if value in allowed: + return value + raise ValueError( + f"Invalid value for argument '{param}' of '{tool_name}': expected one of {list(allowed)!r}, got {value!r}." + ) + try: + return convert_value(value, annotation, field_name=param, use_defaults_on_failure=False) + except Exception as exc: + expected = getattr(annotation, "__name__", None) or str(annotation) + raise ValueError( + f"Invalid value for argument '{param}' of '{tool_name}': expected {expected}, got {value!r} ({exc})." + ) from exc + + +def _json_compatible(value: Any) -> Any: + """Convert a coerced argument value into a JSON-compatible value for + schema validation (Enums → values, datetimes → ISO strings, …).""" + if isinstance(value, Enum): + return _json_compatible(value.value) + if isinstance(value, (datetime, date, time)): + return value.isoformat() + if isinstance(value, BaseModel): + return _json_compatible(value.model_dump()) + if is_dataclass(value) and not isinstance(value, type): + return _json_compatible(asdict(value)) + if isinstance(value, Mapping): + return {str(k): _json_compatible(v) for k, v in value.items()} + if isinstance(value, (list, tuple, set)): + return [_json_compatible(v) for v in value] + return value + + def tool_from_function( fn: Callable[..., Any], *, @@ -365,6 +704,16 @@ def tool_from_function( except Exception: hints = {} + if not param_docs and any( + pname != "self" and hints.get(pname, p.annotation) is not inspect.Parameter.empty + for pname, p in sig.parameters.items() + ): + logger.debug( + "No parseable parameter docs found for %r (Google/NumPy/reST styles supported); " + "falling back to parameter names for descriptions.", + tool_name, + ) + properties: dict[str, Any] = {} required: list[str] = [] @@ -434,8 +783,13 @@ def register( description: str | None = None, max_description_chars: int | None = MAX_TOOL_DESCRIPTION_CHARS, ) -> ToolDefinition: - """Register *fn* as a tool and return the :class:`ToolDefinition`.""" + """Register *fn* as a tool and return the :class:`ToolDefinition`. + + Raises: + ValueError: If the tool name is not valid (``^[a-zA-Z0-9_-]{1,64}$``). + """ td = tool_from_function(fn, name=name, description=description, max_description_chars=max_description_chars) + _validate_tool_name(td.name) self._tools[td.name] = td return td @@ -448,7 +802,12 @@ def tool(self, fn: Callable[..., Any]) -> Callable[..., Any]: return fn def add(self, tool_def: ToolDefinition) -> None: - """Add a pre-built :class:`ToolDefinition`.""" + """Add a pre-built :class:`ToolDefinition`. + + Raises: + ValueError: If the tool name is not valid (``^[a-zA-Z0-9_-]{1,64}$``). + """ + _validate_tool_name(tool_def.name) self._tools[tool_def.name] = tool_def # ------------------------------------------------------------------ @@ -575,6 +934,7 @@ def add_tukuy_skill( td = ToolDefinition(name=name, description=td.description, parameters=td.parameters, function=td.function) if description: td = ToolDefinition(name=td.name, description=description, parameters=td.parameters, function=td.function) + _validate_tool_name(td.name) self._tools[td.name] = td return td @@ -599,8 +959,13 @@ def add_tukuy_skills( # Serialisation # ------------------------------------------------------------------ - def to_openai_format(self) -> list[dict[str, Any]]: - return [td.to_openai_format() for td in self._tools.values()] + def to_openai_format(self, strict: bool = False) -> list[dict[str, Any]]: + """Serialise all tools to the OpenAI ``tools`` array format. + + With ``strict=True`` each tool's parameters schema is normalised for + OpenAI strict tool use (see :meth:`ToolDefinition.to_openai_format`). + """ + return [td.to_openai_format(strict=strict) for td in self._tools.values()] def to_anthropic_format(self) -> list[dict[str, Any]]: return [td.to_anthropic_format() for td in self._tools.values()] @@ -639,6 +1004,66 @@ def _validate_arguments(td: ToolDefinition, arguments: dict[str, Any]) -> str | f"You sent: {json.dumps(arguments) if arguments else '{} (empty)'}" ) + @staticmethod + def _coerce_and_validate_arguments( + td: ToolDefinition, arguments: dict[str, Any] + ) -> tuple[dict[str, Any], str | None]: + """Coerce *arguments* to the function's annotated types and validate them. + + Returns ``(coerced_arguments, error)``. *error* is an LLM-friendly + message (``None`` on success) describing exactly which argument is + wrong and what was expected, so the model can self-correct instead of + a bare ``TypeError`` escaping from the tool function. + """ + schema = td.parameters if isinstance(td.parameters, dict) else {} + properties = schema.get("properties", {}) or {} + + # Reject unknown extra keys (only when the schema declares properties). + if properties: + unknown = [k for k in arguments if k not in properties] + if unknown: + return arguments, ( + f"Unknown argument(s) for '{td.name}': {', '.join(sorted(unknown))}. " + f"Valid arguments are: {', '.join(properties)}. " + f"You sent: {json.dumps(arguments, default=str)}" + ) + + # Coerce each argument towards the function's type annotation. + try: + hints = get_type_hints(td.function) + except Exception: + hints = {} + coerced = dict(arguments) + for key, value in arguments.items(): + annotation = hints.get(key) + if annotation is None or annotation is Any: + continue + try: + coerced[key] = _coerce_argument(td.name, key, value, annotation) + except ValueError as exc: + return arguments, str(exc) + + # Validate the coerced arguments against the tool's JSON Schema. + if jsonschema is not None and schema: + instance = {k: _json_compatible(v) for k, v in coerced.items()} + errors = sorted( + jsonschema.Draft7Validator(schema).iter_errors(instance), + key=lambda e: list(e.path), + ) + if errors: + details = [] + for e in errors[:5]: + where = "/".join(str(p) for p in e.path) or "(root)" + details.append(f" - {where}: {e.message}") + return arguments, ( + f"Invalid argument(s) for '{td.name}':\n" + + "\n".join(details) + + "\nExpected arguments matching schema: " + + json.dumps(schema, default=str) + + f"\nYou sent: {json.dumps(arguments, default=str)}" + ) + return coerced, None + # ------------------------------------------------------------------ # Execution # ------------------------------------------------------------------ @@ -646,6 +1071,11 @@ def _validate_arguments(td: ToolDefinition, arguments: dict[str, Any]) -> str | def execute(self, name: str, arguments: dict[str, Any]) -> Any: """Execute a registered tool by name with the given arguments. + Arguments are coerced to the function's annotated types and validated + against the tool's JSON Schema; on failure an LLM-friendly error + string is returned (so the model can self-correct) instead of a bare + ``TypeError`` escaping. + Raises: KeyError: If no tool with *name* is registered. """ @@ -655,7 +1085,10 @@ def execute(self, name: str, arguments: dict[str, Any]) -> Any: error = self._validate_arguments(td, arguments) if error: return error - return td.function(**arguments) + coerced, error = self._coerce_and_validate_arguments(td, arguments) + if error: + return error + return td.function(**coerced) async def aexecute(self, name: str, arguments: dict[str, Any]) -> Any: """Execute a registered tool, awaiting async tool functions. @@ -673,13 +1106,16 @@ async def aexecute(self, name: str, arguments: dict[str, Any]) -> Any: if td is None: raise KeyError(f"Tool not registered: {name!r}") error = self._validate_arguments(td, arguments) + if error: + return error + coerced, error = self._coerce_and_validate_arguments(td, arguments) if error: return error # Prefer dedicated async wrapper (set by tukuy bridge) async_fn = getattr(td.function, "_async_fn", None) if async_fn is not None: - return await async_fn(**arguments) - result = td.function(**arguments) + return await async_fn(**coerced) + result = td.function(**coerced) if inspect.isawaitable(result): return await result return result diff --git a/tests/test_tools_schema.py b/tests/test_tools_schema.py index 26239498..d99d4052 100644 --- a/tests/test_tools_schema.py +++ b/tests/test_tools_schema.py @@ -2,9 +2,20 @@ from __future__ import annotations +import enum +from dataclasses import dataclass +from datetime import date, datetime +from typing import Any, Literal, Optional, TypedDict, Union + import pytest +from pydantic import BaseModel -from prompture.agents.tools_schema import ToolRegistry +from prompture.agents.tools_schema import ( + ToolDefinition, + ToolRegistry, + _python_type_to_json_schema, + tool_from_function, +) # --------------------------------------------------------------------------- # Helpers @@ -151,14 +162,290 @@ def test_subset_tools_are_executable(self): result = sub.execute("file_read", {"path": "/tmp/test"}) assert result == "/tmp/test" - def test_filter_tools_are_executable(self): - reg = _build_registry() - sub = reg.filter(lambda td: td.name == "python_execute") - result = sub.execute("python_execute", {"code": "print('hi')"}) - assert result == "print('hi')" - def test_exclude_tools_are_executable(self): - reg = _build_registry() - sub = reg.exclude({"file_read"}) - result = sub.execute("file_write", {"path": "/tmp/f", "content": "data"}) - assert result == "/tmp/f: data" +# --------------------------------------------------------------------------- +# Unannotated parameters are unconstrained, not silently "string" +# --------------------------------------------------------------------------- + + +class TestUnannotatedParameters: + """An absent annotation means "type unknown", so the schema must not + constrain it. Claiming ``{"type": "string"}`` both misinforms the model + and — now that arguments are validated against the schema — rejects every + correct non-string call.""" + + def test_missing_annotation_yields_empty_schema(self): + import inspect + + assert _python_type_to_json_schema(inspect.Parameter.empty) == {} + + def test_lambda_tool_accepts_the_type_it_actually_wants(self): + reg = ToolRegistry() + reg.register(lambda x: x + 1, name="inc", description="Increment") + assert reg.execute("inc", {"x": 5}) == 6 + + def test_unannotated_def_parameter_is_not_typed(self): + def echo(value): + """Echo a value back. + + Args: + value: Anything at all. + """ + return value + + td = tool_from_function(echo) + prop = td.parameters["properties"]["value"] + assert "type" not in prop + # The description is still carried through for the model's benefit. + assert prop["description"] == "Anything at all." + + def test_unannotated_parameter_accepts_every_json_type(self): + reg = ToolRegistry() + reg.register(lambda value: value, name="echo", description="Echo") + for payload in (1, "s", 1.5, True, None, [1, 2], {"k": "v"}): + assert reg.execute("echo", {"value": payload}) == payload + + def test_annotated_parameters_still_validate(self): + """Loosening unannotated params must not loosen annotated ones.""" + + def strict(count: int) -> int: + """Take an int. + + Args: + count: A real integer. + """ + return count + + reg = ToolRegistry() + reg.register(strict) + assert reg.execute("strict", {"count": 3}) == 3 + # A coercible string is still accepted... + assert reg.execute("strict", {"count": "7"}) == 7 + # ...but junk is reported back to the model rather than executed. + error = reg.execute("strict", {"count": "not-an-int"}) + assert isinstance(error, str) + assert "expected int" in error + + +# --------------------------------------------------------------------------- +# Type -> JSON Schema coverage +# +# These annotations all used to collapse to {"type": "string"}, so the schema +# the model was handed disagreed with what the tool actually accepted. +# --------------------------------------------------------------------------- + + +class Color(enum.Enum): + RED = "red" + BLUE = "blue" + + +class Point(BaseModel): + x: int + y: int + + +@dataclass +class Box: + w: int + h: int + + +class Movie(TypedDict): + title: str + year: int + + +class TestPythonTypeToJsonSchema: + @pytest.mark.parametrize( + ("annotation", "expected"), + [ + (int, {"type": "integer"}), + (str, {"type": "string"}), + (float, {"type": "number"}), + (bool, {"type": "boolean"}), + ], + ) + def test_scalars(self, annotation, expected): + assert _python_type_to_json_schema(annotation) == expected + + def test_any_is_unconstrained(self): + assert _python_type_to_json_schema(Any) == {} + + def test_unknown_type_falls_back_to_string(self): + assert _python_type_to_json_schema(object) == {"type": "string"} + + # -- unions --------------------------------------------------------- + + def test_pep604_optional(self): + assert _python_type_to_json_schema(int | None) == {"anyOf": [{"type": "integer"}, {"type": "null"}]} + + def test_typing_optional_matches_pep604(self): + # The legacy spelling is the point of this test, so keep it verbatim. + legacy = Optional[int] # noqa: UP045 + assert _python_type_to_json_schema(legacy) == _python_type_to_json_schema(int | None) + + def test_union_without_none(self): + assert _python_type_to_json_schema(Union[int, str]) == {"anyOf": [{"type": "integer"}, {"type": "string"}]} + + def test_multi_member_union_keeps_every_member(self): + assert _python_type_to_json_schema(str | int | None) == { + "anyOf": [{"type": "string"}, {"type": "integer"}, {"type": "null"}] + } + + # -- enumerations --------------------------------------------------- + + def test_literal_becomes_enum(self): + assert _python_type_to_json_schema(Literal["a", "b"]) == {"enum": ["a", "b"], "type": "string"} + + def test_mixed_literal_omits_type(self): + """A Literal spanning types cannot claim a single JSON type.""" + assert _python_type_to_json_schema(Literal["a", 1]) == {"enum": ["a", 1]} + + def test_enum_class_uses_member_values(self): + assert _python_type_to_json_schema(Color) == {"enum": ["red", "blue"], "type": "string"} + + # -- containers ----------------------------------------------------- + + def test_list_of_scalars(self): + assert _python_type_to_json_schema(list[int]) == {"type": "array", "items": {"type": "integer"}} + + def test_bare_list_has_no_item_constraint(self): + assert _python_type_to_json_schema(list) == {"type": "array"} + + def test_homogeneous_tuple(self): + assert _python_type_to_json_schema(tuple[int, ...]) == {"type": "array", "items": {"type": "integer"}} + + def test_fixed_length_tuple_uses_prefix_items(self): + assert _python_type_to_json_schema(tuple[int, str]) == { + "type": "array", + "prefixItems": [{"type": "integer"}, {"type": "string"}], + "items": False, + } + + def test_dict_value_type_becomes_additional_properties(self): + assert _python_type_to_json_schema(dict[str, int]) == { + "type": "object", + "additionalProperties": {"type": "integer"}, + } + + def test_dict_of_any_stays_unconstrained(self): + assert _python_type_to_json_schema(dict[str, Any]) == {"type": "object"} + + # -- formatted scalars ---------------------------------------------- + + @pytest.mark.parametrize( + ("annotation", "fmt"), + [(datetime, "date-time"), (date, "date")], + ) + def test_datetime_scalars_carry_a_format(self, annotation, fmt): + assert _python_type_to_json_schema(annotation) == {"type": "string", "format": fmt} + + # -- nested structured types ---------------------------------------- + + def test_nested_pydantic_model_expands(self): + schema = _python_type_to_json_schema(Point) + assert schema["type"] == "object" + assert set(schema["properties"]) == {"x", "y"} + assert schema["properties"]["x"]["type"] == "integer" + assert sorted(schema["required"]) == ["x", "y"] + + def test_nested_dataclass_expands(self): + schema = _python_type_to_json_schema(Box) + assert set(schema["properties"]) == {"w", "h"} + + def test_typed_dict_expands(self): + schema = _python_type_to_json_schema(Movie) + assert schema["properties"]["year"]["type"] == "integer" + + def test_list_of_models_expands_items(self): + schema = _python_type_to_json_schema(list[Point]) + assert schema["type"] == "array" + assert set(schema["items"]["properties"]) == {"x", "y"} + + +class TestToolFromFunctionRichTypes: + """End-to-end: a tool advertises a schema matching what it accepts.""" + + def test_optional_literal_and_enum_params_are_faithful(self): + def configure( + mode: Literal["fast", "slow"], + color: Color, + retries: int | None = None, + ) -> str: + """Configure something. + + Args: + mode: How fast to go. + color: Which colour. + retries: Optional retry count. + """ + return f"{mode}/{color.value}/{retries}" + + td = tool_from_function(configure) + props = td.parameters["properties"] + + assert props["mode"]["enum"] == ["fast", "slow"] + assert props["color"]["enum"] == ["red", "blue"] + assert props["retries"]["anyOf"] == [{"type": "integer"}, {"type": "null"}] + # Only the non-defaulted params are required. + assert td.parameters["required"] == ["mode", "color"] + # Docstring descriptions survive alongside the richer types. + assert props["mode"]["description"] == "How fast to go." + + def test_nested_model_param_is_advertised_as_an_object(self): + def move(target: Point) -> str: + """Move to a point. + + Args: + target: Where to go. + """ + return f"{target.x},{target.y}" + + props = tool_from_function(move).parameters["properties"] + assert props["target"]["type"] == "object" + assert set(props["target"]["properties"]) == {"x", "y"} + + def test_dict_param_advertises_its_value_type(self): + def tally(counts: dict[str, int]) -> int: + """Sum a mapping. + + Args: + counts: Name to count. + """ + return sum(counts.values()) + + props = tool_from_function(tally).parameters["properties"] + assert props["counts"]["additionalProperties"] == {"type": "integer"} + + def test_registry_executes_a_literal_typed_tool(self): + """The schema is honest, so a valid call runs and an invalid one is + reported back to the model instead of raising.""" + + def pick(mode: Literal["fast", "slow"]) -> str: + """Pick a mode. + + Args: + mode: Which mode. + """ + return f"picked {mode}" + + reg = ToolRegistry() + reg.register(pick) + assert reg.execute("pick", {"mode": "fast"}) == "picked fast" + error = reg.execute("pick", {"mode": "sideways"}) + assert isinstance(error, str) + assert "sideways" in error + + def test_rich_schema_survives_conversion_to_openai_format(self): + def configure(mode: Literal["fast", "slow"]) -> str: + """Configure. + + Args: + mode: Which mode. + """ + return mode + + td: ToolDefinition = tool_from_function(configure) + wire = td.to_openai_format() + assert wire["function"]["parameters"]["properties"]["mode"]["enum"] == ["fast", "slow"] From 4eca1d405f9b2e22e9510b4d88431c12c457a990 Mon Sep 17 00:00:00 2001 From: Juan Denis <13461850+jhd3197@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:17:46 -0400 Subject: [PATCH 5/8] fix(agents): close approval, stop and timeout holes in the tool loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An async `on_approval_needed` handler auto-approved everything. The callback was typed as possibly-awaitable but invoked bare from the sync agent, so an `async def` handler returned a coroutine — always truthy — and every dangerous tool was silently approved, defeating the only human-in-the-loop mechanism. Callbacks now run through a bridge that drives awaitables to completion, carrying contextvars across the thread hop when a loop is already running. `Agent.stop()` was dead code: `_stop_requested` was written in six places and read nowhere, and the one test asserted only that the flag was set, so the suite certified a no-op. `stop()` now forwards to the conversation driving the run. It is tracked in `_active_conversation` rather than `_conversation`, because the latter is only populated when `persistent_conversation=True` — forwarding through it alone left `stop()` a silent no-op for the default agent. A `stop()` arriving before the conversation exists is replayed onto it rather than lost to the race. Exhausting `max_iterations` raised RuntimeError and discarded the whole run. The loop now takes one final turn with tools removed, so the model answers with what it has and `max_rounds_reached` records why — the worst case was exactly the weak local models that prompted tools target. Async tools ran coroutines via `asyncio.run` on a throwaway thread, which blocked the loop, dropped contextvars (tukuy `SecurityContext`, `current_tool_call_id`) and could deadlock; they now await natively. Adds a per-tool timeout, so a hung tool no longer wedges a run forever, exposed as `Agent(tool_timeout=...)` and overridable per call via `options={"tool_timeout": ...}`. Independent async calls execute concurrently via `asyncio.gather`; `zip(tool_calls, results)` pairs results with `strict=True` so a mis-pairing fails loudly instead of attaching a result to the wrong `tool_call_id`. Remaining sync/async drift and loop hardening: history trimming no longer orphans a `tool` message from its assistant `tool_calls` (which providers reject with a 400), malformed and truncated arguments are reported back to the model instead of executed as `{}`, tool results have a default size bound, `current_tool_call_id` is set on the sync path too, depth guards apply to every run path, `RunContext` is rebuilt per tool invocation so tools see live iteration and history, and `AgentStep` records the real tool name and a populated `duration_ms` instead of the call id and a dead field. Agent-level pre-call budget enforcement is removed: it read a freshly created, always-empty UsageSession and so never triggered. Enforcement lives in the Conversation, which sees real usage. Co-Authored-By: Claude Opus 5 (1M context) --- prompture/agents/agent.py | 326 ++++++++++--- prompture/agents/async_agent.py | 522 ++++++++++++++------ prompture/agents/async_conversation.py | 343 ++++++++++--- prompture/agents/conversation.py | 318 ++++++++++-- prompture/agents/types.py | 17 +- tests/test_agent.py | 639 +++++++++++++++++++++++++ tests/test_conversation_robustness.py | 550 +++++++++++++++++++++ tests/test_simulated_tools.py | 16 +- tests/test_tool_use.py | 20 +- 9 files changed, 2425 insertions(+), 326 deletions(-) create mode 100644 tests/test_conversation_robustness.py diff --git a/prompture/agents/agent.py b/prompture/agents/agent.py index d06b2cec..cf81850d 100644 --- a/prompture/agents/agent.py +++ b/prompture/agents/agent.py @@ -26,6 +26,8 @@ from __future__ import annotations +import asyncio +import concurrent.futures import contextvars import inspect import json @@ -39,7 +41,7 @@ from ..drivers.base import Driver from ..extraction.tools import clean_json_text -from ..infra.budget import BudgetPolicy, BudgetState, enforce_budget, resolve_budget_policy +from ..infra.budget import BudgetPolicy, resolve_budget_policy from ..infra.callbacks import DriverCallbacks from ..infra.provider_env import ProviderEnvironment from ..infra.session import UsageSession @@ -131,6 +133,34 @@ def _get_first_param_name(fn: Callable[..., Any]) -> str: return "" +def _invoke_cb_sync(callback: Callable[..., Any], *cb_args: Any) -> Any: + """Invoke *callback* from sync code, driving awaitables to completion. + + Async callbacks (e.g. ``async def on_approval_needed(...)``) return a + coroutine when called bare — which is always truthy and previously + caused silent auto-approval. This bridge detects awaitables and runs + them: directly with :func:`asyncio.run` when no loop is running, or on + a worker thread (carrying over ``contextvars``) when a loop is already + running in the current thread. + """ + result = callback(*cb_args) + if not inspect.isawaitable(result): + return result + + async def _await_result() -> Any: + return await result + + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + if loop is not None and loop.is_running(): + ctx_snapshot = contextvars.copy_context() + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + return pool.submit(ctx_snapshot.run, lambda: asyncio.run(_await_result())).result() + return asyncio.run(_await_result()) + + # ------------------------------------------------------------------ # Agent # ------------------------------------------------------------------ @@ -180,6 +210,11 @@ class Agent(Generic[DepsType]): skill_config: Optional configuration dict injected as a :class:`SkillContext` into tukuy skill ``invoke()`` calls. When ``None`` (default) no config is injected. + tool_timeout: Per-tool wall-clock budget in seconds. A tool that + overruns it yields an error result the model can react to + instead of wedging the run forever. ``None`` (default) means + no timeout. Can be overridden per call via + ``options={"tool_timeout": ...}``. """ def __init__( @@ -209,6 +244,7 @@ def __init__( auto_approve_safe_only: bool = False, skill_config: dict[str, Any] | None = None, max_tool_result_length: int | None = None, + tool_timeout: float | None = None, max_depth: int = _DEFAULT_MAX_AGENT_DEPTH, env: ProviderEnvironment | None = None, ) -> None: @@ -240,7 +276,12 @@ def __init__( self._auto_approve_safe_only = auto_approve_safe_only self._skill_config = skill_config self._max_tool_result_length = max_tool_result_length + self._tool_timeout = tool_timeout self._conversation: Conversation | None = None + # The conversation driving the current run. Tracked separately from + # ``_conversation`` (which only holds persistent ones) so that + # ``stop()`` can reach the loop of a non-persistent run too. + self._active_conversation: Conversation | None = None # Build internal tool registry self._tools = ToolRegistry() @@ -279,8 +320,21 @@ def state(self) -> AgentState: return self._lifecycle def stop(self) -> None: - """Request graceful shutdown after the current iteration.""" + """Request graceful shutdown after the current iteration. + + Sets the agent-level flag and forwards the request to the + conversation driving the current run (cooperative stop) so its + tool-round loop exits gracefully between rounds instead of starting + another round. Works for non-persistent agents too: the in-flight + conversation is tracked in ``_active_conversation`` for the duration + of the run, so the flag is never merely decorative. + + Safe to call from a tool, a callback, or another thread. + """ self._stop_requested = True + conv = self._active_conversation or self._conversation + if conv is not None: + conv.request_stop() @property def callbacks(self) -> AgentCallbacks: @@ -375,6 +429,7 @@ def run(self, prompt: str, *, deps: Any = None) -> AgentResult: token = _agent_depth.set(current_depth + 1) self._lifecycle = AgentState.running self._stop_requested = False + self._active_conversation = None steps: list[AgentStep] = [] try: @@ -409,20 +464,61 @@ def _build_run_context( prompt=prompt, ) + def _make_live_ctx_fn( + self, + prompt: str, + deps: Any, + session: UsageSession, + run_state: dict[str, Any], + ) -> Callable[[], RunContext[Any]]: + """Return a factory that builds a fresh :class:`RunContext` per call. + + The factory reads the current conversation from *run_state* (key + ``"conv"``) so tools invoked in later tool rounds see live + ``iteration``/``messages``/``usage`` instead of the snapshot taken + at the start of the run. ``run_state["conv"]`` must be assigned + once the conversation is built; before that, an empty context is + produced. + """ + + def _live_ctx() -> RunContext[Any]: + conv = run_state.get("conv") + messages = conv.messages if conv is not None else [] + iteration = max( + 0, + sum(1 for m in messages if m.get("role") == "assistant" and m.get("tool_calls")) - 1, + ) + return self._build_run_context(prompt, deps, session, messages, iteration) + + return _live_ctx + # ------------------------------------------------------------------ # Tool wrapping (RunContext injection + ModelRetry + callbacks) # ------------------------------------------------------------------ - def _wrap_tools_with_context(self, ctx: RunContext[Any], session: UsageSession | None = None) -> ToolRegistry: + def _wrap_tools_with_context( + self, + ctx: RunContext[Any], + session: UsageSession | None = None, + ctx_fn: Callable[[], RunContext[Any]] | None = None, + tool_timings: list[dict[str, Any]] | None = None, + ) -> ToolRegistry: """Return a new :class:`ToolRegistry` with wrapped tool functions. For each registered tool: - - If the tool's first param is ``RunContext``, inject *ctx* automatically. + - If the tool's first param is ``RunContext``, inject the live context + automatically. When *ctx_fn* is provided it is called at each tool + invocation so tools see the current iteration/messages; otherwise the + static *ctx* snapshot is used. - Catch :class:`ModelRetry` and convert to an error string. - - Fire ``agent_callbacks.on_tool_start`` / ``on_tool_end``. + - Fire ``agent_callbacks.on_tool_start`` / ``on_tool_end`` (awaitable + callbacks are driven to completion via :func:`_invoke_cb_sync`). - Strip the ``RunContext`` parameter from the JSON schema sent to the LLM. - If the tool wraps a child agent (via ``as_tool``), aggregate its usage into the parent *session*. + - When *tool_timings* is provided, append a + ``{"name", "timestamp", "duration_ms"}`` record per invocation so + step extraction can populate ``AgentStep.duration_ms``. """ if not self._tools: return ToolRegistry() @@ -442,10 +538,14 @@ def _make_wrapper( _name: str, _cb: AgentCallbacks = cb, _session: UsageSession | None = session, + _ctx_fn: Callable[[], RunContext[Any]] | None = ctx_fn, + _timings: list[dict[str, Any]] | None = tool_timings, ) -> Callable[..., Any]: def wrapper(**kwargs: Any) -> Any: + call_ctx = _ctx_fn() if _ctx_fn is not None else ctx + start = time.perf_counter() if _cb.on_tool_start: - _cb.on_tool_start(_name, kwargs) + _invoke_cb_sync(_cb.on_tool_start, _name, kwargs) try: # Inject skill config via SkillContext for tukuy skills if self._skill_config is not None: @@ -474,18 +574,18 @@ def wrapper(**kwargs: Any) -> Any: ) if _wants: - result = _fn(ctx, **kwargs) + result = _fn(call_ctx, **kwargs) else: result = _fn(**kwargs) except ApprovalRequired as exc: # Handle approval request if _cb.on_approval_needed: - approved = _cb.on_approval_needed(exc.tool_name, exc.action, exc.details) + approved = _invoke_cb_sync(_cb.on_approval_needed, exc.tool_name, exc.action, exc.details) if approved: # Retry the tool call after approval try: if _wants: - result = _fn(ctx, **kwargs) + result = _fn(call_ctx, **kwargs) else: result = _fn(**kwargs) except ApprovalRequired: @@ -518,8 +618,17 @@ def wrapper(**kwargs: Any) -> Any: } ) + if _timings is not None: + _timings.append( + { + "name": _name, + "timestamp": time.time(), + "duration_ms": (time.perf_counter() - start) * 1000, + } + ) + if _cb.on_tool_end: - _cb.on_tool_end(_name, result) + _invoke_cb_sync(_cb.on_tool_end, _name, result) return result return wrapper @@ -718,16 +827,30 @@ def _build_conversation( # Reuse existing conversation in persistent mode if self._persistent_conversation and self._conversation is not None: conv = self._conversation - # Update tools and callbacks for this run - if tools is not None: - conv._tools = tools - if driver_callbacks is not None: - conv._driver.callbacks = driver_callbacks - # Propagate before_turn hook (subclasses may set this) - hook = getattr(self, "_before_turn_hook", None) - if hook is not None: - conv._before_turn = hook - return conv + # If the agent's driver was swapped after the conversation was + # built (e.g. a budget fallback forced a rebuild), the cached + # conversation still holds the old driver — discard it and fall + # through to build a fresh one. + driver_stale = self._driver is not None and getattr(conv, "_driver", None) is not self._driver + if not driver_stale: + # Respect the newly resolved system prompt for this run. + if system_prompt is not None: + conv.system_prompt = system_prompt + # Update tools and callbacks for this run. + # NOTE: assigning ``callbacks`` mutates the driver instance, + # which may be shared (e.g. passed to several agents). A + # persistent-conversation agent should treat its driver as + # exclusively owned by the agent. + if tools is not None: + conv._tools = tools + if driver_callbacks is not None: + conv._driver.callbacks = driver_callbacks + # Propagate before_turn hook (subclasses may set this) + hook = getattr(self, "_before_turn_hook", None) + if hook is not None: + conv._before_turn = hook + return self._track_active(conv) + self._conversation = None effective_tools = tools if tools is not None else (self._tools if self._tools else None) @@ -742,6 +865,8 @@ def _build_conversation( kwargs["before_turn"] = hook if self._max_tool_result_length is not None: kwargs["max_tool_result_length"] = self._max_tool_result_length + if self._tool_timeout is not None: + kwargs["tool_timeout"] = self._tool_timeout if self._options: kwargs["options"] = self._options if driver_callbacks is not None: @@ -769,6 +894,19 @@ def _build_conversation( conv = Conversation(**kwargs) if self._persistent_conversation: self._conversation = conv + return self._track_active(conv) + + def _track_active(self, conv: Conversation) -> Conversation: + """Record *conv* as the conversation driving the current run. + + Lets :meth:`stop` reach the tool-round loop of a non-persistent run. + If ``stop()`` was already called between the run starting and the + conversation being built, the request is replayed onto *conv* so it + is not lost to that race. + """ + self._active_conversation = conv + if self._stop_requested: + conv.request_stop() return conv def _execute(self, prompt: str, steps: list[AgentStep], deps: Any) -> AgentResult: @@ -793,8 +931,16 @@ def _execute(self, prompt: str, steps: list[AgentStep], deps: Any) -> AgentResul # 4. Resolve system prompt (call it if callable, passing ctx) resolved_system_prompt = self._resolve_system_prompt(ctx) - # 5. Wrap tools with context (pass session for child agent usage aggregation) - wrapped_tools = self._wrap_tools_with_context(ctx, session) + # 5. Wrap tools with context (pass session for child agent usage + # aggregation; ctx_fn refreshes RunContext per tool round) + run_state: dict[str, Any] = {} + tool_timings: list[dict[str, Any]] = [] + wrapped_tools = self._wrap_tools_with_context( + ctx, + session, + ctx_fn=self._make_live_ctx_fn(prompt, deps, session, run_state), + tool_timings=tool_timings, + ) # 6. Build Conversation conv = self._build_conversation( @@ -802,36 +948,16 @@ def _execute(self, prompt: str, steps: list[AgentStep], deps: Any) -> AgentResul tools=wrapped_tools if wrapped_tools else None, driver_callbacks=driver_callbacks, ) + run_state["conv"] = conv # 7. Fire on_iteration callback if self._agent_callbacks.on_iteration: self._agent_callbacks.on_iteration(0) - # 8. Enforce budget policy at the agent level (pre-call) - if self._budget_policy is not None: - state = BudgetState( - cost_used=session.cost, - tokens_used=session.total_tokens, - max_cost=self._max_cost, - max_tokens=self._max_tokens, - ) - new_model = enforce_budget( - state, - self._budget_policy, - fallback_models=self._fallback_models, - current_model=self._model, - on_model_fallback=self._on_model_fallback, - ) - if new_model is not None: - self._model = new_model - self._driver = None # force rebuild - conv = self._build_conversation( - system_prompt=resolved_system_prompt, - tools=wrapped_tools if wrapped_tools else None, - driver_callbacks=driver_callbacks, - ) - - # 9. Ask the conversation (handles full tool loop internally) + # 8. Ask the conversation (handles full tool loop internally) + # Note: budget policy is enforced inside the Conversation + # (``_check_budget``); the agent-level pre-call check was removed + # because it read a fresh, always-empty UsageSession. agent_name = self.name or self.__class__.__name__ with tracker.agent(agent_name): t0 = time.perf_counter() @@ -856,6 +982,7 @@ def _execute(self, prompt: str, steps: list[AgentStep], deps: Any) -> AgentResul steps, all_tool_calls, getattr(conv, "_full_tool_results", None), + tool_timings, ) # Handle output_type parsing @@ -906,11 +1033,42 @@ def _extract_steps( steps: list[AgentStep], all_tool_calls: list[dict[str, Any]], full_tool_results: dict[str, str] | None = None, + tool_timings: list[dict[str, Any]] | None = None, ) -> None: - """Scan conversation messages and populate steps and tool_calls.""" + """Scan conversation messages and populate steps and tool_calls. + + Args: + messages: Conversation messages to scan. + steps: List to append :class:`AgentStep` records to. + all_tool_calls: List to append tool-call dicts to. + full_tool_results: Optional map of tool_call_id to the full + (pre-truncation) tool result string. + tool_timings: Optional list of ``{"name", "timestamp", + "duration_ms"}`` records captured by the tool wrappers. + Consumed in FIFO order per tool name to populate + ``duration_ms`` and a real execution timestamp on + ``tool_result`` steps. + """ now = time.time() + # Map tool_call_id -> tool name from assistant tool_calls messages so + # tool_result steps record the real tool name, not the call id. + id_to_name: dict[str, str] = {} + for msg in messages: + if msg.get("role") == "assistant": + for tc in msg.get("tool_calls") or []: + tc_fn = tc.get("function", {}) + tc_name = tc_fn.get("name", tc.get("name", "")) + tc_id = tc.get("id", "") + if tc_id and tc_name: + id_to_name[tc_id] = tc_name + + # FIFO queues of timing records per tool name + timings_by_name: dict[str, list[dict[str, Any]]] = {} + for rec in tool_timings or []: + timings_by_name.setdefault(rec.get("name", ""), []).append(rec) + for msg in messages: role = msg.get("role", "") # Extract usage from message meta if present @@ -981,13 +1139,23 @@ def _extract_steps( full_result = None if full_tool_results and tool_call_id: full_result = full_tool_results.get(tool_call_id) + # Resolve the real tool name from the originating assistant + # tool_calls message (fall back to the raw id). + tool_name = id_to_name.get(tool_call_id, tool_call_id) + # Pop the matching timing record (FIFO) if available. + timing = None + if tool_name is not None: + queue = timings_by_name.get(tool_name) + if queue: + timing = queue.pop(0) steps.append( AgentStep( step_type=StepType.tool_result, - timestamp=now, + timestamp=timing["timestamp"] if timing else now, content=msg.get("content", ""), - tool_name=tool_call_id, + tool_name=tool_name, tool_result=full_result, + duration_ms=timing["duration_ms"] if timing else 0.0, ) ) @@ -1081,9 +1249,18 @@ def iter(self, prompt: str, *, deps: Any = None) -> AgentIterator: return AgentIterator(gen) def _execute_iter(self, prompt: str, deps: Any) -> Generator[AgentStep, None, AgentResult]: - """Generator that executes the agent loop and yields each step.""" + """Generator that executes the agent loop and yields each step. + + Raises: + RecursionError: If the agent nesting depth exceeds ``max_depth``. + """ + current_depth = _agent_depth.get() + if current_depth >= self._max_depth: + raise RecursionError(f"Agent recursion depth exceeded: {current_depth} >= {self._max_depth}") + token = _agent_depth.set(current_depth + 1) self._lifecycle = AgentState.running self._stop_requested = False + self._active_conversation = None steps: list[AgentStep] = [] try: @@ -1095,6 +1272,8 @@ def _execute_iter(self, prompt: str, deps: Any) -> Generator[AgentStep, None, Ag except Exception: self._lifecycle = AgentState.errored raise + finally: + _agent_depth.reset(token) # ------------------------------------------------------------------ # run_stream() — streaming output @@ -1108,17 +1287,29 @@ def run_stream(self, prompt: str, *, deps: Any = None) -> StreamedAgentResult: final :class:`AgentResult` is available via :attr:`StreamedAgentResult.result`. - When tools are registered, streaming falls back to non-streaming - ``conv.ask()`` and yields the full response as a single - ``text_delta`` event. + When tools are registered, the tool loop runs via + ``conv.ask_with_tool_events()``: ``tool_call`` and ``tool_result`` + events are emitted as tools execute, and the final LLM response is + yielded as a single ``text_delta`` event (per-turn text is not + token-streamed in this mode). Without tools, the driver's native + streaming is used when available. """ gen = self._execute_stream(prompt, deps) return StreamedAgentResult(gen) def _execute_stream(self, prompt: str, deps: Any) -> Generator[StreamEvent, None, AgentResult]: - """Generator that executes the agent loop and yields stream events.""" + """Generator that executes the agent loop and yields stream events. + + Raises: + RecursionError: If the agent nesting depth exceeds ``max_depth``. + """ + current_depth = _agent_depth.get() + if current_depth >= self._max_depth: + raise RecursionError(f"Agent recursion depth exceeded: {current_depth} >= {self._max_depth}") + token = _agent_depth.set(current_depth + 1) self._lifecycle = AgentState.running self._stop_requested = False + self._active_conversation = None steps: list[AgentStep] = [] try: @@ -1139,7 +1330,14 @@ def _execute_stream(self, prompt: str, deps: Any) -> Generator[StreamEvent, None resolved_system_prompt = self._resolve_system_prompt(ctx) # 5. Wrap tools with context (pass session for child agent usage aggregation) - wrapped_tools = self._wrap_tools_with_context(ctx, session) + run_state: dict[str, Any] = {} + tool_timings: list[dict[str, Any]] = [] + wrapped_tools = self._wrap_tools_with_context( + ctx, + session, + ctx_fn=self._make_live_ctx_fn(prompt, deps, session, run_state), + tool_timings=tool_timings, + ) has_tools = bool(wrapped_tools) # 6. Build Conversation @@ -1148,6 +1346,7 @@ def _execute_stream(self, prompt: str, deps: Any) -> Generator[StreamEvent, None tools=wrapped_tools if wrapped_tools else None, driver_callbacks=driver_callbacks, ) + run_state["conv"] = conv # 7. Fire on_iteration callback if self._agent_callbacks.on_iteration: @@ -1207,6 +1406,7 @@ def _execute_stream(self, prompt: str, deps: Any) -> Generator[StreamEvent, None steps, all_tool_calls, getattr(conv, "_full_tool_results", None), + tool_timings, ) # 9. Parse output @@ -1257,6 +1457,8 @@ def _execute_stream(self, prompt: str, deps: Any) -> Generator[StreamEvent, None except Exception: self._lifecycle = AgentState.errored raise + finally: + _agent_depth.reset(token) # ------------------------------------------------------------------ # run_live() — interleaved tool calling with streaming text deltas @@ -1292,6 +1494,7 @@ def _execute_live(self, prompt: str, deps: Any) -> Generator[Any, None, AgentRes token = _agent_depth.set(current_depth + 1) self._lifecycle = AgentState.running self._stop_requested = False + self._active_conversation = None steps: list[AgentStep] = [] try: @@ -1303,13 +1506,21 @@ def _execute_live(self, prompt: str, deps: Any) -> Generator[Any, None, AgentRes ctx = self._build_run_context(prompt, deps, session, [], 0) effective_prompt = self._run_input_guardrails(ctx, prompt) resolved_system_prompt = self._resolve_system_prompt(ctx) - wrapped_tools = self._wrap_tools_with_context(ctx, session) + run_state: dict[str, Any] = {} + tool_timings: list[dict[str, Any]] = [] + wrapped_tools = self._wrap_tools_with_context( + ctx, + session, + ctx_fn=self._make_live_ctx_fn(prompt, deps, session, run_state), + tool_timings=tool_timings, + ) conv = self._build_conversation( system_prompt=resolved_system_prompt, tools=wrapped_tools if wrapped_tools else None, driver_callbacks=driver_callbacks, ) + run_state["conv"] = conv if self._agent_callbacks.on_iteration: self._agent_callbacks.on_iteration(0) @@ -1346,6 +1557,7 @@ def _execute_live(self, prompt: str, deps: Any) -> Generator[Any, None, AgentRes steps, all_tool_calls, getattr(conv, "_full_tool_results", None), + tool_timings, ) if self._output_type is not None: diff --git a/prompture/agents/async_agent.py b/prompture/agents/async_agent.py index ed4e78fe..2cda9e95 100644 --- a/prompture/agents/async_agent.py +++ b/prompture/agents/async_agent.py @@ -15,6 +15,8 @@ from __future__ import annotations import asyncio +import concurrent.futures +import contextvars import inspect import json import logging @@ -26,7 +28,7 @@ from pydantic import BaseModel from ..extraction.tools import clean_json_text -from ..infra.budget import BudgetPolicy, BudgetState, enforce_budget, resolve_budget_policy +from ..infra.budget import BudgetPolicy, resolve_budget_policy from ..infra.callbacks import DriverCallbacks from ..infra.provider_env import ProviderEnvironment from ..infra.session import UsageSession @@ -70,6 +72,32 @@ def _is_async_callable(fn: Callable[..., Any]) -> bool: return dunder_call is not None and asyncio.iscoroutinefunction(dunder_call) +def _run_awaitable_sync(awaitable: Any) -> Any: + """Drive *awaitable* to completion from sync code. + + Used only as a fallback when a sync caller (``ToolRegistry.execute``) + hits an async tool/callback. When a loop is already running in the + current thread the awaitable runs on a worker thread with the caller's + ``contextvars`` copied over, so values such as tukuy's + ``SecurityContext`` and ``current_tool_call_id`` survive the hop. + The preferred path is the ``_async_fn`` hook awaited by + :meth:`ToolRegistry.aexecute`, which needs no thread bridge at all. + """ + + async def _await_it() -> Any: + return await awaitable + + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + if loop is not None and loop.is_running(): + ctx_snapshot = contextvars.copy_context() + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + return pool.submit(ctx_snapshot.run, lambda: asyncio.run(_await_it())).result() + return asyncio.run(_await_it()) + + def _tool_wants_context(fn: Callable[..., Any]) -> bool: """Check whether *fn*'s first parameter is annotated as :class:`RunContext`.""" sig = inspect.signature(fn) @@ -160,6 +188,11 @@ class AsyncAgent(Generic[DepsType]): skill_config: Optional configuration dict injected as a :class:`SkillContext` into tukuy skill ``invoke()`` calls. When ``None`` (default) no config is injected. + tool_timeout: Per-tool wall-clock budget in seconds. A tool that + overruns it yields an error result the model can react to + instead of wedging the run forever. ``None`` (default) means + no timeout. Can be overridden per call via + ``options={"tool_timeout": ...}``. """ def __init__( @@ -189,6 +222,7 @@ def __init__( auto_approve_safe_only: bool = False, skill_config: dict[str, Any] | None = None, max_tool_result_length: int | None = None, + tool_timeout: float | None = None, max_depth: int = _DEFAULT_MAX_AGENT_DEPTH, env: ProviderEnvironment | None = None, ) -> None: @@ -220,7 +254,12 @@ def __init__( self._auto_approve_safe_only = auto_approve_safe_only self._skill_config = skill_config self._max_tool_result_length = max_tool_result_length + self._tool_timeout = tool_timeout self._conversation: Any = None + # The conversation driving the current run. Tracked separately from + # ``_conversation`` (which only holds persistent ones) so that + # ``stop()`` can reach the loop of a non-persistent run too. + self._active_conversation: Any = None # Build internal tool registry self._tools = ToolRegistry() @@ -256,8 +295,21 @@ def state(self) -> AgentState: return self._lifecycle def stop(self) -> None: - """Request graceful shutdown after the current iteration.""" + """Request graceful shutdown after the current iteration. + + Sets the agent-level flag and forwards the request to the conversation + driving the current run (cooperative stop) so its tool-round loop + exits gracefully between rounds instead of starting another round. + Works for non-persistent agents too: the in-flight conversation is + tracked in ``_active_conversation`` for the duration of the run, so + the flag is never merely decorative. + + Safe to call from a tool, a callback, or another thread. + """ self._stop_requested = True + conv = self._active_conversation or self._conversation + if conv is not None: + conv.request_stop() @property def callbacks(self) -> AgentCallbacks: @@ -317,8 +369,6 @@ def _call_agent(prompt: str) -> str: loop = None if loop is not None and loop.is_running(): - import concurrent.futures - with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: result = pool.submit(asyncio.run, agent.run(prompt)).result() else: @@ -329,8 +379,26 @@ def _call_agent(prompt: str) -> str: return extractor(result) return result.output_text + async def _call_agent_async(prompt: str) -> str: + """Run the wrapped async agent, awaited in the current loop. + + Registered as ``_async_fn`` so :meth:`ToolRegistry.aexecute` + awaits this directly instead of using the thread bridge in + ``_call_agent`` (which drops ``contextvars`` such as tukuy's + ``SecurityContext`` and ``current_tool_call_id``). + """ + result = await agent.run(prompt) + _call_agent._last_agent_result = result # type: ignore[attr-defined] + if extractor is not None: + extracted = extractor(result) + if inspect.isawaitable(extracted): + extracted = await extracted + return extracted + return result.output_text + _call_agent._source_agent = agent # type: ignore[attr-defined] _call_agent._last_agent_result = None # type: ignore[attr-defined] + _call_agent._async_fn = _call_agent_async # type: ignore[attr-defined] return ToolDefinition( name=tool_name, @@ -371,6 +439,7 @@ async def run( token = _agent_depth.set(current_depth + 1) self._lifecycle = AgentState.running self._stop_requested = False + self._active_conversation = None steps: list[AgentStep] = [] try: @@ -444,17 +513,67 @@ def _build_run_context( prompt=prompt, ) + def _make_live_ctx_fn( + self, + prompt: str, + deps: Any, + session: UsageSession, + run_state: dict[str, Any], + ) -> Callable[[], RunContext[Any]]: + """Return a factory that builds a fresh :class:`RunContext` per call. + + The factory reads the current conversation from *run_state* (key + ``"conv"``) so tools invoked in later tool rounds see live + ``iteration``/``messages``/``usage`` instead of the snapshot taken + at the start of the run. ``run_state["conv"]`` must be assigned + once the conversation is built; before that, an empty context is + produced. + """ + + def _live_ctx() -> RunContext[Any]: + conv = run_state.get("conv") + messages = conv.messages if conv is not None else [] + iteration = max( + 0, + sum(1 for m in messages if m.get("role") == "assistant" and m.get("tool_calls")) - 1, + ) + return self._build_run_context(prompt, deps, session, messages, iteration) + + return _live_ctx + # ------------------------------------------------------------------ # Tool wrapping (RunContext injection + ModelRetry + callbacks) # ------------------------------------------------------------------ - def _wrap_tools_with_context(self, ctx: RunContext[Any], session: UsageSession | None = None) -> ToolRegistry: + def _wrap_tools_with_context( + self, + ctx: RunContext[Any], + session: UsageSession | None = None, + ctx_fn: Callable[[], RunContext[Any]] | None = None, + tool_timings: list[dict[str, Any]] | None = None, + ) -> ToolRegistry: """Return a new :class:`ToolRegistry` with wrapped tool functions. - All wrappers are **sync** so they work with ``ToolRegistry.execute()``. - For async tool functions, the wrapper uses - ``asyncio.get_event_loop().run_until_complete()`` as a fallback. - If *session* is provided, child agent usage is aggregated into it. + For each registered tool: + - If the tool's first param is ``RunContext``, inject the live context + automatically. When *ctx_fn* is provided it is called at each tool + invocation so tools see the current iteration/messages; otherwise the + static *ctx* snapshot is used. + - Catch :class:`ModelRetry` and convert to an error string. + - Fire ``agent_callbacks.on_tool_start`` / ``on_tool_end``. + - Strip the ``RunContext`` parameter from the JSON schema sent to the LLM. + - If the tool wraps a child agent (via ``as_tool``), aggregate its + usage into the parent *session*. + - When *tool_timings* is provided, append a + ``{"name", "timestamp", "duration_ms"}`` record per invocation. + + Async tool functions additionally get an ``async`` wrapper attached + as ``_async_fn`` on the sync wrapper. :meth:`ToolRegistry.aexecute` + prefers and awaits that hook, so inside ``AsyncConversation`` the + coroutine runs on the current event loop — no thread/asyncio.run + bridge, and ``contextvars`` (tukuy ``SecurityContext``, + ``current_tool_call_id``) propagate naturally. The sync wrapper + remains as a fallback for plain :meth:`ToolRegistry.execute` calls. """ if not self._tools: return ToolRegistry() @@ -475,100 +594,101 @@ def _make_wrapper( _is_async: bool, _cb: AgentCallbacks = cb, _session: UsageSession | None = session, + _ctx_fn: Callable[[], RunContext[Any]] | None = ctx_fn, + _timings: list[dict[str, Any]] | None = tool_timings, ) -> Callable[..., Any]: def _invoke_cb_sync(callback: Callable[..., Any], *cb_args: Any) -> Any: """Invoke a possibly-async callback from a sync tool wrapper.""" - if _is_async_callable(callback): - try: - loop = asyncio.get_running_loop() - except RuntimeError: - loop = None - if loop is not None and loop.is_running(): - import concurrent.futures - - with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: - return pool.submit(asyncio.run, callback(*cb_args)).result() - else: - return asyncio.run(callback(*cb_args)) - return callback(*cb_args) + result = callback(*cb_args) + if inspect.isawaitable(result): + return _run_awaitable_sync(result) + return result + + def _prepare_call(kwargs: dict[str, Any]) -> dict[str, Any]: + """Apply skill-config injection and the auto-approve gate.""" + # Inject skill config via SkillContext for tukuy skills + if self._skill_config is not None: + _skill_obj = getattr(_fn, "__skill__", None) + if _skill_obj is not None: + from tukuy import SkillContext + + kwargs["context"] = SkillContext(config=self._skill_config) + + # Auto-approve gate: block tukuy skills with side effects + if self._auto_approve_safe_only: + skill_obj = getattr(_fn, "__skill__", None) + if skill_obj is not None: + desc = skill_obj.descriptor + has_side_effects = getattr(desc, "side_effects", False) + has_network = getattr(desc, "requires_network", False) + if has_side_effects or has_network: + raise ApprovalRequired( + tool_name=_name, + action="execute tool with side effects", + details={ + "side_effects": has_side_effects, + "requires_network": has_network, + "skill_name": desc.name, + }, + ) + return kwargs + + def _call_args(call_ctx: RunContext[Any]) -> tuple[Any, ...]: + return (call_ctx,) if _wants else () + + def _aggregate_child_usage() -> None: + """Aggregate child agent usage to the parent session.""" + if _session is not None and hasattr(_fn, "_source_agent"): + agent_result = getattr(_fn, "_last_agent_result", None) + if agent_result is not None and hasattr(agent_result, "run_usage"): + child_usage = agent_result.run_usage + child_name = getattr(_fn._source_agent, "name", "") or _name + _session.record( + { + "meta": { + "prompt_tokens": child_usage.get("prompt_tokens", 0), + "completion_tokens": child_usage.get("completion_tokens", 0), + "total_tokens": child_usage.get("total_tokens", 0), + "cost": child_usage.get("cost", 0.0), + }, + "driver": f"sub-agent:{child_name}", + } + ) + + def _record_timing(start: float) -> None: + if _timings is not None: + _timings.append( + { + "name": _name, + "timestamp": time.time(), + "duration_ms": (time.perf_counter() - start) * 1000, + } + ) def wrapper(**kwargs: Any) -> Any: + call_ctx = _ctx_fn() if _ctx_fn is not None else ctx + start = time.perf_counter() if _cb.on_tool_start: _invoke_cb_sync(_cb.on_tool_start, _name, kwargs) try: - # Inject skill config via SkillContext for tukuy skills - if self._skill_config is not None: - _skill_obj = getattr(_fn, "__skill__", None) - if _skill_obj is not None: - from tukuy import SkillContext - - kwargs["context"] = SkillContext(config=self._skill_config) - - # Auto-approve gate: block tukuy skills with side effects - if self._auto_approve_safe_only: - skill_obj = getattr(_fn, "__skill__", None) - if skill_obj is not None: - desc = skill_obj.descriptor - has_side_effects = getattr(desc, "side_effects", False) - has_network = getattr(desc, "requires_network", False) - if has_side_effects or has_network: - raise ApprovalRequired( - tool_name=_name, - action="execute tool with side effects", - details={ - "side_effects": has_side_effects, - "requires_network": has_network, - "skill_name": desc.name, - }, - ) - - if _wants: - call_args: tuple[Any, ...] = (ctx,) - else: - call_args = () - + _prepare_call(kwargs) if _is_async: - coro = _fn(*call_args, **kwargs) - # Try to get running loop; if none, use asyncio.run() - try: - loop = asyncio.get_running_loop() - except RuntimeError: - loop = None - if loop is not None and loop.is_running(): - # We're inside an async context — create a new thread - import concurrent.futures - - with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: - result = pool.submit(asyncio.run, coro).result() - else: - result = asyncio.run(coro) + result = _run_awaitable_sync(_fn(*_call_args(call_ctx), **kwargs)) else: - result = _fn(*call_args, **kwargs) + result = _fn(*_call_args(call_ctx), **kwargs) except ApprovalRequired as exc: # Handle approval request if _cb.on_approval_needed: approved = _invoke_cb_sync(_cb.on_approval_needed, exc.tool_name, exc.action, exc.details) if approved: + # Retry the tool call after approval try: if _is_async: - coro = _fn(*call_args, **kwargs) if not _wants else _fn(ctx, **kwargs) - try: - loop = asyncio.get_running_loop() - except RuntimeError: - loop = None - if loop is not None and loop.is_running(): - import concurrent.futures - - with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: - result = pool.submit(asyncio.run, coro).result() - else: - result = asyncio.run(coro) + result = _run_awaitable_sync(_fn(*_call_args(call_ctx), **kwargs)) else: - if _wants: - result = _fn(ctx, **kwargs) - else: - result = _fn(**kwargs) + result = _fn(*_call_args(call_ctx), **kwargs) except ApprovalRequired: + # Tool raised ApprovalRequired again - don't loop result = f"Error: Tool '{_name}' requires approval but approval was already granted" except ModelRetry as retry_exc: result = f"Error: {retry_exc.message}" @@ -579,28 +699,53 @@ def wrapper(**kwargs: Any) -> Any: except ModelRetry as exc: result = f"Error: {exc.message}" - # Aggregate child agent usage to parent session - if _session is not None and hasattr(_fn, "_source_agent"): - agent_result = getattr(_fn, "_last_agent_result", None) - if agent_result is not None and hasattr(agent_result, "run_usage"): - child_usage = agent_result.run_usage - child_name = getattr(_fn._source_agent, "name", "") or _name - _session.record( - { - "meta": { - "prompt_tokens": child_usage.get("prompt_tokens", 0), - "completion_tokens": child_usage.get("completion_tokens", 0), - "total_tokens": child_usage.get("total_tokens", 0), - "cost": child_usage.get("cost", 0.0), - }, - "driver": f"sub-agent:{child_name}", - } - ) + _aggregate_child_usage() + _record_timing(start) if _cb.on_tool_end: _invoke_cb_sync(_cb.on_tool_end, _name, result) return result + async def async_wrapper(**kwargs: Any) -> Any: + """Fully-async wrapper: awaited directly via ``_async_fn``.""" + call_ctx = _ctx_fn() if _ctx_fn is not None else ctx + start = time.perf_counter() + if _cb.on_tool_start: + await AsyncAgent._invoke_callback(_cb.on_tool_start, _name, kwargs) + try: + _prepare_call(kwargs) + result = await _fn(*_call_args(call_ctx), **kwargs) + except ApprovalRequired as exc: + # Handle approval request + if _cb.on_approval_needed: + approved = await AsyncAgent._invoke_callback( + _cb.on_approval_needed, exc.tool_name, exc.action, exc.details + ) + if approved: + # Retry the tool call after approval + try: + result = await _fn(*_call_args(call_ctx), **kwargs) + except ApprovalRequired: + # Tool raised ApprovalRequired again - don't loop + result = f"Error: Tool '{_name}' requires approval but approval was already granted" + except ModelRetry as retry_exc: + result = f"Error: {retry_exc.message}" + else: + result = f"Error: Tool '{_name}' execution denied - approval required: {exc.action}" + else: + result = f"Error: Tool '{_name}' requires approval but no approval handler is configured" + except ModelRetry as exc: + result = f"Error: {exc.message}" + + _aggregate_child_usage() + _record_timing(start) + + if _cb.on_tool_end: + await AsyncAgent._invoke_callback(_cb.on_tool_end, _name, result) + return result + + if _is_async: + wrapper._async_fn = async_wrapper # type: ignore[attr-defined] return wrapper wrapped = _make_wrapper(original_fn, wants_ctx, tool_name, is_async) @@ -758,14 +903,28 @@ def _build_conversation( # Reuse existing conversation in persistent mode if self._persistent_conversation and self._conversation is not None: conv = self._conversation - if tools is not None: - conv._tools = tools - if driver_callbacks is not None: - conv._driver.callbacks = driver_callbacks - hook = getattr(self, "_before_turn_hook", None) - if hook is not None: - conv._before_turn = hook - return conv + # If the agent's driver was swapped after the conversation was + # built (e.g. a budget fallback forced a rebuild), the cached + # conversation still holds the old driver — discard it and fall + # through to build a fresh one. + driver_stale = self._driver is not None and getattr(conv, "_driver", None) is not self._driver + if not driver_stale: + # Respect the newly resolved system prompt for this run. + if system_prompt is not None: + conv.system_prompt = system_prompt + # NOTE: assigning ``callbacks`` mutates the driver instance, + # which may be shared (e.g. passed to several agents). A + # persistent-conversation agent should treat its driver as + # exclusively owned by the agent. + if tools is not None: + conv._tools = tools + if driver_callbacks is not None: + conv._driver.callbacks = driver_callbacks + hook = getattr(self, "_before_turn_hook", None) + if hook is not None: + conv._before_turn = hook + return self._track_active(conv) + self._conversation = None effective_tools = tools if tools is not None else (self._tools if self._tools else None) @@ -779,6 +938,8 @@ def _build_conversation( kwargs["before_turn"] = hook if self._max_tool_result_length is not None: kwargs["max_tool_result_length"] = self._max_tool_result_length + if self._tool_timeout is not None: + kwargs["tool_timeout"] = self._tool_timeout if self._options: kwargs["options"] = self._options if driver_callbacks is not None: @@ -806,6 +967,19 @@ def _build_conversation( conv = AsyncConversation(**kwargs) if self._persistent_conversation: self._conversation = conv + return self._track_active(conv) + + def _track_active(self, conv: Any) -> Any: + """Record *conv* as the conversation driving the current run. + + Lets :meth:`stop` reach the tool-round loop of a non-persistent run. + If ``stop()`` was already called between the run starting and the + conversation being built, the request is replayed onto *conv* so it + is not lost to that race. + """ + self._active_conversation = conv + if self._stop_requested: + conv.request_stop() return conv async def _execute( @@ -837,8 +1011,16 @@ async def _execute( # 4. Resolve system prompt resolved_system_prompt = self._resolve_system_prompt(ctx) - # 5. Wrap tools with context (pass session for child agent usage aggregation) - wrapped_tools = self._wrap_tools_with_context(ctx, session) + # 5. Wrap tools with context (pass session for child agent usage + # aggregation; ctx_fn refreshes RunContext per tool round) + run_state: dict[str, Any] = {} + tool_timings: list[dict[str, Any]] = [] + wrapped_tools = self._wrap_tools_with_context( + ctx, + session, + ctx_fn=self._make_live_ctx_fn(prompt, deps, session, run_state), + tool_timings=tool_timings, + ) # 6. Build AsyncConversation conv = self._build_conversation( @@ -846,36 +1028,16 @@ async def _execute( tools=wrapped_tools if wrapped_tools else None, driver_callbacks=driver_callbacks, ) + run_state["conv"] = conv # 7. Fire on_iteration callback if self._agent_callbacks.on_iteration: await self._invoke_callback(self._agent_callbacks.on_iteration, 0) - # 8. Enforce budget policy at the agent level (pre-call) - if self._budget_policy is not None: - state = BudgetState( - cost_used=session.cost, - tokens_used=session.total_tokens, - max_cost=self._max_cost, - max_tokens=self._max_tokens, - ) - new_model = enforce_budget( - state, - self._budget_policy, - fallback_models=self._fallback_models, - current_model=self._model, - on_model_fallback=self._on_model_fallback, - ) - if new_model is not None: - self._model = new_model - self._driver = None # force rebuild - conv = self._build_conversation( - system_prompt=resolved_system_prompt, - tools=wrapped_tools if wrapped_tools else None, - driver_callbacks=driver_callbacks, - ) - - # 9. Ask the conversation (handles full tool loop internally) + # 8. Ask the conversation (handles full tool loop internally) + # Note: budget policy is enforced inside the conversation; the + # agent-level pre-call check was removed because it read a fresh, + # always-empty UsageSession. agent_name = self.name or self.__class__.__name__ with tracker.agent(agent_name): t0 = time.perf_counter() @@ -896,7 +1058,7 @@ async def _execute( # 9. Extract steps and tool calls all_tool_calls: list[dict[str, Any]] = [] full_results = getattr(conv, "_full_tool_results", None) - self._extract_steps(conv.messages, steps, all_tool_calls, full_results) + self._extract_steps(conv.messages, steps, all_tool_calls, full_results, tool_timings) # Handle output_type parsing if self._output_type is not None: @@ -944,10 +1106,41 @@ def _extract_steps( steps: list[AgentStep], all_tool_calls: list[dict[str, Any]], full_tool_results: dict[str, str] | None = None, + tool_timings: list[dict[str, Any]] | None = None, ) -> None: - """Scan conversation messages and populate steps and tool_calls.""" + """Scan conversation messages and populate steps and tool_calls. + + Args: + messages: Conversation messages to scan. + steps: List to append :class:`AgentStep` records to. + all_tool_calls: List to append tool-call dicts to. + full_tool_results: Optional map of tool_call_id to the full + (pre-truncation) tool result string. + tool_timings: Optional list of ``{"name", "timestamp", + "duration_ms"}`` records captured by the tool wrappers. + Consumed in FIFO order per tool name to populate + ``duration_ms`` and a real execution timestamp on + ``tool_result`` steps. + """ now = time.time() + # Map tool_call_id -> tool name from assistant tool_calls messages so + # tool_result steps record the real tool name, not the call id. + id_to_name: dict[str, str] = {} + for msg in messages: + if msg.get("role") == "assistant": + for tc in msg.get("tool_calls") or []: + tc_fn = tc.get("function", {}) + tc_name = tc_fn.get("name", tc.get("name", "")) + tc_id = tc.get("id", "") + if tc_id and tc_name: + id_to_name[tc_id] = tc_name + + # FIFO queues of timing records per tool name + timings_by_name: dict[str, list[dict[str, Any]]] = {} + for rec in tool_timings or []: + timings_by_name.setdefault(rec.get("name", ""), []).append(rec) + for msg in messages: role = msg.get("role", "") # Extract usage from message meta if present @@ -1018,13 +1211,23 @@ def _extract_steps( full_result = None if full_tool_results and tool_call_id: full_result = full_tool_results.get(tool_call_id) + # Resolve the real tool name from the originating assistant + # tool_calls message (fall back to the raw id). + tool_name = id_to_name.get(tool_call_id, tool_call_id) + # Pop the matching timing record (FIFO) if available. + timing = None + if tool_name is not None: + queue = timings_by_name.get(tool_name) + if queue: + timing = queue.pop(0) steps.append( AgentStep( step_type=StepType.tool_result, - timestamp=now, + timestamp=timing["timestamp"] if timing else now, content=msg.get("content", ""), - tool_name=tool_call_id, + tool_name=tool_name, tool_result=full_result, + duration_ms=timing["duration_ms"] if timing else 0.0, ) ) @@ -1104,9 +1307,18 @@ async def _execute_iter( *, images: list[Any] | None = None, ) -> AsyncGenerator[AgentStep]: - """Async generator that executes the agent loop and yields each step.""" + """Async generator that executes the agent loop and yields each step. + + Raises: + RecursionError: If the agent nesting depth exceeds ``max_depth``. + """ + current_depth = _agent_depth.get() + if current_depth >= self._max_depth: + raise RecursionError(f"Agent recursion depth exceeded: {current_depth} >= {self._max_depth}") + token = _agent_depth.set(current_depth + 1) self._lifecycle = AgentState.running self._stop_requested = False + self._active_conversation = None steps: list[AgentStep] = [] try: @@ -1119,6 +1331,8 @@ async def _execute_iter( except Exception: self._lifecycle = AgentState.errored raise + finally: + _agent_depth.reset(token) # ------------------------------------------------------------------ # run_stream() — async streaming @@ -1142,6 +1356,7 @@ async def _execute_stream( token = _agent_depth.set(current_depth + 1) self._lifecycle = AgentState.running self._stop_requested = False + self._active_conversation = None steps: list[AgentStep] = [] try: @@ -1154,7 +1369,14 @@ async def _execute_stream( ctx = self._build_run_context(prompt, deps, session, [], 0) effective_prompt = self._run_input_guardrails(ctx, prompt) resolved_system_prompt = self._resolve_system_prompt(ctx) - wrapped_tools = self._wrap_tools_with_context(ctx, session) + run_state: dict[str, Any] = {} + tool_timings: list[dict[str, Any]] = [] + wrapped_tools = self._wrap_tools_with_context( + ctx, + session, + ctx_fn=self._make_live_ctx_fn(prompt, deps, session, run_state), + tool_timings=tool_timings, + ) has_tools = bool(wrapped_tools) conv = self._build_conversation( @@ -1162,6 +1384,7 @@ async def _execute_stream( tools=wrapped_tools if wrapped_tools else None, driver_callbacks=driver_callbacks, ) + run_state["conv"] = conv if self._agent_callbacks.on_iteration: await self._invoke_callback(self._agent_callbacks.on_iteration, 0) @@ -1199,7 +1422,7 @@ async def _execute_stream( # Extract steps all_tool_calls: list[dict[str, Any]] = [] full_results = getattr(conv, "_full_tool_results", None) - self._extract_steps(conv.messages, steps, all_tool_calls, full_results) + self._extract_steps(conv.messages, steps, all_tool_calls, full_results, tool_timings) # Parse output if self._output_type is not None: @@ -1280,6 +1503,7 @@ async def _execute_live( token = _agent_depth.set(current_depth + 1) self._lifecycle = AgentState.running self._stop_requested = False + self._active_conversation = None steps: list[AgentStep] = [] try: @@ -1291,13 +1515,21 @@ async def _execute_live( ctx = self._build_run_context(prompt, deps, session, [], 0) effective_prompt = self._run_input_guardrails(ctx, prompt) resolved_system_prompt = self._resolve_system_prompt(ctx) - wrapped_tools = self._wrap_tools_with_context(ctx, session) + run_state: dict[str, Any] = {} + tool_timings: list[dict[str, Any]] = [] + wrapped_tools = self._wrap_tools_with_context( + ctx, + session, + ctx_fn=self._make_live_ctx_fn(prompt, deps, session, run_state), + tool_timings=tool_timings, + ) conv = self._build_conversation( system_prompt=resolved_system_prompt, tools=wrapped_tools if wrapped_tools else None, driver_callbacks=driver_callbacks, ) + run_state["conv"] = conv if self._agent_callbacks.on_iteration: await self._invoke_callback(self._agent_callbacks.on_iteration, 0) @@ -1329,7 +1561,7 @@ async def _execute_live( response_text = "".join(response_text_parts) all_tool_calls: list[dict[str, Any]] = [] full_results = getattr(conv, "_full_tool_results", None) - self._extract_steps(conv.messages, steps, all_tool_calls, full_results) + self._extract_steps(conv.messages, steps, all_tool_calls, full_results, tool_timings) if self._output_type is not None: output, output_text = await self._parse_output(conv, response_text, steps, all_tool_calls, 0.0, session) diff --git a/prompture/agents/async_conversation.py b/prompture/agents/async_conversation.py index 0d61e408..f675c005 100644 --- a/prompture/agents/async_conversation.py +++ b/prompture/agents/async_conversation.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import json import logging import time @@ -59,8 +60,10 @@ def __init__( callbacks: DriverCallbacks | None = None, tools: ToolRegistry | None = None, max_tool_rounds: int = 10, - max_tool_result_length: int | None = None, + max_tool_result_length: int | None = 16000, simulated_tools: bool | Literal["auto"] = "auto", + tool_timeout: float | None = None, + sequential_tools: bool = False, conversation_id: str | None = None, auto_save: str | Path | None = None, tags: list[str] | None = None, @@ -126,6 +129,12 @@ def __init__( self._max_tool_result_length = max_tool_result_length self._simulated_tools = simulated_tools self._max_history_messages = max_history_messages + self._tool_timeout = tool_timeout + self._sequential_tools = sequential_tools + + # Cooperative stop (C2) and graceful max-rounds (C4) state. + self._stop_requested = False + self._max_rounds_reached = False # Budget enforcement self._max_cost = max_cost @@ -191,6 +200,21 @@ def usage(self) -> dict[str, Any]: def clear(self) -> None: """Reset message history (keeps system_prompt and driver).""" self._messages.clear() + self._full_tool_results.clear() + + def request_stop(self) -> None: + """Request a cooperative stop of the tool loop. + + The loop checks the flag between rounds; when set, it finishes + with one final no-tools answer instead of executing more tools. + """ + self._stop_requested = True + + @property + def max_rounds_reached(self) -> bool: + """``True`` when the last tool loop hit ``max_tool_rounds`` and + completed with a graceful final answer instead of raising.""" + return self._max_rounds_reached def add_context(self, role: str, content: str, images: list[ImageInput] | None = None) -> None: """Seed the history with a user or assistant message.""" @@ -425,6 +449,130 @@ def _truncate_tool_result(self, result_str: str) -> str: result_str[: self._max_tool_result_length] + f"\n\n[... result truncated ({len(result_str):,} chars total)]" ) + def _trim_history(self) -> None: + """Trim history to the sliding window without orphaning tool pairs. + + A naive tail-slice can cut between an assistant ``tool_calls`` + message and its ``tool`` results, which providers reject with a + 400. After slicing, drop any leading ``tool`` messages whose + matching assistant message was cut. + """ + if self._max_history_messages is None or len(self._messages) <= self._max_history_messages: + return + self._messages = self._messages[-self._max_history_messages :] + while self._messages and self._messages[0].get("role") == "tool": + self._messages.pop(0) + + @staticmethod + def _malformed_arguments_message(tc: dict[str, Any]) -> str | None: + """Return an error tool-result string when the driver flagged the + tool call's arguments as malformed (``arguments_error``, C1) or + truncated (C3); ``None`` when the call is safe to execute.""" + args_error = tc.get("arguments_error") + if args_error: + return ( + f"Error: arguments for tool '{tc.get('name')}' could not be parsed ({args_error}). " + "Retry the tool call with valid JSON arguments." + ) + if tc.get("truncated"): + return ( + f"Error: arguments for tool '{tc.get('name')}' were truncated. " + "Retry the tool call with smaller, valid arguments." + ) + return None + + async def _execute_tool_call(self, tc: dict[str, Any], timeout: float | None) -> tuple[str, bool]: + """Execute one tool call, returning ``(result_str, is_error)``. + + Never executes the tool when its arguments were flagged as + malformed/truncated — the error is fed back as the tool result so + the model can retry with valid arguments. + """ + malformed = self._malformed_arguments_message(tc) + if malformed is not None: + return malformed, True + from ..extraction.tukuy_bridge import current_tool_call_id + + token = current_tool_call_id.set(tc["id"]) + try: + if timeout is not None: + result = await asyncio.wait_for(self._tools.aexecute(tc["name"], tc["arguments"]), timeout) + else: + result = await self._tools.aexecute(tc["name"], tc["arguments"]) + return (json.dumps(result) if not isinstance(result, str) else result), False + except (TimeoutError, asyncio.TimeoutError): + return f"Error: tool '{tc['name']}' timed out after {timeout}s", True + except Exception as exc: + return ( + f"Error ({type(exc).__name__}): {exc}" if str(exc) else f"Error: {type(exc).__name__}: {exc!r}" + ), True + finally: + current_tool_call_id.reset(token) + + async def _run_tool_calls(self, tool_calls: list[dict[str, Any]], merged: dict[str, Any]) -> list[tuple[str, bool]]: + """Execute one turn's tool calls, returning ``(result_str, is_error)`` + pairs in the same order as *tool_calls*. + + Independent calls run concurrently via ``asyncio.gather`` (result + order is preserved regardless of completion order) unless + ``sequential_tools`` is set on the conversation or in *merged* + options. + """ + timeout = merged.get("tool_timeout", self._tool_timeout) + sequential = merged.get("sequential_tools", self._sequential_tools) + if sequential or len(tool_calls) <= 1: + return [await self._execute_tool_call(tc, timeout) for tc in tool_calls] + return list(await asyncio.gather(*(self._execute_tool_call(tc, timeout) for tc in tool_calls))) + + async def _final_answer_without_tools(self, msgs: list[dict[str, Any]], merged: dict[str, Any]) -> str: + """One final driver call with tools removed, asking the model to + answer from the tool results it already has. + + Used for graceful exits: max-rounds exhaustion (C4) and + cooperative stop (C2). + """ + self._check_budget() + final_msgs = [ + *msgs, + { + "role": "user", + "content": ( + "You cannot call any more tools. Answer the user's question now, " + "using only the tool results already gathered above." + ), + }, + ] + resp = await self._driver.generate_messages_with_hooks(final_msgs, merged) + text: str = resp.get("text", "") + self._last_reasoning = resp.get("reasoning_content") + self._accumulate_usage(resp.get("meta", {})) + self._messages.append({"role": "assistant", "content": text}) + return text + + async def _final_answer_simulated(self, augmented_system: str, merged: dict[str, Any]) -> str: + """Graceful-exit final call for the simulated-tools loop (C4/C2).""" + from .simulated_tools import parse_simulated_response + + self._check_budget() + msgs: list[dict[str, Any]] = [{"role": "system", "content": augmented_system}] + msgs.extend(self._messages) + msgs.append( + { + "role": "user", + "content": ( + "You cannot call any more tools. Provide your final answer now " + "(in the final_answer format) based on the tool results already gathered." + ), + } + ) + resp = await self._driver.generate_messages_with_hooks(msgs, merged) + text = resp.get("text", "") + self._accumulate_usage(resp.get("meta", {})) + parsed = parse_simulated_response(text, self._tools) + answer: str = parsed["content"] if parsed["type"] == "final_answer" else text + self._messages.append({"role": "assistant", "content": answer}) + return answer + def _build_messages(self, user_content: str, images: list[ImageInput] | None = None) -> list[dict[str, Any]]: """Build the full messages array for an API call.""" msgs: list[dict[str, Any]] = [] @@ -455,8 +603,7 @@ def _accumulate_usage(self, meta: dict[str, Any]) -> None: self._usage["turns"], ) # Trim history to sliding window (system prompt lives outside _messages) - if self._max_history_messages is not None and len(self._messages) > self._max_history_messages: - self._messages = self._messages[-self._max_history_messages :] + self._trim_history() self._maybe_auto_save() async def ask( @@ -479,6 +626,10 @@ async def ask( return await self._ask_with_simulated_tools(content, options, images=images) elif use_native and self._simulated_tools is not True: # type: ignore[comparison-overlap] return await self._ask_with_tools(content, options, images=images) + logger.warning( + "Tools are registered but the driver does not support tool use and " + "simulated_tools is disabled; tools are being ignored for this call." + ) self._check_budget() merged = {**self._options, **(options or {})} @@ -503,6 +654,8 @@ async def _ask_with_tools( images: list[ImageInput] | None = None, ) -> str: """Async tool-use loop: send -> check tool_calls -> execute -> re-send.""" + self._stop_requested = False + self._max_rounds_reached = False merged = {**self._options, **(options or {})} tool_defs = self._tools.to_openai_format() @@ -511,6 +664,9 @@ async def _ask_with_tools( msgs = self._build_messages_raw() for _round in range(self._max_tool_rounds): + if self._stop_requested: + logger.info("Stop requested; finishing with a final answer (no more tool calls)") + return await self._final_answer_without_tools(msgs, merged) self._check_budget() if await self._run_before_turn(): msgs = self._build_messages_raw() @@ -527,6 +683,12 @@ async def _ask_with_tools( self._messages.append({"role": "assistant", "content": text}) return text + # Ensure every tool call has a usable id so parallel calls + # can't collide in _full_tool_results. + for tc in tool_calls: + if not tc.get("id"): + tc["id"] = f"call_{uuid.uuid4().hex}" + assistant_msg: dict[str, Any] = {"role": "assistant", "content": text} assistant_msg["tool_calls"] = [ { @@ -544,20 +706,8 @@ async def _ask_with_tools( self._messages.append(assistant_msg) msgs.append(assistant_msg) - for tc in tool_calls: - from ..extraction.tukuy_bridge import current_tool_call_id as _ask_tc_id - - _ask_token = _ask_tc_id.set(tc["id"]) - try: - result = await self._tools.aexecute(tc["name"], tc["arguments"]) - result_str = json.dumps(result) if not isinstance(result, str) else result - except Exception as exc: - result_str = ( - f"Error ({type(exc).__name__}): {exc}" if str(exc) else f"Error: {type(exc).__name__}: {exc!r}" - ) - finally: - _ask_tc_id.reset(_ask_token) - + results = await self._run_tool_calls(tool_calls, merged) + for tc, (result_str, _is_error) in zip(tool_calls, results, strict=True): # Preserve full result for step extraction before truncating self._full_tool_results[tc["id"]] = result_str @@ -569,7 +719,14 @@ async def _ask_with_tools( self._messages.append(tool_result_msg) msgs.append(tool_result_msg) - raise RuntimeError(f"Tool execution loop exceeded {self._max_tool_rounds} rounds") + # Max rounds exhausted — one graceful final answer without tools + # instead of raising (C4). + logger.warning( + "Tool loop reached max_tool_rounds=%d; requesting a final answer without tools", + self._max_tool_rounds, + ) + self._max_rounds_reached = True + return await self._final_answer_without_tools(msgs, merged) async def ask_with_tool_events( self, @@ -597,6 +754,8 @@ async def ask_with_tool_events( return # Native tool calling with event emission + self._stop_requested = False + self._max_rounds_reached = False merged = {**self._options, **(options or {})} tool_defs = self._tools.to_openai_format() @@ -605,6 +764,12 @@ async def ask_with_tool_events( msgs = self._build_messages_raw() for _round in range(self._max_tool_rounds): + if self._stop_requested: + logger.info("Stop requested; finishing with a final answer (no more tool calls)") + final_text = await self._final_answer_without_tools(msgs, merged) + yield {"type": "text_delta", "text": final_text} + return + self._check_budget() if await self._run_before_turn(): msgs = self._build_messages_raw() resp = await self._driver.generate_messages_with_tools_with_hooks(msgs, tool_defs, merged) @@ -621,6 +786,12 @@ async def ask_with_tool_events( yield {"type": "text_delta", "text": text} return + # Ensure every tool call has a usable id so parallel calls + # can't collide in _full_tool_results. + for tc in tool_calls: + if not tc.get("id"): + tc["id"] = f"call_{uuid.uuid4().hex}" + assistant_msg: dict[str, Any] = {"role": "assistant", "content": text} assistant_msg["tool_calls"] = [ { @@ -643,22 +814,11 @@ async def ask_with_tool_events( "arguments": tc["arguments"], "id": tc["id"], } - # Set the current tool_call_id so downstream code (e.g. the - # tukuy bridge instruction wrapper) can associate streaming - # deltas with the correct tool call. - from ..extraction.tukuy_bridge import current_tool_call_id - - _tc_token = current_tool_call_id.set(tc["id"]) - try: - result = await self._tools.aexecute(tc["name"], tc["arguments"]) - result_str = json.dumps(result) if not isinstance(result, str) else result - except Exception as exc: - result_str = ( - f"Error ({type(exc).__name__}): {exc}" if str(exc) else f"Error: {type(exc).__name__}: {exc!r}" - ) - finally: - current_tool_call_id.reset(_tc_token) + # Independent calls run concurrently; results stay in call order + # so tool_call_id <-> result matching holds. + results = await self._run_tool_calls(tool_calls, merged) + for tc, (result_str, _is_error) in zip(tool_calls, results, strict=True): # Preserve full result for step extraction before truncating self._full_tool_results[tc["id"]] = result_str @@ -679,7 +839,14 @@ async def ask_with_tool_events( self._messages.append(tool_result_msg) msgs.append(tool_result_msg) - raise RuntimeError(f"Tool execution loop exceeded {self._max_tool_rounds} rounds") + # Max rounds exhausted — graceful final answer without tools (C4). + logger.warning( + "Tool loop reached max_tool_rounds=%d; requesting a final answer without tools", + self._max_tool_rounds, + ) + self._max_rounds_reached = True + final_text = await self._final_answer_without_tools(msgs, merged) + yield {"type": "text_delta", "text": final_text} async def ask_live( self, @@ -734,6 +901,8 @@ async def ask_live( yield TurnComplete(usage=dict(self._usage)) return + self._stop_requested = False + self._max_rounds_reached = False merged = {**self._options, **(options or {})} tool_defs = self._tools.to_openai_format() @@ -742,6 +911,12 @@ async def ask_live( msgs = self._build_messages_raw() for round_idx in range(self._max_tool_rounds): + if self._stop_requested: + logger.info("Stop requested; finishing with a final answer (no more tool calls)") + final_text = await self._final_answer_without_tools(msgs, merged) + yield TextDelta(text=final_text) + yield TurnComplete(usage=dict(self._usage)) + return self._check_budget() if await self._run_before_turn(): msgs = self._build_messages_raw() @@ -751,7 +926,7 @@ async def ask_live( assistant_text_parts: list[str] = [] assistant_thinking_parts: list[str] = [] tool_calls_in_turn: list[dict[str, Any]] = [] - pending_tools: list[tuple[str, str, dict[str, Any]]] = [] + pending_tools: list[dict[str, Any]] = [] turn_usage: dict[str, Any] = {} stream = self._driver.generate_messages_with_tools_stream(msgs, tool_defs, merged) @@ -763,10 +938,18 @@ async def ask_live( elif et == "thinking_delta": assistant_thinking_parts.append(event.text) elif et == "tool_use_stop": - pending_tools.append((event.id, event.name, event.input)) + tool_use_id = event.id or f"call_{uuid.uuid4().hex}" + pending_tools.append( + { + "id": tool_use_id, + "name": event.name, + "arguments": event.input, + "truncated": getattr(event, "truncated", False), + } + ) tool_calls_in_turn.append( { - "id": event.id, + "id": tool_use_id, "type": "function", "function": {"name": event.name, "arguments": json.dumps(event.input)}, } @@ -793,35 +976,31 @@ async def ask_live( yield TurnComplete(usage=dict(self._usage)) return - for tool_id, tool_name, tool_input in pending_tools: - from ..extraction.tukuy_bridge import current_tool_call_id - - _tok = current_tool_call_id.set(tool_id) - try: - result = await self._tools.aexecute(tool_name, tool_input) - result_str = json.dumps(result) if not isinstance(result, str) else result - is_error = False - except Exception as exc: - result_str = ( - f"Error ({type(exc).__name__}): {exc}" if str(exc) else f"Error: {type(exc).__name__}: {exc!r}" - ) - is_error = True - finally: - current_tool_call_id.reset(_tok) - - self._full_tool_results[tool_id] = result_str - yield ToolResult(id=tool_id, name=tool_name, output=result_str, is_error=is_error) + # Independent calls run concurrently; results stay in call order + # so tool_call_id <-> result matching holds. + results = await self._run_tool_calls(pending_tools, merged) + for tc, (result_str, is_error) in zip(pending_tools, results, strict=True): + self._full_tool_results[tc["id"]] = result_str + yield ToolResult(id=tc["id"], name=tc["name"], output=result_str, is_error=is_error) - truncated = self._truncate_tool_result(result_str) + truncated_result = self._truncate_tool_result(result_str) tool_result_msg: dict[str, Any] = { "role": "tool", - "tool_call_id": tool_id, - "content": truncated, + "tool_call_id": tc["id"], + "content": truncated_result, } self._messages.append(tool_result_msg) msgs.append(tool_result_msg) - raise RuntimeError(f"ask_live exceeded {self._max_tool_rounds} rounds") + # Max rounds exhausted — graceful final answer without tools (C4). + logger.warning( + "ask_live reached max_tool_rounds=%d; requesting a final answer without tools", + self._max_tool_rounds, + ) + self._max_rounds_reached = True + final_text = await self._final_answer_without_tools(msgs, merged) + yield TextDelta(text=final_text) + yield TurnComplete(usage=dict(self._usage)) async def _ask_with_simulated_tool_events( self, @@ -832,6 +1011,8 @@ async def _ask_with_simulated_tool_events( """Async simulated tool calling with event emission.""" from .simulated_tools import build_tool_prompt, format_tool_result, parse_simulated_response + self._stop_requested = False + self._max_rounds_reached = False merged = {**self._options, **(options or {})} tool_prompt = build_tool_prompt(self._tools) @@ -843,6 +1024,12 @@ async def _ask_with_simulated_tool_events( self._messages.append({"role": "user", "content": user_content}) for _round in range(self._max_tool_rounds): + if self._stop_requested: + logger.info("Stop requested; finishing with a final answer (no more tool calls)") + answer = await self._final_answer_simulated(augmented_system, merged) + yield {"type": "text_delta", "text": answer} + return + self._check_budget() await self._run_before_turn() msgs: list[dict[str, Any]] = [] msgs.append({"role": "system", "content": augmented_system}) @@ -868,12 +1055,18 @@ async def _ask_with_simulated_tool_events( yield {"type": "tool_call", "name": tool_name, "arguments": tool_args, "id": ""} + timeout = merged.get("tool_timeout", self._tool_timeout) from ..extraction.tukuy_bridge import current_tool_call_id as _sim_tc_id _sim_token = _sim_tc_id.set("") try: - result = await self._tools.aexecute(tool_name, tool_args) + if timeout is not None: + result = await asyncio.wait_for(self._tools.aexecute(tool_name, tool_args), timeout) + else: + result = await self._tools.aexecute(tool_name, tool_args) result_msg = format_tool_result(tool_name, result) + except (TimeoutError, asyncio.TimeoutError): + result_msg = format_tool_result(tool_name, f"Error: tool '{tool_name}' timed out after {timeout}s") except Exception as exc: result_msg = format_tool_result(tool_name, f"Error: {exc}") finally: @@ -883,7 +1076,14 @@ async def _ask_with_simulated_tool_events( self._messages.append({"role": "user", "content": self._truncate_tool_result(result_msg)}) - raise RuntimeError(f"Simulated tool execution loop exceeded {self._max_tool_rounds} rounds") + # Max rounds exhausted — graceful final answer without tools (C4). + logger.warning( + "Simulated tool loop reached max_tool_rounds=%d; requesting a final answer without tools", + self._max_tool_rounds, + ) + self._max_rounds_reached = True + answer = await self._final_answer_simulated(augmented_system, merged) + yield {"type": "text_delta", "text": answer} async def _ask_with_simulated_tools( self, @@ -894,6 +1094,8 @@ async def _ask_with_simulated_tools( """Async prompt-based tool calling for drivers without native tool use.""" from .simulated_tools import build_tool_prompt, format_tool_result, parse_simulated_response + self._stop_requested = False + self._max_rounds_reached = False merged = {**self._options, **(options or {})} tool_prompt = build_tool_prompt(self._tools) @@ -907,6 +1109,9 @@ async def _ask_with_simulated_tools( self._messages.append({"role": "user", "content": user_content}) for _round in range(self._max_tool_rounds): + if self._stop_requested: + logger.info("Stop requested; finishing with a final answer (no more tool calls)") + return await self._final_answer_simulated(augmented_system, merged) self._check_budget() await self._run_before_turn() # Build messages with the augmented system prompt @@ -933,16 +1138,28 @@ async def _ask_with_simulated_tools( # Record assistant's tool call as an assistant message self._messages.append({"role": "assistant", "content": text}) + timeout = merged.get("tool_timeout", self._tool_timeout) try: - result = await self._tools.aexecute(tool_name, tool_args) + if timeout is not None: + result = await asyncio.wait_for(self._tools.aexecute(tool_name, tool_args), timeout) + else: + result = await self._tools.aexecute(tool_name, tool_args) result_msg = format_tool_result(tool_name, result) + except (TimeoutError, asyncio.TimeoutError): + result_msg = format_tool_result(tool_name, f"Error: tool '{tool_name}' timed out after {timeout}s") except Exception as exc: result_msg = format_tool_result(tool_name, f"Error: {exc}") # Record tool result as a user message (truncated for the LLM) self._messages.append({"role": "user", "content": self._truncate_tool_result(result_msg)}) - raise RuntimeError(f"Simulated tool execution loop exceeded {self._max_tool_rounds} rounds") + # Max rounds exhausted — graceful final answer without tools (C4). + logger.warning( + "Simulated tool loop reached max_tool_rounds=%d; requesting a final answer without tools", + self._max_tool_rounds, + ) + self._max_rounds_reached = True + return await self._final_answer_simulated(augmented_system, merged) def _build_messages_raw(self) -> list[dict[str, Any]]: """Build messages array from system prompt + full history (including tool messages).""" diff --git a/prompture/agents/conversation.py b/prompture/agents/conversation.py index ffb6bad9..b192754c 100644 --- a/prompture/agents/conversation.py +++ b/prompture/agents/conversation.py @@ -7,6 +7,8 @@ import time import uuid from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import TimeoutError as FuturesTimeoutError from datetime import date, datetime, timezone from decimal import Decimal from pathlib import Path @@ -60,8 +62,9 @@ def __init__( callbacks: DriverCallbacks | None = None, tools: ToolRegistry | None = None, max_tool_rounds: int = 10, - max_tool_result_length: int | None = None, + max_tool_result_length: int | None = 16000, simulated_tools: bool | Literal["auto"] = "auto", + tool_timeout: float | None = None, conversation_id: str | None = None, auto_save: str | Path | None = None, tags: list[str] | None = None, @@ -129,6 +132,11 @@ def __init__( self._max_tool_result_length = max_tool_result_length self._simulated_tools = simulated_tools self._max_history_messages = max_history_messages + self._tool_timeout = tool_timeout + + # Cooperative stop (C2) and graceful max-rounds (C4) state. + self._stop_requested = False + self._max_rounds_reached = False # Budget enforcement self._max_cost = max_cost @@ -194,6 +202,21 @@ def usage(self) -> dict[str, Any]: def clear(self) -> None: """Reset message history (keeps system_prompt and driver).""" self._messages.clear() + self._full_tool_results.clear() + + def request_stop(self) -> None: + """Request a cooperative stop of the tool loop. + + The loop checks the flag between rounds; when set, it finishes + with one final no-tools answer instead of executing more tools. + """ + self._stop_requested = True + + @property + def max_rounds_reached(self) -> bool: + """``True`` when the last tool loop hit ``max_tool_rounds`` and + completed with a graceful final answer instead of raising.""" + return self._max_rounds_reached def add_context(self, role: str, content: str, images: list[ImageInput] | None = None) -> None: """Seed the history with a user or assistant message.""" @@ -432,6 +455,128 @@ def _truncate_tool_result(self, result_str: str) -> str: result_str[: self._max_tool_result_length] + f"\n\n[... result truncated ({len(result_str):,} chars total)]" ) + def _trim_history(self) -> None: + """Trim history to the sliding window without orphaning tool pairs. + + A naive tail-slice can cut between an assistant ``tool_calls`` + message and its ``tool`` results, which providers reject with a + 400. After slicing, drop any leading ``tool`` messages whose + matching assistant message was cut. + """ + if self._max_history_messages is None or len(self._messages) <= self._max_history_messages: + return + self._messages = self._messages[-self._max_history_messages :] + while self._messages and self._messages[0].get("role") == "tool": + self._messages.pop(0) + + @staticmethod + def _malformed_arguments_message(tc: dict[str, Any]) -> str | None: + """Return an error tool-result string when the driver flagged the + tool call's arguments as malformed (``arguments_error``, C1) or + truncated (C3); ``None`` when the call is safe to execute.""" + args_error = tc.get("arguments_error") + if args_error: + return ( + f"Error: arguments for tool '{tc.get('name')}' could not be parsed ({args_error}). " + "Retry the tool call with valid JSON arguments." + ) + if tc.get("truncated"): + return ( + f"Error: arguments for tool '{tc.get('name')}' were truncated. " + "Retry the tool call with smaller, valid arguments." + ) + return None + + def _call_tool(self, name: str, arguments: dict[str, Any], timeout: float | None) -> Any: + """Execute a registered tool, enforcing *timeout* seconds when set. + + Raises ``TimeoutError`` when the tool does not finish in time. A + timed-out worker thread cannot be killed in Python; it is detached + (``shutdown(wait=False)``) so the loop moves on. + """ + if timeout is None: + return self._tools.execute(name, arguments) + pool = ThreadPoolExecutor(max_workers=1) + try: + future = pool.submit(self._tools.execute, name, arguments) + return future.result(timeout=timeout) + finally: + pool.shutdown(wait=False) + + def _execute_tool_call(self, tc: dict[str, Any], timeout: float | None) -> tuple[str, bool]: + """Execute one tool call, returning ``(result_str, is_error)``. + + Never executes the tool when its arguments were flagged as + malformed/truncated — the error is fed back as the tool result so + the model can retry with valid arguments. + """ + malformed = self._malformed_arguments_message(tc) + if malformed is not None: + return malformed, True + from ..extraction.tukuy_bridge import current_tool_call_id + + token = current_tool_call_id.set(tc["id"]) + try: + result = self._call_tool(tc["name"], tc["arguments"], timeout) + return (json.dumps(result) if not isinstance(result, str) else result), False + except (TimeoutError, FuturesTimeoutError): + return f"Error: tool '{tc['name']}' timed out after {timeout}s", True + except Exception as exc: + return ( + f"Error ({type(exc).__name__}): {exc}" if str(exc) else f"Error: {type(exc).__name__}: {exc!r}" + ), True + finally: + current_tool_call_id.reset(token) + + def _final_answer_without_tools(self, msgs: list[dict[str, Any]], merged: dict[str, Any]) -> str: + """One final driver call with tools removed, asking the model to + answer from the tool results it already has. + + Used for graceful exits: max-rounds exhaustion (C4) and + cooperative stop (C2). + """ + self._check_budget() + final_msgs = [ + *msgs, + { + "role": "user", + "content": ( + "You cannot call any more tools. Answer the user's question now, " + "using only the tool results already gathered above." + ), + }, + ] + resp = self._driver.generate_messages_with_hooks(final_msgs, merged) + text: str = resp.get("text", "") + self._last_reasoning = resp.get("reasoning_content") + self._accumulate_usage(resp.get("meta", {})) + self._messages.append({"role": "assistant", "content": text}) + return text + + def _final_answer_simulated(self, augmented_system: str, merged: dict[str, Any]) -> str: + """Graceful-exit final call for the simulated-tools loop (C4/C2).""" + from .simulated_tools import parse_simulated_response + + self._check_budget() + msgs: list[dict[str, Any]] = [{"role": "system", "content": augmented_system}] + msgs.extend(self._messages) + msgs.append( + { + "role": "user", + "content": ( + "You cannot call any more tools. Provide your final answer now " + "(in the final_answer format) based on the tool results already gathered." + ), + } + ) + resp = self._driver.generate_messages_with_hooks(msgs, merged) + text = resp.get("text", "") + self._accumulate_usage(resp.get("meta", {})) + parsed = parse_simulated_response(text, self._tools) + answer: str = parsed["content"] if parsed["type"] == "final_answer" else text + self._messages.append({"role": "assistant", "content": answer}) + return answer + def _build_messages(self, user_content: str, images: list[ImageInput] | None = None) -> list[dict[str, Any]]: """Build the full messages array for an API call.""" msgs: list[dict[str, Any]] = [] @@ -464,8 +609,7 @@ def _accumulate_usage(self, meta: dict[str, Any]) -> None: self._usage["turns"], ) # Trim history to sliding window (system prompt lives outside _messages) - if self._max_history_messages is not None and len(self._messages) > self._max_history_messages: - self._messages = self._messages[-self._max_history_messages :] + self._trim_history() self._maybe_auto_save() def ask( @@ -495,6 +639,10 @@ def ask( return self._ask_with_simulated_tools(content, options, images=images) elif use_native and self._simulated_tools is not True: # type: ignore[comparison-overlap] return self._ask_with_tools(content, options, images=images) + logger.warning( + "Tools are registered but the driver does not support tool use and " + "simulated_tools is disabled; tools are being ignored for this call." + ) self._check_budget() merged = {**self._options, **(options or {})} @@ -520,6 +668,8 @@ def _ask_with_tools( images: list[ImageInput] | None = None, ) -> str: """Execute the tool-use loop: send -> check tool_calls -> execute -> re-send.""" + self._stop_requested = False + self._max_rounds_reached = False merged = {**self._options, **(options or {})} tool_defs = self._tools.to_openai_format() @@ -529,6 +679,9 @@ def _ask_with_tools( msgs = self._build_messages_raw() for _round in range(self._max_tool_rounds): + if self._stop_requested: + logger.info("Stop requested; finishing with a final answer (no more tool calls)") + return self._final_answer_without_tools(msgs, merged) self._check_budget() if self._run_before_turn(): msgs = self._build_messages_raw() @@ -546,6 +699,12 @@ def _ask_with_tools( self._messages.append({"role": "assistant", "content": text}) return text + # Ensure every tool call has a usable id so parallel calls + # can't collide in _full_tool_results. + for tc in tool_calls: + if not tc.get("id"): + tc["id"] = f"call_{uuid.uuid4().hex}" + # Record assistant message with tool_calls assistant_msg: dict[str, Any] = {"role": "assistant", "content": text} assistant_msg["tool_calls"] = [ @@ -565,14 +724,9 @@ def _ask_with_tools( msgs.append(assistant_msg) # Execute each tool call and append results + timeout = merged.get("tool_timeout", self._tool_timeout) for tc in tool_calls: - try: - result = self._tools.execute(tc["name"], tc["arguments"]) - result_str = json.dumps(result) if not isinstance(result, str) else result - except Exception as exc: - result_str = ( - f"Error ({type(exc).__name__}): {exc}" if str(exc) else f"Error: {type(exc).__name__}: {exc!r}" - ) + result_str, _is_error = self._execute_tool_call(tc, timeout) # Preserve full result for step extraction before truncating self._full_tool_results[tc["id"]] = result_str @@ -585,7 +739,14 @@ def _ask_with_tools( self._messages.append(tool_result_msg) msgs.append(tool_result_msg) - raise RuntimeError(f"Tool execution loop exceeded {self._max_tool_rounds} rounds") + # Max rounds exhausted — one graceful final answer without tools + # instead of raising (C4). + logger.warning( + "Tool loop reached max_tool_rounds=%d; requesting a final answer without tools", + self._max_tool_rounds, + ) + self._max_rounds_reached = True + return self._final_answer_without_tools(msgs, merged) def ask_with_tool_events( self, @@ -618,6 +779,8 @@ def ask_with_tool_events( return # Native tool calling with event emission + self._stop_requested = False + self._max_rounds_reached = False merged = {**self._options, **(options or {})} tool_defs = self._tools.to_openai_format() @@ -626,6 +789,12 @@ def ask_with_tool_events( msgs = self._build_messages_raw() for _round in range(self._max_tool_rounds): + if self._stop_requested: + logger.info("Stop requested; finishing with a final answer (no more tool calls)") + final_text = self._final_answer_without_tools(msgs, merged) + yield {"type": "text_delta", "text": final_text} + return + self._check_budget() if self._run_before_turn(): msgs = self._build_messages_raw() resp = self._driver.generate_messages_with_tools_with_hooks(msgs, tool_defs, merged) @@ -644,6 +813,12 @@ def ask_with_tool_events( yield {"type": "text_delta", "text": text} return + # Ensure every tool call has a usable id so parallel calls + # can't collide in _full_tool_results. + for tc in tool_calls: + if not tc.get("id"): + tc["id"] = f"call_{uuid.uuid4().hex}" + # Record assistant message with tool_calls assistant_msg: dict[str, Any] = {"role": "assistant", "content": text} assistant_msg["tool_calls"] = [ @@ -661,6 +836,7 @@ def ask_with_tool_events( msgs.append(assistant_msg) # Execute each tool and yield events + timeout = merged.get("tool_timeout", self._tool_timeout) for tc in tool_calls: yield { "type": "tool_call", @@ -668,13 +844,10 @@ def ask_with_tool_events( "arguments": tc["arguments"], "id": tc["id"], } - try: - result = self._tools.execute(tc["name"], tc["arguments"]) - result_str = json.dumps(result) if not isinstance(result, str) else result - except Exception as exc: - result_str = ( - f"Error ({type(exc).__name__}): {exc}" if str(exc) else f"Error: {type(exc).__name__}: {exc!r}" - ) + result_str, _is_error = self._execute_tool_call(tc, timeout) + + # Preserve full result for step extraction (parity with _ask_with_tools) + self._full_tool_results[tc["id"]] = result_str # Yield the FULL result for UI consumers yield { @@ -693,7 +866,14 @@ def ask_with_tool_events( self._messages.append(tool_result_msg) msgs.append(tool_result_msg) - raise RuntimeError(f"Tool execution loop exceeded {self._max_tool_rounds} rounds") + # Max rounds exhausted — graceful final answer without tools (C4). + logger.warning( + "Tool loop reached max_tool_rounds=%d; requesting a final answer without tools", + self._max_tool_rounds, + ) + self._max_rounds_reached = True + final_text = self._final_answer_without_tools(msgs, merged) + yield {"type": "text_delta", "text": final_text} def ask_live( self, @@ -761,6 +941,8 @@ def ask_live( yield TurnComplete(usage=dict(self._usage)) return + self._stop_requested = False + self._max_rounds_reached = False merged = {**self._options, **(options or {})} tool_defs = self._tools.to_openai_format() @@ -769,6 +951,12 @@ def ask_live( msgs = self._build_messages_raw() for round_idx in range(self._max_tool_rounds): + if self._stop_requested: + logger.info("Stop requested; finishing with a final answer (no more tool calls)") + final_text = self._final_answer_without_tools(msgs, merged) + yield TextDelta(text=final_text) + yield TurnComplete(usage=dict(self._usage)) + return self._check_budget() if self._run_before_turn(): msgs = self._build_messages_raw() @@ -778,7 +966,7 @@ def ask_live( assistant_text_parts: list[str] = [] assistant_thinking_parts: list[str] = [] tool_calls_in_turn: list[dict[str, Any]] = [] - pending_tools: list[tuple[str, str, dict[str, Any]]] = [] + pending_tools: list[dict[str, Any]] = [] turn_usage: dict[str, Any] = {} stop_reason: str = "end_turn" @@ -791,10 +979,18 @@ def ask_live( elif et == "thinking_delta": assistant_thinking_parts.append(event.text) elif et == "tool_use_stop": - pending_tools.append((event.id, event.name, event.input)) + tool_use_id = event.id or f"call_{uuid.uuid4().hex}" + pending_tools.append( + { + "id": tool_use_id, + "name": event.name, + "arguments": event.input, + "truncated": getattr(event, "truncated", False), + } + ) tool_calls_in_turn.append( { - "id": event.id, + "id": tool_use_id, "type": "function", "function": {"name": event.name, "arguments": json.dumps(event.input)}, } @@ -822,25 +1018,18 @@ def ask_live( yield TurnComplete(usage=dict(self._usage)) return - for tool_id, tool_name, tool_input in pending_tools: - try: - result = self._tools.execute(tool_name, tool_input) - result_str = json.dumps(result) if not isinstance(result, str) else result - is_error = False - except Exception as exc: - result_str = ( - f"Error ({type(exc).__name__}): {exc}" if str(exc) else f"Error: {type(exc).__name__}: {exc!r}" - ) - is_error = True + timeout = merged.get("tool_timeout", self._tool_timeout) + for tc in pending_tools: + result_str, is_error = self._execute_tool_call(tc, timeout) - self._full_tool_results[tool_id] = result_str - yield ToolResult(id=tool_id, name=tool_name, output=result_str, is_error=is_error) + self._full_tool_results[tc["id"]] = result_str + yield ToolResult(id=tc["id"], name=tc["name"], output=result_str, is_error=is_error) - truncated = self._truncate_tool_result(result_str) + truncated_result = self._truncate_tool_result(result_str) tool_result_msg: dict[str, Any] = { "role": "tool", - "tool_call_id": tool_id, - "content": truncated, + "tool_call_id": tc["id"], + "content": truncated_result, } self._messages.append(tool_result_msg) msgs.append(tool_result_msg) @@ -849,7 +1038,15 @@ def ask_live( # the real signal is whether pending_tools was non-empty. del stop_reason - raise RuntimeError(f"ask_live exceeded {self._max_tool_rounds} rounds") + # Max rounds exhausted — graceful final answer without tools (C4). + logger.warning( + "ask_live reached max_tool_rounds=%d; requesting a final answer without tools", + self._max_tool_rounds, + ) + self._max_rounds_reached = True + final_text = self._final_answer_without_tools(msgs, merged) + yield TextDelta(text=final_text) + yield TurnComplete(usage=dict(self._usage)) def _ask_with_simulated_tool_events( self, @@ -860,6 +1057,8 @@ def _ask_with_simulated_tool_events( """Simulated tool calling with event emission.""" from .simulated_tools import build_tool_prompt, format_tool_result, parse_simulated_response + self._stop_requested = False + self._max_rounds_reached = False merged = {**self._options, **(options or {})} tool_prompt = build_tool_prompt(self._tools) @@ -871,6 +1070,12 @@ def _ask_with_simulated_tool_events( self._messages.append({"role": "user", "content": user_content}) for _round in range(self._max_tool_rounds): + if self._stop_requested: + logger.info("Stop requested; finishing with a final answer (no more tool calls)") + answer = self._final_answer_simulated(augmented_system, merged) + yield {"type": "text_delta", "text": answer} + return + self._check_budget() self._run_before_turn() msgs: list[dict[str, Any]] = [] msgs.append({"role": "system", "content": augmented_system}) @@ -896,17 +1101,32 @@ def _ask_with_simulated_tool_events( yield {"type": "tool_call", "name": tool_name, "arguments": tool_args, "id": ""} + timeout = merged.get("tool_timeout", self._tool_timeout) + from ..extraction.tukuy_bridge import current_tool_call_id as _sim_tc_id + + _sim_token = _sim_tc_id.set("") try: - result = self._tools.execute(tool_name, tool_args) + result = self._call_tool(tool_name, tool_args, timeout) result_msg = format_tool_result(tool_name, result) + except (TimeoutError, FuturesTimeoutError): + result_msg = format_tool_result(tool_name, f"Error: tool '{tool_name}' timed out after {timeout}s") except Exception as exc: result_msg = format_tool_result(tool_name, f"Error: {exc}") + finally: + _sim_tc_id.reset(_sim_token) yield {"type": "tool_result", "name": tool_name, "result": result_msg, "id": ""} self._messages.append({"role": "user", "content": self._truncate_tool_result(result_msg)}) - raise RuntimeError(f"Simulated tool execution loop exceeded {self._max_tool_rounds} rounds") + # Max rounds exhausted — graceful final answer without tools (C4). + logger.warning( + "Simulated tool loop reached max_tool_rounds=%d; requesting a final answer without tools", + self._max_tool_rounds, + ) + self._max_rounds_reached = True + answer = self._final_answer_simulated(augmented_system, merged) + yield {"type": "text_delta", "text": answer} def _ask_with_simulated_tools( self, @@ -917,6 +1137,8 @@ def _ask_with_simulated_tools( """Prompt-based tool calling for drivers without native tool use.""" from .simulated_tools import build_tool_prompt, format_tool_result, parse_simulated_response + self._stop_requested = False + self._max_rounds_reached = False merged = {**self._options, **(options or {})} tool_prompt = build_tool_prompt(self._tools) @@ -930,6 +1152,9 @@ def _ask_with_simulated_tools( self._messages.append({"role": "user", "content": user_content}) for _round in range(self._max_tool_rounds): + if self._stop_requested: + logger.info("Stop requested; finishing with a final answer (no more tool calls)") + return self._final_answer_simulated(augmented_system, merged) self._check_budget() self._run_before_turn() # Build messages with the augmented system prompt @@ -956,16 +1181,25 @@ def _ask_with_simulated_tools( # Record assistant's tool call as an assistant message self._messages.append({"role": "assistant", "content": text}) + timeout = merged.get("tool_timeout", self._tool_timeout) try: - result = self._tools.execute(tool_name, tool_args) + result = self._call_tool(tool_name, tool_args, timeout) result_msg = format_tool_result(tool_name, result) + except (TimeoutError, FuturesTimeoutError): + result_msg = format_tool_result(tool_name, f"Error: tool '{tool_name}' timed out after {timeout}s") except Exception as exc: result_msg = format_tool_result(tool_name, f"Error: {exc}") # Record tool result as a user message (truncated for the LLM) self._messages.append({"role": "user", "content": self._truncate_tool_result(result_msg)}) - raise RuntimeError(f"Simulated tool execution loop exceeded {self._max_tool_rounds} rounds") + # Max rounds exhausted — graceful final answer without tools (C4). + logger.warning( + "Simulated tool loop reached max_tool_rounds=%d; requesting a final answer without tools", + self._max_tool_rounds, + ) + self._max_rounds_reached = True + return self._final_answer_simulated(augmented_system, merged) def _build_messages_raw(self) -> list[dict[str, Any]]: """Build messages array from system prompt + full history (including tool messages).""" diff --git a/prompture/agents/types.py b/prompture/agents/types.py index 41fc246e..5c8ff70a 100644 --- a/prompture/agents/types.py +++ b/prompture/agents/types.py @@ -86,9 +86,9 @@ def __init__( class RunContext(Generic[DepsType]): """Dependency-injection context available to tools and guardrails. - Built at the start of each :meth:`Agent.run` invocation and passed - automatically to tools whose first parameter is annotated as - ``RunContext``. + Rebuilt for each tool invocation during a run (so tools see the live + iteration, message history, and usage) and passed automatically to + tools whose first parameter is annotated as ``RunContext``. Attributes: deps: User-supplied dependencies (database handles, API clients, etc.). @@ -116,14 +116,17 @@ class AgentCallbacks: fires at the HTTP/driver layer. Attributes: - on_step: Called for each step during execution. + on_step: Called with each :class:`AgentStep` after the run + completes (steps are collected during execution and replayed + at the end of the run; this is not a live, mid-loop hook). on_tool_start: Called before a tool is invoked with (name, args). on_tool_end: Called after a tool completes with (name, result). on_iteration: Called at the start of each iteration with the index. on_output: Called when the agent produces final output. - on_thinking: Called when the agent emits thinking/reasoning content. - The callback receives the thinking text (e.g., content within - tags for models that support chain-of-thought). + on_thinking: Called after the run completes, once per thinking + step, with the thinking text (e.g., content within + tags for models that support chain-of-thought). Like + ``on_step`` this fires at the end of the run, not mid-loop. on_approval_needed: Called when a tool raises ApprovalRequired. The callback receives (tool_name, action, details) and should return True to approve execution or False to deny. diff --git a/tests/test_agent.py b/tests/test_agent.py index e6da9496..ddef921a 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -2172,3 +2172,642 @@ async def _test(): assert result.run_usage["call_count"] >= 1 asyncio.run(_test()) + + +# =========================================================================== +# Tool-calling audit fixes (A1, H2, C2, C4, M8/A11, L6/A12, L10/A5, M10) +# =========================================================================== + + +def _tool_call_response(call_id: str, name: str, arguments: dict) -> dict: + """Build a driver response dict requesting a single tool call.""" + return { + "text": "", + "meta": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "cost": 0.001}, + "tool_calls": [{"id": call_id, "name": name, "arguments": arguments}], + "stop_reason": "tool_use", + } + + +def _text_response(text: str) -> dict: + """Build a plain final-text driver response dict.""" + return { + "text": text, + "meta": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "cost": 0.001}, + "tool_calls": [], + "stop_reason": "end_turn", + } + + +def _tool_message_contents(result: AgentResult) -> list[str]: + return [m.get("content", "") for m in result.messages if m.get("role") == "tool"] + + +# --------------------------------------------------------------------------- +# A1: sync approval flow (incl. awaitable approval handlers) +# --------------------------------------------------------------------------- + + +class TestSyncApprovalFlow: + def test_approval_granted_retries_and_executes(self): + """Granted approval retries the tool and uses its real result.""" + from prompture.agents.types import ApprovalRequired + + calls = {"n": 0} + approval_log: list[tuple[str, str]] = [] + + def risky_tool(command: str) -> str: + """Run a risky command.""" + calls["n"] += 1 + if calls["n"] == 1: + raise ApprovalRequired("risky_tool", f"Execute: {command}") + return f"ran: {command}" + + def approve(tool_name: str, action: str, details: dict) -> bool: + approval_log.append((tool_name, action)) + return True + + responses = [ + _tool_call_response("call_1", "risky_tool", {"command": "ls"}), + _text_response("Done."), + ] + cb = AgentCallbacks(on_approval_needed=approve) + agent = Agent("test/model", driver=MockToolDriver(responses), tools=[risky_tool], agent_callbacks=cb) + result = agent.run("run ls") + + assert approval_log == [("risky_tool", "Execute: ls")] + assert calls["n"] == 2 # initial attempt + approved retry + assert any("ran: ls" in c for c in _tool_message_contents(result)) + + def test_approval_denied_returns_denial_string(self): + """Denied approval feeds a denial string back as the tool result.""" + from prompture.agents.types import ApprovalRequired + + def risky_tool(command: str) -> str: + """Run a risky command.""" + raise ApprovalRequired("risky_tool", f"Execute: {command}") + + responses = [ + _tool_call_response("call_1", "risky_tool", {"command": "rm -rf"}), + _text_response("Understood."), + ] + cb = AgentCallbacks(on_approval_needed=lambda name, action, details: False) + agent = Agent("test/model", driver=MockToolDriver(responses), tools=[risky_tool], agent_callbacks=cb) + result = agent.run("delete everything") + + assert any("execution denied" in c for c in _tool_message_contents(result)) + + def test_no_approval_handler_returns_error_string(self): + """Without an approval handler, an explanatory error is fed back.""" + from prompture.agents.types import ApprovalRequired + + def risky_tool(command: str) -> str: + """Run a risky command.""" + raise ApprovalRequired("risky_tool", f"Execute: {command}") + + responses = [ + _tool_call_response("call_1", "risky_tool", {"command": "ls"}), + _text_response("OK."), + ] + agent = Agent("test/model", driver=MockToolDriver(responses), tools=[risky_tool]) + result = agent.run("run ls") + + assert any("no approval handler is configured" in c for c in _tool_message_contents(result)) + + def test_approval_granted_but_tool_raises_again_no_infinite_loop(self): + """A tool that re-raises ApprovalRequired after approval stops after one retry.""" + from prompture.agents.types import ApprovalRequired + + approval_calls = {"n": 0} + + def always_risky(command: str) -> str: + """Always requires approval.""" + raise ApprovalRequired("always_risky", f"Execute: {command}") + + def approve(tool_name: str, action: str, details: dict) -> bool: + approval_calls["n"] += 1 + return True + + responses = [ + _tool_call_response("call_1", "always_risky", {"command": "ls"}), + _text_response("Done."), + ] + cb = AgentCallbacks(on_approval_needed=approve) + agent = Agent("test/model", driver=MockToolDriver(responses), tools=[always_risky], agent_callbacks=cb) + result = agent.run("run ls") + + assert approval_calls["n"] == 1 # handler consulted exactly once + assert any("approval was already granted" in c for c in _tool_message_contents(result)) + + def test_async_approval_handler_is_awaited(self): + """An ``async def`` approval handler returning False must DENY. + + Regression for A1: calling the handler bare returns a truthy + coroutine, which silently auto-approved everything. + """ + from prompture.agents.types import ApprovalRequired + + def risky_tool(command: str) -> str: + """Run a risky command.""" + raise ApprovalRequired("risky_tool", f"Execute: {command}") + + async def deny_async(tool_name: str, action: str, details: dict) -> bool: + return False + + responses = [ + _tool_call_response("call_1", "risky_tool", {"command": "rm -rf"}), + _text_response("Understood."), + ] + cb = AgentCallbacks(on_approval_needed=deny_async) + agent = Agent("test/model", driver=MockToolDriver(responses), tools=[risky_tool], agent_callbacks=cb) + result = agent.run("delete everything") + + assert any("execution denied" in c for c in _tool_message_contents(result)) + + +# --------------------------------------------------------------------------- +# H2: genuinely async tools through AsyncAgent +# --------------------------------------------------------------------------- + + +class TestAsyncToolExecution: + def test_async_def_tool_executed_through_async_agent(self): + """An ``async def`` tool runs to completion in AsyncAgent.run().""" + + async def async_echo(text: str) -> str: + """Echo text asynchronously.""" + await asyncio.sleep(0) + return f"async: {text}" + + responses = [ + _tool_call_response("call_1", "async_echo", {"text": "hi"}), + _text_response("Done."), + ] + agent = AsyncAgent("test/model", driver=MockAsyncToolDriver(responses), tools=[async_echo]) + result = asyncio.run(agent.run("echo hi")) + + assert result.output == "Done." + assert any("async: hi" in c for c in _tool_message_contents(result)) + + def test_async_tool_wrapper_registers_async_fn(self): + """Wrapped async tools expose ``_async_fn`` for ToolRegistry.aexecute.""" + from prompture.infra.session import UsageSession + + async def a_fn(x: str) -> str: + """Async tool.""" + return f"got {x}" + + agent = AsyncAgent("test/model", driver=MockAsyncDriver(), tools=[a_fn]) + ctx = agent._build_run_context("test", None, UsageSession(), [], 0) + wrapped = agent._wrap_tools_with_context(ctx) + + td = wrapped.get("a_fn") + assert td is not None + async_fn = getattr(td.function, "_async_fn", None) + assert async_fn is not None + assert asyncio.run(async_fn(x="hey")) == "got hey" + + def test_async_tool_sees_current_tool_call_id_contextvar(self): + """Contextvars propagate into async tools (no thread/asyncio.run bridge).""" + from prompture.extraction.tukuy_bridge import current_tool_call_id + + seen: list[str | None] = [] + + async def ctx_probe() -> str: + """Record the ambient tool_call_id.""" + seen.append(current_tool_call_id.get()) + return "ok" + + responses = [ + _tool_call_response("call_42", "ctx_probe", {}), + _text_response("Done."), + ] + agent = AsyncAgent("test/model", driver=MockAsyncToolDriver(responses), tools=[ctx_probe]) + asyncio.run(agent.run("probe")) + + assert seen == ["call_42"] + + def test_async_tool_with_async_approval_handler(self): + """Async approval handler is awaited in the async wrapper path.""" + from prompture.agents.types import ApprovalRequired + + calls = {"n": 0} + + async def risky_async(command: str) -> str: + """Risky async tool.""" + calls["n"] += 1 + if calls["n"] == 1: + raise ApprovalRequired("risky_async", f"Execute: {command}") + return f"ran: {command}" + + async def approve_async(tool_name: str, action: str, details: dict) -> bool: + return True + + responses = [ + _tool_call_response("call_1", "risky_async", {"command": "ls"}), + _text_response("Done."), + ] + cb = AgentCallbacks(on_approval_needed=approve_async) + agent = AsyncAgent("test/model", driver=MockAsyncToolDriver(responses), tools=[risky_async], agent_callbacks=cb) + result = asyncio.run(agent.run("run ls")) + + assert calls["n"] == 2 + assert any("ran: ls" in c for c in _tool_message_contents(result)) + + +# --------------------------------------------------------------------------- +# C2: Agent.stop() cooperative shutdown +# --------------------------------------------------------------------------- + + +class TestAgentStopDelegation: + def test_stop_delegates_to_conversation_request_stop(self): + """stop() sets the agent flag and forwards to the conversation.""" + + class FakeConv: + def __init__(self): + self.calls = 0 + + def request_stop(self): + self.calls += 1 + + agent = Agent("test/model", driver=MockDriver()) + fake = FakeConv() + agent._conversation = fake # type: ignore[assignment] + agent.stop() + + assert agent._stop_requested is True + assert fake.calls == 1 + + def test_stop_without_conversation_is_safe(self): + """stop() with no active conversation only sets the agent flag.""" + agent = Agent("test/model", driver=MockDriver()) + agent.stop() + assert agent._stop_requested is True + + def test_stop_halts_running_tool_loop(self): + """A tool calling agent.stop() ends the loop with a graceful answer.""" + holder: dict[str, Any] = {} + + def stopper() -> str: + """Request the agent to stop.""" + holder["agent"].stop() + return "stop requested" + + responses = [ + _tool_call_response("call_1", "stopper", {}), + _text_response("final answer after stop"), + ] + driver = MockToolDriver(responses) + agent = Agent( + "test/model", + driver=driver, + tools=[stopper], + max_iterations=5, + persistent_conversation=True, + ) + holder["agent"] = agent + result = agent.run("start") + + assert result.output == "final answer after stop" + # 1 tool-call round + 1 final answer call; the loop did not run all 5 rounds + assert driver._call_idx == 2 + + def test_async_stop_delegates_to_conversation_request_stop(self): + """AsyncAgent.stop() forwards to the conversation as well.""" + + class FakeConv: + def __init__(self): + self.calls = 0 + + def request_stop(self): + self.calls += 1 + + agent = AsyncAgent("test/model", driver=MockAsyncDriver()) + fake = FakeConv() + agent._conversation = fake + agent.stop() + + assert agent._stop_requested is True + assert fake.calls == 1 + + def test_stop_halts_loop_without_persistent_conversation(self): + """stop() reaches the loop of a default (non-persistent) agent. + + Regression guard: ``_conversation`` is only populated when + ``persistent_conversation=True``, so forwarding solely through it made + ``stop()`` a silent no-op for the default agent. + """ + holder: dict[str, Any] = {} + + def stopper() -> str: + """Request the agent to stop.""" + holder["agent"].stop() + return "stop requested" + + responses = [ + _tool_call_response("call_1", "stopper", {}), + _text_response("final answer after stop"), + ] + driver = MockToolDriver(responses) + agent = Agent("test/model", driver=driver, tools=[stopper], max_iterations=5) + holder["agent"] = agent + result = agent.run("start") + + assert result.output == "final answer after stop" + # 1 tool-call round + 1 final answer; the loop did not run all 5 rounds. + assert driver._call_idx == 2 + + def test_stop_before_run_is_not_lost_to_the_race(self): + """A stop() landing before the conversation exists is replayed onto it.""" + + def ping() -> str: + """Ping.""" + return "pong" + + agent = Agent("test/model", driver=MockToolDriver([_text_response("hi")]), tools=[ping]) + conv = agent._build_conversation() + assert conv._stop_requested is False + + agent.stop() + # Rebuilding mid-run must carry the pending request over. + conv2 = agent._build_conversation() + assert conv2._stop_requested is True + + def test_async_stop_halts_loop_without_persistent_conversation(self): + """AsyncAgent.stop() also reaches a non-persistent run's loop.""" + holder: dict[str, Any] = {} + + def stopper() -> str: + """Request the agent to stop.""" + holder["agent"].stop() + return "stop requested" + + async def _test(): + responses = [ + _tool_call_response("call_1", "stopper", {}), + _text_response("async final answer after stop"), + ] + driver = MockAsyncToolDriver(responses) + agent = AsyncAgent("test/model", driver=driver, tools=[stopper], max_iterations=5) + holder["agent"] = agent + result = await agent.run("start") + + assert result.output == "async final answer after stop" + assert driver._call_idx == 2 + + asyncio.run(_test()) + + +# --------------------------------------------------------------------------- +# Per-tool timeout plumbed from the Agent +# --------------------------------------------------------------------------- + + +class TestAgentToolTimeout: + def test_tool_timeout_forwarded_to_conversation(self): + """Agent(tool_timeout=...) reaches the Conversation that runs tools.""" + + def ping() -> str: + """Ping.""" + return "pong" + + agent = Agent("test/model", driver=MockToolDriver([_text_response("hi")]), tools=[ping], tool_timeout=2.5) + conv = agent._build_conversation() + assert conv._tool_timeout == 2.5 + + def test_tool_timeout_defaults_to_none(self): + """No timeout unless asked for, preserving the previous behaviour.""" + agent = Agent("test/model", driver=MockToolDriver([_text_response("hi")])) + assert agent._build_conversation()._tool_timeout is None + + def test_slow_tool_times_out_and_reports_to_the_model(self): + """A tool overrunning the budget yields an error result, not a hang.""" + import time as _time + + def slow() -> str: + """Sleep past the timeout.""" + _time.sleep(5) + return "never returned" + + responses = [ + _tool_call_response("call_1", "slow", {}), + _text_response("gave up on the slow tool"), + ] + agent = Agent( + "test/model", + driver=MockToolDriver(responses), + tools=[slow], + tool_timeout=0.05, + ) + result = agent.run("call the slow tool") + + assert result.output == "gave up on the slow tool" + assert any("timed out" in c for c in _tool_message_contents(result)) + + def test_async_tool_timeout_forwarded_to_conversation(self): + """AsyncAgent(tool_timeout=...) reaches the AsyncConversation too.""" + agent = AsyncAgent("test/model", driver=MockAsyncToolDriver([_text_response("hi")]), tool_timeout=1.5) + assert agent._build_conversation()._tool_timeout == 1.5 + + +# --------------------------------------------------------------------------- +# C4: graceful finish when max_iterations is exhausted +# --------------------------------------------------------------------------- + + +class TestMaxIterationsGraceful: + def test_max_iterations_graceful_finish(self): + """Exhausting max_iterations yields a final answer instead of RuntimeError.""" + + def ping() -> str: + """Ping.""" + return "pong" + + responses = [ + _tool_call_response("call_1", "ping", {}), + _tool_call_response("call_2", "ping", {}), + _text_response("Answering from what I have."), + ] + driver = MockToolDriver(responses) + agent = Agent( + "test/model", + driver=driver, + tools=[ping], + max_iterations=2, + persistent_conversation=True, + ) + result = agent.run("keep pinging") + + assert result.output == "Answering from what I have." + assert result.state == AgentState.idle + assert agent.conversation is not None + assert agent.conversation.max_rounds_reached is True + + def test_async_max_iterations_graceful_finish(self): + """AsyncAgent also finishes gracefully when tool rounds are exhausted.""" + + def ping() -> str: + """Ping.""" + return "pong" + + async def _test(): + responses = [ + _tool_call_response("call_1", "ping", {}), + _tool_call_response("call_2", "ping", {}), + _text_response("Async final answer."), + ] + driver = MockAsyncToolDriver(responses) + agent = AsyncAgent( + "test/model", + driver=driver, + tools=[ping], + max_iterations=2, + persistent_conversation=True, + ) + result = await agent.run("keep pinging") + assert result.output == "Async final answer." + assert agent.conversation is not None + assert agent.conversation.max_rounds_reached is True + + asyncio.run(_test()) + + +# --------------------------------------------------------------------------- +# L10/A5 + M8/A11 + L6/A12: live RunContext, real tool names, durations +# --------------------------------------------------------------------------- + + +class TestLiveRunContextAndSteps: + def test_run_context_refreshed_per_tool_round(self): + """Tools see live iteration/messages instead of the run-start snapshot.""" + snapshots: list[tuple[int, int]] = [] + + def probe(ctx: RunContext, note: str) -> str: + """Record iteration and message count.""" + snapshots.append((ctx.iteration, len(ctx.messages))) + return f"noted {note}" + + responses = [ + _tool_call_response("call_1", "probe", {"note": "a"}), + _tool_call_response("call_2", "probe", {"note": "b"}), + _text_response("done"), + ] + agent = Agent("test/model", driver=MockToolDriver(responses), tools=[probe], max_iterations=5) + agent.run("go") + + assert len(snapshots) == 2 + assert snapshots[0][0] == 0 # first tool round + assert snapshots[1][0] == 1 # second tool round + assert snapshots[1][1] > snapshots[0][1] # message history grew + + def test_tool_result_step_records_real_tool_name(self): + """tool_result steps carry the tool name, not the tool_call_id.""" + + def my_tool() -> str: + """A named tool.""" + return "result" + + responses = [ + _tool_call_response("call_xyz", "my_tool", {}), + _text_response("ok"), + ] + agent = Agent("test/model", driver=MockToolDriver(responses), tools=[my_tool]) + result = agent.run("go") + + tr_steps = [s for s in result.steps if s.step_type == StepType.tool_result] + assert len(tr_steps) == 1 + assert tr_steps[0].tool_name == "my_tool" + + def test_tool_result_step_has_duration_ms(self): + """tool_result steps carry a measured duration.""" + import time as _time + + def slow_tool() -> str: + """A slow tool.""" + _time.sleep(0.02) + return "done" + + responses = [ + _tool_call_response("call_9", "slow_tool", {}), + _text_response("ok"), + ] + agent = Agent("test/model", driver=MockToolDriver(responses), tools=[slow_tool]) + result = agent.run("go") + + tr_steps = [s for s in result.steps if s.step_type == StepType.tool_result] + assert len(tr_steps) == 1 + assert tr_steps[0].duration_ms >= 10 + + def test_async_run_context_refreshed_per_tool_round(self): + """AsyncAgent tools also see live iteration/messages.""" + snapshots: list[tuple[int, int]] = [] + + async def aprobe(ctx: RunContext, note: str) -> str: + """Record iteration and message count.""" + snapshots.append((ctx.iteration, len(ctx.messages))) + return f"noted {note}" + + async def _test(): + responses = [ + _tool_call_response("call_1", "aprobe", {"note": "a"}), + _tool_call_response("call_2", "aprobe", {"note": "b"}), + _text_response("done"), + ] + agent = AsyncAgent("test/model", driver=MockAsyncToolDriver(responses), tools=[aprobe], max_iterations=5) + await agent.run("go") + + asyncio.run(_test()) + + assert len(snapshots) == 2 + assert snapshots[0][0] == 0 + assert snapshots[1][0] == 1 + assert snapshots[1][1] > snapshots[0][1] + + +# --------------------------------------------------------------------------- +# M10: persistent conversation reuse respects new system prompt / driver +# --------------------------------------------------------------------------- + + +class TestPersistentConversationReuse: + def test_reuse_updates_system_prompt(self): + """A reused conversation picks up the newly resolved system prompt.""" + driver = MockDriver(["one", "two"]) + agent = Agent("test/model", driver=driver, system_prompt="prompt A", persistent_conversation=True) + + agent.run("hi") + assert agent.conversation is not None + assert agent.conversation.system_prompt == "prompt A" + + agent._system_prompt = "prompt B" + agent.run("hi again") + assert agent.conversation.system_prompt == "prompt B" + + def test_reuse_rebuilds_when_driver_swapped(self): + """Swapping the agent's driver discards the stale conversation.""" + d1 = MockDriver(["one"]) + d2 = MockDriver(["two"]) + agent = Agent("test/model", driver=d1, persistent_conversation=True) + + agent.run("hi") + conv1 = agent.conversation + assert conv1 is not None + + agent._driver = d2 + result = agent.run("hi again") + + assert result.output == "two" + assert agent.conversation is not conv1 + assert agent.conversation._driver is d2 + + def test_reuse_keeps_conversation_when_driver_unchanged(self): + """Same driver + same config reuses the conversation (history grows).""" + driver = MockDriver(["one", "two"]) + agent = Agent("test/model", driver=driver, persistent_conversation=True) + + agent.run("hi") + conv1 = agent.conversation + agent.run("hi again") + + assert agent.conversation is conv1 diff --git a/tests/test_conversation_robustness.py b/tests/test_conversation_robustness.py new file mode 100644 index 00000000..c1d83f46 --- /dev/null +++ b/tests/test_conversation_robustness.py @@ -0,0 +1,550 @@ +"""Robustness tests for Conversation/AsyncConversation tool loops. + +Covers the tool-calling audit contracts: malformed/truncated argument +feedback (C1/C3), cooperative stop (C2), graceful max-rounds final +answer (C4), per-tool timeouts, full-result preservation, budget +enforcement in the event path, parallel async execution ordering, and +boundary-aware history trimming. +""" + +from __future__ import annotations + +import asyncio +import time +from typing import Any + +import pytest + +from prompture.agents.async_conversation import AsyncConversation +from prompture.agents.conversation import Conversation +from prompture.agents.tools_schema import ToolRegistry +from prompture.drivers.async_base import AsyncDriver +from prompture.drivers.base import Driver +from prompture.exceptions import BudgetExceededError + +# --------------------------------------------------------------------------- +# Mock drivers +# --------------------------------------------------------------------------- + + +def _meta(cost: float = 0.0) -> dict[str, Any]: + return {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15, "cost": cost} + + +def _tool_call_resp(*tool_calls: dict[str, Any], cost: float = 0.0) -> dict[str, Any]: + return {"text": "", "meta": _meta(cost), "tool_calls": list(tool_calls), "stop_reason": "tool_use"} + + +def _text_resp(text: str, cost: float = 0.0) -> dict[str, Any]: + return {"text": text, "meta": _meta(cost), "tool_calls": [], "stop_reason": "end_turn"} + + +class MockToolDriver(Driver): + supports_messages = True + supports_tool_use = True + + def __init__(self, responses: list[dict[str, Any]]): + self._responses = list(responses) + self._call_idx = 0 + + def generate(self, prompt, options): + return self._get_next() + + def generate_messages(self, messages, options): + return self._get_next() + + def generate_messages_with_tools(self, messages, tools, options): + return self._get_next() + + def _get_next(self): + resp = self._responses[self._call_idx] + self._call_idx += 1 + return resp + + +class MockAsyncToolDriver(AsyncDriver): + supports_messages = True + supports_tool_use = True + + def __init__(self, responses: list[dict[str, Any]]): + self._responses = list(responses) + self._call_idx = 0 + + async def generate(self, prompt, options): + return self._get_next() + + async def generate_messages(self, messages, options): + return self._get_next() + + async def generate_messages_with_tools(self, messages, tools, options): + return self._get_next() + + def _get_next(self): + resp = self._responses[self._call_idx] + self._call_idx += 1 + return resp + + +def _echo_registry() -> ToolRegistry: + reg = ToolRegistry() + + def echo(text: str) -> str: + """Echo text back.""" + return text + + reg.register(echo) + return reg + + +def _tool_messages(conv) -> list[dict[str, Any]]: + return [m for m in conv.messages if m.get("role") == "tool"] + + +# --------------------------------------------------------------------------- +# Unknown tool / missing argument feedback +# --------------------------------------------------------------------------- + + +class TestErrorFeedback: + def test_unknown_tool_feeds_error_and_completes(self): + """Driver requests an unregistered tool -> error tool message -> loop continues.""" + responses = [ + _tool_call_resp({"id": "call_1", "name": "ghost_tool", "arguments": {}}), + _text_resp("Sorry, that tool does not exist."), + ] + conv = Conversation(driver=MockToolDriver(responses), tools=_echo_registry()) + result = conv.ask("Call ghost_tool") + + assert result == "Sorry, that tool does not exist." + tool_msgs = _tool_messages(conv) + assert len(tool_msgs) == 1 + assert tool_msgs[0]["tool_call_id"] == "call_1" + assert "not registered" in tool_msgs[0]["content"] + + def test_missing_required_argument_feedback(self): + """A call missing a required arg gets the validation error fed back, not a crash.""" + responses = [ + _tool_call_resp({"id": "call_1", "name": "echo", "arguments": {}}), + _text_resp("I need the text argument."), + ] + conv = Conversation(driver=MockToolDriver(responses), tools=_echo_registry()) + result = conv.ask("Call echo with nothing") + + assert result == "I need the text argument." + tool_msgs = _tool_messages(conv) + assert len(tool_msgs) == 1 + assert "Missing required argument(s) for 'echo'" in tool_msgs[0]["content"] + + +# --------------------------------------------------------------------------- +# C1/C3: malformed / truncated arguments are not executed +# --------------------------------------------------------------------------- + + +class TestMalformedArguments: + def test_arguments_error_skips_execution(self): + called = False + + def echo(text: str) -> str: + """Echo text back.""" + nonlocal called + called = True + return text + + reg = ToolRegistry() + reg.register(echo) + + responses = [ + _tool_call_resp( + { + "id": "call_1", + "name": "echo", + "arguments": {}, + "arguments_error": "invalid JSON at offset 12", + } + ), + _text_resp("Retrying with valid arguments."), + ] + conv = Conversation(driver=MockToolDriver(responses), tools=reg) + result = conv.ask("Call echo") + + assert result == "Retrying with valid arguments." + assert called is False + tool_msgs = _tool_messages(conv) + assert len(tool_msgs) == 1 + assert "could not be parsed" in tool_msgs[0]["content"] + assert "invalid JSON at offset 12" in tool_msgs[0]["content"] + # Full results still keyed by the tool call id + assert conv._full_tool_results["call_1"] == tool_msgs[0]["content"] + + def test_truncated_arguments_skip_execution(self): + called = False + + def echo(text: str) -> str: + """Echo text back.""" + nonlocal called + called = True + return text + + reg = ToolRegistry() + reg.register(echo) + + responses = [ + _tool_call_resp({"id": "call_1", "name": "echo", "arguments": {}, "truncated": True}), + _text_resp("Recovered."), + ] + conv = Conversation(driver=MockToolDriver(responses), tools=reg) + result = conv.ask("Call echo") + + assert result == "Recovered." + assert called is False + assert "were truncated" in _tool_messages(conv)[0]["content"] + + +# --------------------------------------------------------------------------- +# Tool result truncation vs. full-result preservation +# --------------------------------------------------------------------------- + + +class TestToolResultTruncation: + def test_oversized_result_truncated_in_history_full_in_store(self): + big = "x" * 500 + + def big_tool() -> str: + """Return a large payload.""" + return big + + reg = ToolRegistry() + reg.register(big_tool) + + responses = [ + _tool_call_resp({"id": "call_1", "name": "big_tool", "arguments": {}}), + _text_resp("Done."), + ] + conv = Conversation(driver=MockToolDriver(responses), tools=reg, max_tool_result_length=100) + result = conv.ask("Call big_tool") + + assert result == "Done." + tool_msg = _tool_messages(conv)[0] + # History sees the truncated version + assert len(tool_msg["content"]) < 200 + assert "result truncated" in tool_msg["content"] + # The full result is preserved for step extraction + assert conv._full_tool_results["call_1"] == big + + def test_default_max_tool_result_length(self): + conv = Conversation(driver=MockToolDriver([_text_resp("hi")])) + assert conv._max_tool_result_length == 16000 + + +# --------------------------------------------------------------------------- +# clear() resets full tool results (L8) +# --------------------------------------------------------------------------- + + +class TestClear: + def test_clear_resets_full_tool_results(self): + responses = [ + _tool_call_resp({"id": "call_1", "name": "echo", "arguments": {"text": "hi"}}), + _text_resp("Done."), + ] + conv = Conversation(driver=MockToolDriver(responses), tools=_echo_registry()) + conv.ask("Call echo") + assert conv._full_tool_results + + conv.clear() + assert conv.messages == [] + assert conv._full_tool_results == {} + + +# --------------------------------------------------------------------------- +# Budget enforcement in the event path (L2) +# --------------------------------------------------------------------------- + + +class TestBudgetInEventPath: + def test_budget_enforced_in_ask_with_tool_events(self): + responses = [ + _tool_call_resp({"id": "call_1", "name": "echo", "arguments": {"text": "hi"}}, cost=0.02), + _text_resp("should not get here"), + ] + conv = Conversation( + driver=MockToolDriver(responses), + tools=_echo_registry(), + max_cost=0.01, + budget_policy="hard_stop", + ) + with pytest.raises(BudgetExceededError): + for _event in conv.ask_with_tool_events("Call echo"): + pass + + +# --------------------------------------------------------------------------- +# Cooperative stop mid-loop (C2) +# --------------------------------------------------------------------------- + + +class TestCooperativeStop: + def test_request_stop_mid_loop(self): + executions = 0 + + def echo(text: str) -> str: + """Echo text back.""" + nonlocal executions + executions += 1 + conv.request_stop() + return text + + # Round 1: tool call (tool requests the stop). Round 2 is the + # graceful no-tools final answer. + responses = [ + _tool_call_resp({"id": "call_1", "name": "echo", "arguments": {"text": "hi"}}), + _text_resp("Wrapping up without more tools."), + ] + conv = Conversation(driver=MockToolDriver(responses)) + conv.register_tool(echo) + result = conv.ask("Call echo then keep going") + + assert result == "Wrapping up without more tools." + assert executions == 1 + assert conv.max_rounds_reached is False + + async def test_request_stop_mid_loop_async(self): + executions = 0 + + async def echo(text: str) -> str: + """Echo text back.""" + nonlocal executions + executions += 1 + conv.request_stop() + return text + + responses = [ + _tool_call_resp({"id": "call_1", "name": "echo", "arguments": {"text": "hi"}}), + _text_resp("Wrapping up without more tools."), + ] + conv = AsyncConversation(driver=MockAsyncToolDriver(responses)) + conv.register_tool(echo) + result = await conv.ask("Call echo then keep going") + + assert result == "Wrapping up without more tools." + assert executions == 1 + + +# --------------------------------------------------------------------------- +# Per-tool timeout +# --------------------------------------------------------------------------- + + +class TestToolTimeout: + def test_sync_tool_timeout_feeds_error(self): + def slow() -> str: + """Sleep too long.""" + time.sleep(5) + return "never" + + reg = ToolRegistry() + reg.register(slow) + + responses = [ + _tool_call_resp({"id": "call_1", "name": "slow", "arguments": {}}), + _text_resp("Tool was too slow."), + ] + conv = Conversation(driver=MockToolDriver(responses), tools=reg, tool_timeout=0.1) + result = conv.ask("Call slow") + + assert result == "Tool was too slow." + assert "timed out after 0.1s" in _tool_messages(conv)[0]["content"] + + async def test_async_tool_timeout_feeds_error(self): + async def slow() -> str: + """Sleep too long.""" + await asyncio.sleep(5) + return "never" + + reg = ToolRegistry() + reg.register(slow) + + responses = [ + _tool_call_resp({"id": "call_1", "name": "slow", "arguments": {}}), + _text_resp("Tool was too slow."), + ] + conv = AsyncConversation(driver=MockAsyncToolDriver(responses), tools=reg, tool_timeout=0.1) + result = await conv.ask("Call slow") + + assert result == "Tool was too slow." + assert "timed out after 0.1s" in _tool_messages(conv)[0]["content"] + + +# --------------------------------------------------------------------------- +# Graceful max-rounds final answer (C4) +# --------------------------------------------------------------------------- + + +class TestGracefulMaxRounds: + def test_native_loop_graceful_final_answer(self): + responses = [ + _tool_call_resp({"id": f"call_{i}", "name": "echo", "arguments": {"text": "hi"}}) for i in range(2) + ] + responses.append(_text_resp("Final answer from gathered results.")) + conv = Conversation(driver=MockToolDriver(responses), tools=_echo_registry(), max_tool_rounds=2) + result = conv.ask("Loop forever") + + assert result == "Final answer from gathered results." + assert conv.max_rounds_reached is True + + async def test_async_native_loop_graceful_final_answer(self): + responses = [ + _tool_call_resp({"id": f"call_{i}", "name": "echo", "arguments": {"text": "hi"}}) for i in range(2) + ] + responses.append(_text_resp("Final answer from gathered results.")) + conv = AsyncConversation(driver=MockAsyncToolDriver(responses), tools=_echo_registry(), max_tool_rounds=2) + result = await conv.ask("Loop forever") + + assert result == "Final answer from gathered results." + assert conv.max_rounds_reached is True + + def test_event_path_graceful_final_answer(self): + responses = [ + _tool_call_resp({"id": "call_0", "name": "echo", "arguments": {"text": "hi"}}), + _text_resp("Final answer via events."), + ] + conv = Conversation(driver=MockToolDriver(responses), tools=_echo_registry(), max_tool_rounds=1) + events = list(conv.ask_with_tool_events("Loop forever")) + + assert conv.max_rounds_reached is True + assert events[-1] == {"type": "text_delta", "text": "Final answer via events."} + + +# --------------------------------------------------------------------------- +# Async parallel execution preserves call order +# --------------------------------------------------------------------------- + + +class TestAsyncParallelTools: + async def test_gather_preserves_result_order(self): + finished: list[str] = [] + + async def slow_tool() -> str: + """Slow tool.""" + await asyncio.sleep(0.2) + finished.append("slow_tool") + return "slow-result" + + async def fast_tool() -> str: + """Fast tool.""" + finished.append("fast_tool") + return "fast-result" + + reg = ToolRegistry() + reg.register(slow_tool) + reg.register(fast_tool) + + responses = [ + _tool_call_resp( + {"id": "call_slow", "name": "slow_tool", "arguments": {}}, + {"id": "call_fast", "name": "fast_tool", "arguments": {}}, + ), + _text_resp("Both done."), + ] + conv = AsyncConversation(driver=MockAsyncToolDriver(responses), tools=reg) + result = await conv.ask("Call both") + + assert result == "Both done." + # Completion order differs from call order (parallel execution) + assert finished == ["fast_tool", "slow_tool"] + # Results are recorded in call order so tool_call_id matching holds + tool_msgs = _tool_messages(conv) + assert [m["tool_call_id"] for m in tool_msgs] == ["call_slow", "call_fast"] + assert [m["content"] for m in tool_msgs] == ["slow-result", "fast-result"] + assert conv._full_tool_results == {"call_slow": "slow-result", "call_fast": "fast-result"} + + async def test_sequential_tools_opt_out(self): + started: list[str] = [] + finished: list[str] = [] + + async def slow_tool() -> str: + """Slow tool.""" + started.append("slow_tool") + await asyncio.sleep(0.1) + finished.append("slow_tool") + return "slow-result" + + async def fast_tool() -> str: + """Fast tool.""" + started.append("fast_tool") + finished.append("fast_tool") + return "fast-result" + + reg = ToolRegistry() + reg.register(slow_tool) + reg.register(fast_tool) + + responses = [ + _tool_call_resp( + {"id": "call_slow", "name": "slow_tool", "arguments": {}}, + {"id": "call_fast", "name": "fast_tool", "arguments": {}}, + ), + _text_resp("Both done."), + ] + conv = AsyncConversation(driver=MockAsyncToolDriver(responses), tools=reg, sequential_tools=True) + result = await conv.ask("Call both") + + assert result == "Both done." + # Fully sequential: slow finishes before fast starts + assert started == ["slow_tool", "fast_tool"] + assert finished == ["slow_tool", "fast_tool"] + tool_msgs = _tool_messages(conv) + assert [m["tool_call_id"] for m in tool_msgs] == ["call_slow", "call_fast"] + + +# --------------------------------------------------------------------------- +# History trimming never orphans tool pairs (M11) +# --------------------------------------------------------------------------- + + +class TestHistoryTrimming: + def test_trim_never_leaves_leading_tool_message(self): + conv = Conversation(driver=MockToolDriver([_text_resp("hi")]), max_history_messages=3) + conv._messages = [ + {"role": "user", "content": "q"}, + { + "role": "assistant", + "content": "", + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "echo", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "result"}, + {"role": "user", "content": "q2"}, + {"role": "assistant", "content": "a2"}, + ] + conv._accumulate_usage({}) + + # Naive tail-slice would start with the orphaned `tool` message; + # boundary-aware trimming drops it. + assert len(conv.messages) == 2 + assert conv.messages[0]["role"] == "user" + assert all(m.get("role") != "tool" for m in conv.messages) + + +# --------------------------------------------------------------------------- +# Missing tool-call ids are generated before execution (L4) +# --------------------------------------------------------------------------- + + +class TestMissingToolCallId: + def test_empty_id_gets_generated(self): + responses = [ + _tool_call_resp({"id": "", "name": "echo", "arguments": {"text": "hi"}}), + _text_resp("Done."), + ] + conv = Conversation(driver=MockToolDriver(responses), tools=_echo_registry()) + result = conv.ask("Call echo") + + assert result == "Done." + tool_msg = _tool_messages(conv)[0] + assert tool_msg["tool_call_id"].startswith("call_") + assert len(tool_msg["tool_call_id"]) == len("call_") + 32 + # Assistant tool_calls entry matches the tool result id + assistant_msg = next(m for m in conv.messages if m.get("tool_calls")) + assert assistant_msg["tool_calls"][0]["id"] == tool_msg["tool_call_id"] + assert tool_msg["tool_call_id"] in conv._full_tool_results diff --git a/tests/test_simulated_tools.py b/tests/test_simulated_tools.py index f315a91f..69d7f8a9 100644 --- a/tests/test_simulated_tools.py +++ b/tests/test_simulated_tools.py @@ -324,16 +324,20 @@ def test_single_round(self, tools): assert result == "The weather in London is 22 celsius." assert driver._call_count == 2 - def test_max_rounds_exceeded(self, tools): - """Should raise RuntimeError when max rounds exceeded.""" + def test_max_rounds_graceful_final_answer(self, tools): + """Exhausting max rounds produces a graceful final answer instead of raising.""" from prompture.agents.conversation import Conversation - # Always returns tool calls, never a final answer - responses = [json.dumps({"type": "tool_call", "name": "get_weather", "arguments": {"city": "London"}})] * 5 + # Always returns tool calls during the loop; the final graceful + # call (tools removed) returns a final_answer. + responses = [json.dumps({"type": "tool_call", "name": "get_weather", "arguments": {"city": "London"}})] * 3 + responses.append(json.dumps({"type": "final_answer", "content": "Giving up gracefully."})) driver = MockDriver(responses) conv = Conversation(driver=driver, tools=tools, simulated_tools=True, max_tool_rounds=3) - with pytest.raises(RuntimeError, match="exceeded 3 rounds"): - conv.ask("What is the weather?") + result = conv.ask("What is the weather?") + assert result == "Giving up gracefully." + assert conv.max_rounds_reached is True + assert driver._call_count == 4 def test_tool_error_becomes_message(self, tools): """Tool execution errors are sent back as user messages.""" diff --git a/tests/test_tool_use.py b/tests/test_tool_use.py index 35302ccb..31ffd664 100644 --- a/tests/test_tool_use.py +++ b/tests/test_tool_use.py @@ -318,14 +318,15 @@ def generate_messages(self, messages, options): result = conv.ask("Hi") assert result == "Hello!" - def test_max_tool_rounds_exceeded(self): - """Raise RuntimeError when tool loop exceeds max rounds.""" + def test_max_tool_rounds_graceful_final_answer(self): + """Exhausting max_tool_rounds yields one final no-tools answer instead of raising.""" def noop() -> str: """Do nothing.""" return "ok" - # Every response has tool_calls + # Every loop response has tool_calls; the final graceful call (no + # tools) returns plain text. responses = [ { "text": "", @@ -333,8 +334,14 @@ def noop() -> str: "tool_calls": [{"id": f"call_{i}", "name": "noop", "arguments": {}}], "stop_reason": "tool_use", } - for i in range(5) + for i in range(3) ] + responses.append( + { + "text": "Final answer without tools.", + "meta": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2, "cost": 0.0}, + } + ) reg = ToolRegistry() reg.register(noop) @@ -342,8 +349,9 @@ def noop() -> str: driver = MockToolDriver(responses) conv = Conversation(driver=driver, tools=reg, max_tool_rounds=3) - with pytest.raises(RuntimeError, match="exceeded"): - conv.ask("Do something") + result = conv.ask("Do something") + assert result == "Final answer without tools." + assert conv.max_rounds_reached is True def test_register_tool_method(self): """Conversation.register_tool convenience method works.""" From cc9ed34fff680d7259b494a2954339b8bc30020e Mon Sep 17 00:00:00 2001 From: Juan Denis <13461850+jhd3197@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:18:10 -0400 Subject: [PATCH 6/8] feat(mcp): consume tools from MCP servers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Prompture had no MCP support, while every peer framework (OpenAI Agents SDK, Pydantic AI, LangChain) now consumes MCP servers. It maps cleanly onto the existing ToolRegistry interface — `list_tools()` becomes ToolDefinitions and `call_tool()` becomes the tool function — so MCP tools are called through the normal tool path with no special casing in the agent loop. `register_mcp_tools(registry, session)` registers every tool an initialized `mcp.ClientSession` advertises, with an optional `prefix` to avoid collisions across servers. `mcp_session_from_stdio()` is a convenience context manager for stdio servers, and `register_mcp_tools_sync()` covers sync `Agent` code outside an event loop. Results are flattened to strings: text blocks joined, non-text blocks (images, resources) reduced to a short placeholder, and MCP call errors returned as `Error: ...` so the model can react to them. Behind the optional `prompture[mcp]` extra, exported lazily so importing `prompture.integrations` never requires the dependency. Co-Authored-By: Claude Opus 5 (1M context) --- docs/INTEGRATIONS.md | 33 ++++ prompture/integrations/__init__.py | 19 +++ prompture/integrations/mcp_bridge.py | 224 ++++++++++++++++++++++++ pyproject.toml | 3 +- tests/test_mcp_bridge.py | 243 +++++++++++++++++++++++++++ 5 files changed, 521 insertions(+), 1 deletion(-) create mode 100644 prompture/integrations/mcp_bridge.py create mode 100644 tests/test_mcp_bridge.py diff --git a/docs/INTEGRATIONS.md b/docs/INTEGRATIONS.md index cf96462a..c31f0568 100644 --- a/docs/INTEGRATIONS.md +++ b/docs/INTEGRATIONS.md @@ -118,6 +118,39 @@ except ValidationError: pass ``` +### MCP (Model Context Protocol) Tools + +Prompture can consume tools from any MCP server. MCP tools are registered +into a `ToolRegistry` just like native tools or skills, so agents call them +through the normal tool-calling path. Requires the optional extra: + +```bash +pip install 'prompture[mcp]' +``` + +```python +import asyncio +from prompture import AsyncAgent, ToolRegistry +from prompture.integrations.mcp_bridge import mcp_session_from_stdio, register_mcp_tools + +async def main(): + registry = ToolRegistry() + async with mcp_session_from_stdio("npx", ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]) as session: + await register_mcp_tools(registry, session, prefix="fs") + agent = AsyncAgent("openai/gpt-4o", tools=registry) + result = await agent.run("List the files in /tmp") + print(result.output_text) + +asyncio.run(main()) +``` + +With an already-initialized `mcp.ClientSession` (any transport), use +`await register_mcp_tools(registry, session)` directly. For sync `Agent` +code outside an event loop, `register_mcp_tools_sync(registry, session)` +does the same thing. MCP tool results are flattened to strings: text blocks +are joined, non-text blocks (images, resources) become a short placeholder, +and MCP call errors are returned to the model as `Error: ...` strings. + --- ## Extending Prompture diff --git a/prompture/integrations/__init__.py b/prompture/integrations/__init__.py index bfe25ba5..1d001347 100644 --- a/prompture/integrations/__init__.py +++ b/prompture/integrations/__init__.py @@ -28,8 +28,27 @@ "discover_and_register_plugins", "filter_available_skills", "make_transform_chain", + "mcp_session_from_stdio", + "register_mcp_tools", + "register_mcp_tools_sync", "registry_to_skill_dict", "skill_to_tool_definition", "skills_to_registry", "tool_definition_to_skill", ] + +# MCP bridge is exported lazily so importing this package never requires +# the optional ``mcp`` dependency (pip install 'prompture[mcp]'). +_LAZY_MCP_EXPORTS = { + "register_mcp_tools", + "register_mcp_tools_sync", + "mcp_session_from_stdio", +} + + +def __getattr__(name: str): + if name in _LAZY_MCP_EXPORTS: + from . import mcp_bridge + + return getattr(mcp_bridge, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/prompture/integrations/mcp_bridge.py b/prompture/integrations/mcp_bridge.py new file mode 100644 index 00000000..ff04f82e --- /dev/null +++ b/prompture/integrations/mcp_bridge.py @@ -0,0 +1,224 @@ +"""Bridge between MCP (Model Context Protocol) servers and Prompture's tool registry. + +Discovers tools exposed by an MCP server over a ``mcp.ClientSession`` and +registers them as Prompture :class:`ToolDefinition` objects, so agents can +call MCP tools exactly like native Prompture tools or tukuy skills. + +The ``mcp`` package is an optional dependency — install it with:: + + pip install 'prompture[mcp]' + +All ``mcp`` imports are lazy to avoid import-time errors when the package +is not installed. +""" + +from __future__ import annotations + +import asyncio +import re +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from typing import TYPE_CHECKING, Any + +from ..agents.tools_schema import ToolDefinition, ToolRegistry + +if TYPE_CHECKING: # pragma: no cover - typing only, never imported at runtime + import mcp + +_MCP_IMPORT_MESSAGE = "The 'mcp' package is required for MCP integration. Install it with: pip install 'prompture[mcp]'" + +# Tool names must match this pattern to be accepted by LLM providers. +_NAME_PATTERN = re.compile(r"^[a-zA-Z0-9_-]{1,64}$") +_INVALID_NAME_CHARS = re.compile(r"[^a-zA-Z0-9_-]") + + +def _import_mcp() -> Any: + """Import the ``mcp`` package lazily, raising a clear error if missing.""" + try: + import mcp + except ImportError as exc: # pragma: no cover - depends on environment + raise ImportError(_MCP_IMPORT_MESSAGE) from exc + return mcp + + +def _sanitize_name(name: str, prefix: str | None = None) -> str: + """Sanitize an MCP tool name to ``^[a-zA-Z0-9_-]{1,64}$``. + + Invalid characters are replaced with ``_`` and the result is truncated + to 64 characters. When *prefix* is given, it is prepended as + ``{prefix}_{name}`` before sanitization. + """ + raw = f"{prefix}_{name}" if prefix else name + sanitized = _INVALID_NAME_CHARS.sub("_", raw)[:64] + if not sanitized: + sanitized = "mcp_tool" + if not _NAME_PATTERN.match(sanitized): # pragma: no cover - defensive + sanitized = "mcp_tool" + return sanitized + + +def _serialize_content_block(block: Any) -> str: + """Serialize a single MCP content block to a string. + + Text blocks contribute their text; non-text blocks (images, resources, + ...) are represented by a short placeholder since the tool-result path + is string-only today. + """ + text = getattr(block, "text", None) + if isinstance(text, str): + return text + if isinstance(block, str): + return block + block_type = getattr(block, "type", None) or type(block).__name__ + return f"[{block_type} content block]" + + +def _serialize_call_result(result: Any) -> str: + """Serialize an MCP ``CallToolResult`` to a single string for the tool-result path.""" + if isinstance(result, str): + return result + content = getattr(result, "content", None) + if content is None and isinstance(result, dict): + content = result.get("content") + if not content: + return "" + if not isinstance(content, (list, tuple)): + content = [content] + return "\n".join(_serialize_content_block(block) for block in content) + + +def _make_mcp_tool_function(session: mcp.ClientSession, mcp_name: str) -> Any: + """Build the sync/async callables that dispatch a tool call to the MCP session.""" + + async def _acall(**arguments: Any) -> str: + try: + result = await session.call_tool(mcp_name, arguments) + except Exception as exc: + # MCP call errors become LLM-friendly error strings (registry convention). + return f"Error calling MCP tool '{mcp_name}': {exc}" + text = _serialize_call_result(result) + if getattr(result, "isError", False): + return f"Error from MCP tool '{mcp_name}': {text}" + return text + + def _call(**arguments: Any) -> str: + try: + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run(_acall(**arguments)) + raise RuntimeError( + f"MCP tool '{mcp_name}' was called synchronously from within a running event loop. " + "Use AsyncAgent/AsyncConversation (which awaits the tool via its async path), " + "or call it from a thread without a running loop." + ) + + _call._async_fn = _acall # type: ignore[attr-defined] + return _call + + +async def register_mcp_tools( + registry: ToolRegistry, + session: mcp.ClientSession, + *, + prefix: str | None = None, +) -> list[str]: + """Register every tool exposed by an MCP session into *registry*. + + Calls ``session.list_tools()``, converts each MCP tool to a + :class:`ToolDefinition` (the MCP ``inputSchema`` is already JSON Schema + and is passed through unchanged, with names sanitized to + ``^[a-zA-Z0-9_-]{1,64}$`` and optionally prefixed), and registers it so + execution dispatches to ``session.call_tool(name, arguments)``. + + Each tool is registered with an ``_async_fn`` hook, so + :class:`AsyncAgent`/:class:`AsyncConversation` await it natively. The + sync execution path uses ``asyncio.run`` and raises a clear error when + called from within a running event loop. + + Args: + registry: The :class:`ToolRegistry` to register tools into. + session: An initialized ``mcp.ClientSession``. + prefix: Optional prefix prepended to every tool name as + ``{prefix}_{name}`` (useful to namespace several MCP servers). + + Returns: + The list of registered (sanitized) tool names. + """ + list_result = await session.list_tools() + tools = getattr(list_result, "tools", list_result) + + registered: list[str] = [] + for tool in tools: + mcp_name = getattr(tool, "name", None) + if not mcp_name: + continue + name = _sanitize_name(str(mcp_name), prefix) + description = getattr(tool, "description", None) or f"MCP tool {mcp_name}" + parameters = getattr(tool, "inputSchema", None) + if not isinstance(parameters, dict): + parameters = {"type": "object", "properties": {}} + registry.add( + ToolDefinition( + name=name, + description=str(description), + parameters=parameters, + function=_make_mcp_tool_function(session, str(mcp_name)), + ) + ) + registered.append(name) + return registered + + +def register_mcp_tools_sync( + registry: ToolRegistry, + session: mcp.ClientSession, + *, + prefix: str | None = None, +) -> list[str]: + """Synchronous wrapper around :func:`register_mcp_tools`. + + Uses ``asyncio.run`` when no event loop is running; raises a clear + error otherwise. + """ + try: + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run(register_mcp_tools(registry, session, prefix=prefix)) + raise RuntimeError( + "register_mcp_tools_sync() cannot be called from within a running event loop. " + "Use 'await register_mcp_tools(...)' instead." + ) + + +@asynccontextmanager +async def mcp_session_from_stdio( + command: str, + args: list[str] | None = None, + env: dict | None = None, +) -> AsyncIterator[mcp.ClientSession]: + """One-liner async context manager yielding an initialized MCP stdio session. + + Example:: + + async with mcp_session_from_stdio("npx", ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]) as session: + await register_mcp_tools(registry, session) + + Args: + command: The command used to start the MCP server process. + args: Optional command arguments. + env: Optional environment variables for the server process. + """ + mcp = _import_mcp() + from mcp.client.stdio import stdio_client + + server_params = mcp.StdioServerParameters(command=command, args=args or [], env=env) + async with stdio_client(server_params) as (read, write), mcp.ClientSession(read, write) as session: + await session.initialize() + yield session + + +__all__ = [ + "mcp_session_from_stdio", + "register_mcp_tools", + "register_mcp_tools_sync", +] diff --git a/pyproject.toml b/pyproject.toml index dc416428..75c641e0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,8 +34,9 @@ groq = ["groq>=0.4.0"] toon = ["python-toon>=0.1.0", "tukuy==0.0.30"] pandas = ["pandas>=1.3.0"] sandbox = ["tukuy>=0.0.30"] +mcp = ["mcp>=1.0"] all = [ - "prompture[openai,anthropic,google,groq,toon,pandas,sandbox]", + "prompture[openai,anthropic,google,groq,toon,pandas,sandbox,mcp]", ] test = ["pytest>=7.0", "pytest-asyncio>=0.23.0", "prompture[all]"] dev = ["pytest>=7.0", "pytest-asyncio>=0.23.0", "ruff>=0.8.0", "prompture[all]"] diff --git a/tests/test_mcp_bridge.py b/tests/test_mcp_bridge.py new file mode 100644 index 00000000..09aa0ad5 --- /dev/null +++ b/tests/test_mcp_bridge.py @@ -0,0 +1,243 @@ +"""Tests for the MCP (Model Context Protocol) bridge. + +All MCP objects (ClientSession, tools, content blocks) are mocked — no real +MCP server or ``mcp`` package dependency is required. +""" + +import sys +from types import SimpleNamespace + +import pytest + +from prompture.agents.tools_schema import ToolRegistry +from prompture.integrations import mcp_bridge +from prompture.integrations.mcp_bridge import ( + _sanitize_name, + register_mcp_tools, + register_mcp_tools_sync, +) + + +def _tool(name, description=None, input_schema=None): + return SimpleNamespace( + name=name, + description=description, + inputSchema=input_schema, + ) + + +def _text_block(text): + return SimpleNamespace(type="text", text=text) + + +def _image_block(): + return SimpleNamespace(type="image", data="...", mimeType="image/png") + + +class FakeSession: + """Mock of ``mcp.ClientSession`` with async list_tools/call_tool.""" + + def __init__(self, tools=None, call_result=None, call_error=None): + self._tools = tools or [] + self._call_result = call_result + self._call_error = call_error + self.calls = [] + + async def list_tools(self): + return SimpleNamespace(tools=self._tools) + + async def call_tool(self, name, arguments): + self.calls.append((name, arguments)) + if self._call_error is not None: + raise self._call_error + return self._call_result + + +WEATHER_SCHEMA = { + "type": "object", + "properties": {"city": {"type": "string", "description": "City name"}}, + "required": ["city"], +} + + +class TestRegistration: + async def test_registers_tools_with_schema_passthrough(self): + session = FakeSession(tools=[_tool("get_weather", "Get weather", WEATHER_SCHEMA)]) + registry = ToolRegistry() + + names = await register_mcp_tools(registry, session) + + assert names == ["get_weather"] + td = registry.get("get_weather") + assert td is not None + assert td.description == "Get weather" + assert td.parameters == WEATHER_SCHEMA # JSON Schema passed through unchanged + + async def test_prefix_applied_to_names(self): + session = FakeSession(tools=[_tool("read_file", "Read a file")]) + registry = ToolRegistry() + + names = await register_mcp_tools(registry, session, prefix="fs") + + assert names == ["fs_read_file"] + assert "fs_read_file" in registry + + async def test_missing_description_and_schema_get_defaults(self): + session = FakeSession(tools=[_tool("ping")]) + registry = ToolRegistry() + + await register_mcp_tools(registry, session) + + td = registry.get("ping") + assert td.description == "MCP tool ping" + assert td.parameters == {"type": "object", "properties": {}} + + async def test_empty_tool_list(self): + registry = ToolRegistry() + names = await register_mcp_tools(registry, FakeSession()) + assert names == [] + assert len(registry) == 0 + + +class TestNameSanitization: + def test_valid_name_unchanged(self): + assert _sanitize_name("get_weather-2") == "get_weather-2" + + def test_invalid_chars_replaced(self): + assert _sanitize_name("get weather.now!") == "get_weather_now_" + + def test_long_name_truncated_to_64(self): + name = _sanitize_name("a" * 100) + assert len(name) == 64 + + def test_prefix_counts_toward_limit(self): + name = _sanitize_name("b" * 100, prefix="srv") + assert len(name) == 64 + assert name.startswith("srv_") + + async def test_sanitization_applied_on_registration(self): + session = FakeSession(tools=[_tool("my tool/v2")]) + registry = ToolRegistry() + + names = await register_mcp_tools(registry, session) + + assert names == ["my_tool_v2"] + assert "my_tool_v2" in registry + + +class TestExecutionDispatch: + async def test_aexecute_dispatches_to_session_call_tool(self): + result = SimpleNamespace(content=[_text_block("sunny"), _text_block("25C")], isError=False) + session = FakeSession(tools=[_tool("get_weather", "Get weather", WEATHER_SCHEMA)], call_result=result) + registry = ToolRegistry() + await register_mcp_tools(registry, session, prefix="w") + + output = await registry.aexecute("w_get_weather", {"city": "Paris"}) + + # Original MCP name (not the sanitized/prefixed one) is used server-side. + assert session.calls == [("get_weather", {"city": "Paris"})] + assert output == "sunny\n25C" # text blocks joined + + async def test_async_fn_hook_attached(self): + session = FakeSession(tools=[_tool("ping")], call_result=SimpleNamespace(content=[], isError=False)) + registry = ToolRegistry() + await register_mcp_tools(registry, session) + + td = registry.get("ping") + assert getattr(td.function, "_async_fn", None) is not None + + def test_sync_execute_runs_outside_event_loop(self): + result = SimpleNamespace(content=[_text_block("pong")], isError=False) + session = FakeSession(tools=[_tool("ping")], call_result=result) + registry = ToolRegistry() + register_mcp_tools_sync(registry, session) + + assert registry.execute("ping", {}) == "pong" + assert session.calls == [("ping", {})] + + async def test_sync_execute_inside_event_loop_raises_clear_error(self): + session = FakeSession(tools=[_tool("ping")]) + registry = ToolRegistry() + await register_mcp_tools(registry, session) + + with pytest.raises(RuntimeError, match="running event loop"): + registry.execute("ping", {}) + + +class TestErrorMapping: + async def test_call_exception_becomes_error_string(self): + session = FakeSession( + tools=[_tool("boom")], + call_error=ConnectionError("server unreachable"), + ) + registry = ToolRegistry() + await register_mcp_tools(registry, session) + + output = await registry.aexecute("boom", {}) + + assert "Error calling MCP tool 'boom'" in output + assert "server unreachable" in output + + async def test_is_error_result_becomes_error_string(self): + result = SimpleNamespace(content=[_text_block("no such file")], isError=True) + session = FakeSession(tools=[_tool("read_file")], call_result=result) + registry = ToolRegistry() + await register_mcp_tools(registry, session) + + output = await registry.aexecute("read_file", {}) + + assert output == "Error from MCP tool 'read_file': no such file" + + +class TestContentSerialization: + async def test_non_text_blocks_become_placeholder(self): + result = SimpleNamespace( + content=[_text_block("here you go"), _image_block()], + isError=False, + ) + session = FakeSession(tools=[_tool("screenshot")], call_result=result) + registry = ToolRegistry() + await register_mcp_tools(registry, session) + + output = await registry.aexecute("screenshot", {}) + + assert "here you go" in output + assert "[image content block]" in output + assert "mimeType" not in output # raw block not leaked + + async def test_empty_content_returns_empty_string(self): + session = FakeSession(tools=[_tool("noop")], call_result=SimpleNamespace(content=[], isError=False)) + registry = ToolRegistry() + await register_mcp_tools(registry, session) + + assert await registry.aexecute("noop", {}) == "" + + +class TestSyncRegistrationWrapper: + def test_works_outside_event_loop(self): + session = FakeSession(tools=[_tool("ping")]) + registry = ToolRegistry() + + names = register_mcp_tools_sync(registry, session) + + assert names == ["ping"] + assert "ping" in registry + + async def test_raises_inside_running_loop(self): + registry = ToolRegistry() + with pytest.raises(RuntimeError, match="register_mcp_tools"): + register_mcp_tools_sync(registry, FakeSession()) + + +class TestImportGuard: + async def test_clear_error_when_mcp_not_installed(self, monkeypatch): + monkeypatch.setitem(sys.modules, "mcp", None) # makes `import mcp` raise ImportError + + with pytest.raises(ImportError, match=r"pip install 'prompture\[mcp\]'"): + async with mcp_bridge.mcp_session_from_stdio("some-server-command"): + pass + + def test_bridge_module_importable_without_mcp(self): + # The bridge module is already imported at the top of this file without + # the ``mcp`` package needing to be installed (lazy import). + assert mcp_bridge.register_mcp_tools is not None From b6ea30252eb5ab9c513145be8d4afa5b7781af8d Mon Sep 17 00:00:00 2001 From: Juan Denis <13461850+jhd3197@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:18:20 -0400 Subject: [PATCH 7/8] test(serve): skip SSE streaming tests without the serve extra Both streaming tests asserted a 200 but got 501, which is the endpoint correctly reporting that `sse-starlette` is unavailable rather than a defect. `sse-starlette` ships in the optional `serve` extra, so a valid install without it left the suite red. Skip via `importorskip`, matching the `importorskip("fastapi")` already at the top of both modules. The non-streaming tests need no extra and keep running. Co-Authored-By: Claude Opus 5 (1M context) --- tests/test_openai_server.py | 3 +++ tests/test_server.py | 4 ++++ 2 files changed, 7 insertions(+) diff --git a/tests/test_openai_server.py b/tests/test_openai_server.py index 036c348c..7826cf4c 100644 --- a/tests/test_openai_server.py +++ b/tests/test_openai_server.py @@ -252,6 +252,9 @@ def test_messages_without_user_role_rejected(self, client): class TestChatCompletionsStreaming: def test_stream_produces_sse_chunks(self, client): + # SSE streaming needs the `serve` extra; the endpoint answers 501 + # without it, which is correct behaviour rather than a failure. + pytest.importorskip("sse_starlette") with client.stream( "POST", "/v1/chat/completions", diff --git a/tests/test_server.py b/tests/test_server.py index ddfc892e..c9ed551f 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -257,6 +257,10 @@ async def _raise(*args, **kwargs): def test_run_endpoint_streams_sse_events(self, client): """When stream=true, the endpoint emits SSE-framed events.""" + # SSE streaming needs the `serve` extra; the endpoint answers 501 + # without it, which is correct behaviour rather than a failure. + pytest.importorskip("sse_starlette") + from prompture.infra.coding_agent_events import CodingAgentEvent async def _fake_stream(*args, **kwargs): From 5654ca484393c4a9e39ef6e39a79e9b7c3aa0ca7 Mon Sep 17 00:00:00 2001 From: Juan Denis <13461850+jhd3197@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:29:53 -0400 Subject: [PATCH 8/8] fix(tools): keep structured tool params expanded on Python 3.10/3.11 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI caught `test_typed_dict_expands` failing on Python 3.11: the schema had no `properties` at all. pydantic refuses `typing.TypedDict` below 3.12 (it requires the `typing_extensions` variant), so `TypeAdapter` raised, the except branch logged at debug level and fell through to the generic `{"type": "string"}` return. So on two of the four supported Python versions a TypedDict tool parameter was advertised to the model as a bare string — the exact silent degradation the richer type mapping was written to eliminate. It only looked fixed because local development runs 3.14, where pydantic accepts it. Rebuild the object schema from resolved type hints when TypeAdapter cannot, so field names and types survive on every supported version. Optionality is read per kind — `__required_keys__` for TypedDict, absent default/ default_factory for dataclasses, `is_required()` for models — and `required` is omitted entirely when it cannot be determined, so nothing is over-constrained. The fallback is exercised on all versions by forcing TypeAdapter to raise, rather than being covered only where pydantic happens to refuse; otherwise the path would stay untested on 3.12+ and regress unnoticed again. Co-Authored-By: Claude Opus 5 (1M context) --- prompture/agents/tools_schema.py | 68 ++++++++++++++++++++++-- tests/test_tools_schema.py | 89 ++++++++++++++++++++++++++++++++ 2 files changed, 153 insertions(+), 4 deletions(-) diff --git a/prompture/agents/tools_schema.py b/prompture/agents/tools_schema.py index 4ddddf20..5e91a708 100644 --- a/prompture/agents/tools_schema.py +++ b/prompture/agents/tools_schema.py @@ -27,7 +27,8 @@ def get_weather(city: str, units: str = "celsius") -> str: import re import uuid from collections.abc import Callable, Mapping -from dataclasses import asdict, dataclass, field, is_dataclass +from dataclasses import MISSING, asdict, dataclass, field, is_dataclass +from dataclasses import fields as dataclass_fields from datetime import date, datetime, time from enum import Enum from typing import Any, Literal, get_args, get_origin, get_type_hints, is_typeddict @@ -98,6 +99,60 @@ def _is_structured_type(annotation: Any) -> bool: return issubclass(annotation, BaseModel) +def _structured_required_fields(annotation: Any, hints: dict[str, Any]) -> list[str] | None: + """Required field names of a structured type, or ``None`` if undeterminable. + + ``None`` (rather than "all of them") keeps us from over-constraining a + type whose optionality we cannot read. + """ + if is_typeddict(annotation): + required = getattr(annotation, "__required_keys__", None) + if required is None: + return None + return [name for name in hints if name in required] + if is_dataclass(annotation): + return [ + f.name + for f in dataclass_fields(annotation) + if f.default is MISSING and f.default_factory is MISSING # type: ignore[misc] + ] + if isinstance(annotation, type) and issubclass(annotation, BaseModel): + return [name for name, f in annotation.model_fields.items() if f.is_required()] + return None + + +def _structured_fallback_schema(annotation: Any) -> dict[str, Any] | None: + """Object schema for a structured type :class:`TypeAdapter` cannot handle. + + pydantic refuses ``typing.TypedDict`` on Python < 3.12 (it requires the + ``typing_extensions`` variant), which on two of the four supported Python + versions left such a parameter advertised as a bare string — the exact + silent degradation the richer type mapping exists to prevent. Rebuild the + object schema from resolved hints so the model still sees the real field + names and types. + + Returns ``None`` when the type exposes nothing introspectable, leaving the + caller's own fallback in charge. + """ + try: + hints = get_type_hints(annotation) + except Exception: + logger.debug("Could not resolve type hints for structured type %r", annotation, exc_info=True) + return None + hints.pop("return", None) + if not hints: + return None + + schema: dict[str, Any] = { + "type": "object", + "properties": {name: _python_type_to_json_schema(hint) for name, hint in hints.items()}, + } + required = _structured_required_fields(annotation, hints) + if required: + schema["required"] = required + return schema + + def _python_type_to_json_schema(annotation: Any) -> dict[str, Any]: """Convert a Python type annotation to a JSON Schema snippet. @@ -106,8 +161,10 @@ def _python_type_to_json_schema(annotation: Any) -> dict[str, Any]: ``enum``, ``datetime``/``date``/``time``/``UUID`` as formatted strings, ``list``/``tuple``/``dict`` containers (``dict[str, X]`` gets ``additionalProperties``), nested pydantic models, dataclasses and - ``TypedDict`` via :class:`pydantic.TypeAdapter`, and ``Any`` as the empty - schema. Unknown types fall back to ``{"type": "string"}``. + ``TypedDict`` via :class:`pydantic.TypeAdapter` (falling back to + :func:`_structured_fallback_schema` on the Python versions where pydantic + refuses a given structured type), and ``Any`` as the empty schema. Unknown + types fall back to ``{"type": "string"}``. A *missing* annotation yields the empty (unconstrained) schema rather than a string: an unannotated parameter means "type unknown", and since @@ -176,8 +233,11 @@ def _python_type_to_json_schema(annotation: Any) -> dict[str, Any]: if _is_structured_type(annotation): try: return TypeAdapter(annotation).json_schema() - except Exception as exc: # pragma: no cover - defensive + except Exception as exc: logger.debug("TypeAdapter schema generation failed for %r: %s", annotation, exc) + manual = _structured_fallback_schema(annotation) + if manual is not None: + return manual # Simple types json_type = _TYPE_MAP.get(annotation, "string") diff --git a/tests/test_tools_schema.py b/tests/test_tools_schema.py index d99d4052..e4801ea4 100644 --- a/tests/test_tools_schema.py +++ b/tests/test_tools_schema.py @@ -449,3 +449,92 @@ def configure(mode: Literal["fast", "slow"]) -> str: td: ToolDefinition = tool_from_function(configure) wire = td.to_openai_format() assert wire["function"]["parameters"]["properties"]["mode"]["enum"] == ["fast", "slow"] + + +# --------------------------------------------------------------------------- +# Structured types when pydantic's TypeAdapter refuses them +# +# pydantic rejects `typing.TypedDict` on Python < 3.12 (it wants the +# typing_extensions variant), so on two of the four supported Python versions +# TypeAdapter raises and the parameter used to be advertised as a bare string. +# The failure is forced here so the fallback is covered on every version. +# --------------------------------------------------------------------------- + + +class PartialMovie(TypedDict, total=False): + title: str + year: int + + +@dataclass +class BoxWithDefaults: + w: int + h: int = 3 + + +class PointWithOptional(BaseModel): + x: int + y: int | None = None + + +@pytest.fixture +def type_adapter_unavailable(monkeypatch): + """Make TypeAdapter raise, as it does for typing.TypedDict on Py < 3.12.""" + + def _boom(*args, **kwargs): + raise RuntimeError("simulated: unsupported on this Python version") + + monkeypatch.setattr("prompture.agents.tools_schema.TypeAdapter", _boom) + + +class TestStructuredTypeFallback: + def test_typed_dict_still_expands(self, type_adapter_unavailable): + schema = _python_type_to_json_schema(Movie) + assert schema["type"] == "object" + assert schema["properties"]["year"] == {"type": "integer"} + assert schema["properties"]["title"] == {"type": "string"} + assert sorted(schema["required"]) == ["title", "year"] + + def test_never_degrades_to_a_bare_string(self, type_adapter_unavailable): + """The regression this guards: a whole object described as a string.""" + assert _python_type_to_json_schema(Movie) != {"type": "string"} + + def test_total_false_typed_dict_has_no_required_keys(self, type_adapter_unavailable): + schema = _python_type_to_json_schema(PartialMovie) + assert set(schema["properties"]) == {"title", "year"} + assert "required" not in schema + + def test_dataclass_defaults_are_not_required(self, type_adapter_unavailable): + schema = _python_type_to_json_schema(BoxWithDefaults) + assert schema["required"] == ["w"] + assert schema["properties"]["h"] == {"type": "integer"} + + def test_model_optional_fields_are_not_required(self, type_adapter_unavailable): + schema = _python_type_to_json_schema(PointWithOptional) + assert schema["required"] == ["x"] + assert schema["properties"]["y"] == {"anyOf": [{"type": "integer"}, {"type": "null"}]} + + def test_nested_annotations_are_resolved_recursively(self, type_adapter_unavailable): + @dataclass + class Outer: + tags: list[str] + where: Movie + + schema = _python_type_to_json_schema(Outer) + assert schema["properties"]["tags"] == {"type": "array", "items": {"type": "string"}} + assert schema["properties"]["where"]["properties"]["year"] == {"type": "integer"} + + def test_tool_parameter_keeps_its_object_shape(self, type_adapter_unavailable): + def rate(movie: Movie) -> str: + """Rate a movie. + + Args: + movie: The movie to rate. + """ + return movie["title"] + + prop = tool_from_function(rate).parameters["properties"]["movie"] + assert prop["type"] == "object" + assert set(prop["properties"]) == {"title", "year"} + # The docstring description survives the fallback. + assert prop["description"] == "The movie to rate."