#!/usr/bin/env python3
"""
Tool-related utilities for Klea
File: klea_utils/tools.py
Copyright 2026 Ankur Sinha
Author: Ankur Sinha <sanjay DOT ankur AT gmail DOT com>
"""
import json
import logging
import re
from collections.abc import Sequence
from typing import Any
from fastmcp.client.client import CallToolResult
from mcp.types import EmbeddedResource, TextContent, Tool
from klea_utils.tool_args import is_path_argument_name, is_uri_argument_name
logger = logging.getLogger(__name__)
def _textualize_content_block(block: Any) -> str:
"""Extract text from a single ContentBlock."""
if isinstance(block, TextContent):
return block.text
elif isinstance(block, EmbeddedResource):
resource = block.resource
if hasattr(resource, "blob") and resource.blob:
logger.warning("Blob resource not processed: %s", resource.uri)
text = getattr(resource, "text", None) or getattr(resource, "blob", "")
return f"[Resource: {resource.uri}]\n{text}"
else:
logger.warning("Unhandled content block type: %s", type(block).__name__)
return str(block)
[docs]
def textualize_tool_results(
tool_results: list[CallToolResult],
max_len_per_tool: int | None = None,
include_header: bool = True,
) -> str:
"""Format tool call results as LLM-ready text for use in prompt context.
Tool return values are typically JSON-structured data (dicts, lists) which
LLMs handle naturally. We wrap them in markdown code blocks for clear
separation from surrounding prompt text.
:param tool_results: List of tool call results
:param max_len_per_tool: If set, truncate each tool's textual content to
this many characters (per-tool cap, not total). Ensures a large first
result doesn't starve later tools' outputs.
:param include_header: When ``True`` (default) emit the batch header
(``## Tool Results``) and a ``### Result i/n`` label per result. Set
``False`` to render the bare content of a single result, e.g. when a
caller supplies its own per-result heading.
:returns: Formatted string suitable for inclusion in an LLM prompt
"""
if not tool_results:
return ""
text = "## Tool Results\n" if include_header else ""
for i, result in enumerate(tool_results, 1):
if include_header:
text += f"\n### Result {i}/{len(tool_results)}\n"
parts = [_textualize_content_block(c) for c in result.content]
content = "\n".join(parts)
if max_len_per_tool is not None and len(content) > max_len_per_tool:
content = content[:max_len_per_tool] + f"\n[tool result {i} truncated]"
if result.is_error:
text += "**Tool call failed:** " + content + "\n"
else:
text += "**Tool call succeeded:** " + content + "\n"
return text
[docs]
def last_tool_error_text(
tool_results: list[CallToolResult] | None,
max_len: int = 500,
) -> str:
"""Return the text of the last error result in a batch, or ``""``.
Used to give the tool picker concrete feedback when a round failed (for
example an ``edit_file`` ambiguity), so it can adjust the call instead of
repeating it. Only the last error is returned: the picker retries one
round at a time, so the most recent failure is the one to correct.
:param tool_results: The batch's tool results (``None`` treated as empty).
:param max_len: Truncate the extracted text to this many characters.
:returns: The final error result's textual content, or ``""`` if none.
"""
if not tool_results:
return ""
for result in reversed(tool_results):
if not getattr(result, "is_error", False):
continue
parts = [_textualize_content_block(c) for c in result.content]
text = "\n".join(parts).strip()
if max_len is not None and len(text) > max_len:
text = text[:max_len] + " [truncated]"
return text
return ""
#: Additional argument names that are useful in a failure digest but are not
#: necessarily filesystem paths (for example a transfer source or an output
#: destination). Arbitrary free-text arguments are deliberately not included.
_DIGEST_EXTRA_ARG_NAMES = frozenset({"source", "target", "destination", "output"})
def _is_digest_argument_name(name: str) -> bool:
"""Return whether *name* is safe/compact enough for a failure digest."""
return (
is_path_argument_name(name)
or is_uri_argument_name(name)
or name.lower() in _DIGEST_EXTRA_ARG_NAMES
)
_URI_USERINFO_RE = re.compile(r"(?i)(https?://)[^/@\s]+@")
_URI_SECRET_QUERY_RE = re.compile(
r"(?i)([?&](?:token|api[_-]?key|key|secret|password|signature|sig|auth)=)[^&\s]+"
)
def _redact_uri_secrets(text: str) -> str:
"""Redact credentials and common secret query parameters in a URI."""
text = _URI_USERINFO_RE.sub(r"\1***@", text)
return _URI_SECRET_QUERY_RE.sub(r"\1***", text)
def _format_digest_value(value: Any, max_chars: int) -> str:
"""Return a compact, redacted, truncated representation of an arg value."""
if isinstance(value, (list, tuple, set)):
text = ", ".join(str(item) for item in value)
else:
text = str(value)
text = " ".join(_redact_uri_secrets(text).split())
if max_chars is not None and len(text) > max_chars:
text = text[: max(0, max_chars - 3)].rstrip() + "..."
return text
def _collapse_schema_type(schema: dict[str, Any]) -> str:
"""Collapse a JSON Schema type spec into a single compact type name.
Handles ``anyOf`` unions (e.g. ``[{"type": "integer"}, {"type": "null"}]``
collapses to ``integer``) and plain ``type`` strings.
"""
if "anyOf" in schema:
types = [
t.get("type")
for t in schema["anyOf"]
if isinstance(t, dict) and t.get("type") != "null"
]
if types:
return "|".join(str(t) for t in types)
return schema.get("type") or "any"
def _format_tool_parameters(input_schema: dict[str, Any] | None) -> str:
"""Return a compact one-line-per-parameter summary of an MCP tool schema.
The raw JSON Schema ``properties`` dict is verbose (length and range
validators, ``anyOf`` unions); only the parameter name, type, whether it
is required or optional, its default, and the description are useful for
argument binding, so the rest is dropped and whitespace is normalised.
Optionality and defaults are stated explicitly rather than left implicit:
a model should not have to infer that a parameter without ``required`` is
optional, and it cannot see that an omitted filter means "all" unless the
default is shown (for example ``pattern (string, optional, default "*")``).
A JSON ``null`` default is omitted: it means "unset" for these tools and
rendering it as ``default null`` invites the model to pass a literal null.
"""
if not input_schema:
return ""
properties = input_schema.get("properties") or {}
if not properties:
return ""
required = set(input_schema.get("required") or [])
lines = []
for name, schema in properties.items():
if not isinstance(schema, dict):
continue
ptype = _collapse_schema_type(schema)
desc = " ".join((schema.get("description") or "").split())
flag = "required" if name in required else "optional"
default = schema.get("default")
default_part = (
f", default {json.dumps(default, ensure_ascii=False)}"
if "default" in schema and default is not None
else ""
)
lines.append(f"- {name} ({ptype}, {flag}{default_part}): {desc}".rstrip())
if not lines:
return ""
return "Parameters:\n" + "\n".join(lines)