monkeybot 2.1.1__py3-none-any.whl

monkeybot/core/context/memory_prompt.py

@@ -0,0 +1,222 @@
+"""Memory index selection for the system prompt: recent window, curator when token-heavy."""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+from dataclasses import dataclass
+
+from monkeybot.core.context import TurnContext
+from monkeybot.core.context.curator import (
+    CuratedPromptParts,
+    _env_int,
+    curation_enabled_from_env,
+    curator_model_id,
+    memory_index_token_estimate,
+    run_context_curator,
+)
+from monkeybot.core.llm.provider import Provider
+from monkeybot.core.logging_utils import kv
+from monkeybot.core.memory.index_format import memory_window_slice
+
+_log = logging.getLogger(__name__)
+
+# Process-local: thread_id -> (cache_key, curated_lines). Evicted on session teardown.
+# cache_key covers index content + user message so query-aware picks are not reused.
+_curation_cache: dict[str, tuple[str, list[str]]] = {}
+
+
+def reset_curation_cache_for_tests() -> None:
+    _curation_cache.clear()
+
+
+def evict_curation_cache(thread_id: str) -> None:
+    """Drop the cached curator selection for a thread (call on session removal)."""
+    _curation_cache.pop(thread_id, None)
+
+
+@dataclass(frozen=True)
+class MemoryPromptSelection:
+    """Memory lines and structural coverage for volatile system-prompt injection."""
+
+    lines: list[str]
+    total_lines: int
+    coverage: float
+    confidence: float
+    nudge_search: bool
+    use_custom_lines: bool
+
+
+def memory_window_lines_from_env() -> int:
+    return max(1, _env_int("CONTEXT_CURATION_MEMORY_WINDOW_LINES", 12))
+
+
+def memory_token_threshold_from_env() -> int:
+    return max(1, _env_int("CONTEXT_CURATION_MEMORY_TOKEN_THRESHOLD", 2000))
+
+
+def memory_index_fingerprint(lines: list[str]) -> str:
+    """Stable hash of index lines; used to detect mid-turn INDEX refreshes."""
+    payload = "\n".join(lines)
+    return hashlib.sha256(payload.encode("utf-8")).hexdigest()
+
+
+def _curator_cache_key(lines: list[str], user_message: str) -> str:
+    """Fingerprint for curator cache: index + query (curator is query-aware)."""
+    payload = "\n".join(lines) + "\0" + user_message
+    return hashlib.sha256(payload.encode("utf-8")).hexdigest()
+
+
+def memory_coverage(injected_count: int, total_count: int) -> float:
+    if total_count <= 0:
+        return 1.0
+    return min(1.0, injected_count / total_count)
+
+
+def memory_confidence(*, coverage: float, truncated: bool) -> float:
+    """Structural confidence (not LLM-reported). Full index → 1.0; else equals coverage."""
+    if not truncated:
+        return 1.0
+    return round(coverage, 4)
+
+
+def _selection_from_lines(
+    injected: list[str],
+    total: list[str],
+    *,
+    use_custom_lines: bool,
+) -> MemoryPromptSelection:
+    total_n = len(total)
+    injected_n = len(injected)
+    truncated = injected_n < total_n
+    coverage = memory_coverage(injected_n, total_n)
+    confidence = memory_confidence(coverage=coverage, truncated=truncated)
+    return MemoryPromptSelection(
+        lines=list(injected),
+        total_lines=total_n,
+        coverage=coverage,
+        confidence=confidence,
+        nudge_search=truncated,
+        use_custom_lines=use_custom_lines,
+    )
+
+
+def _cached_curator_lines(thread_id: str, cache_key: str) -> list[str] | None:
+    cached = _curation_cache.get(thread_id)
+    if cached is None:
+        return None
+    cached_key, lines = cached
+    if cached_key != cache_key:
+        return None
+    return list(lines)
+
+
+def _store_curator_cache(thread_id: str, cache_key: str, lines: list[str]) -> None:
+    _curation_cache[thread_id] = (cache_key, list(lines))
+
+
+async def _run_curator_cached(
+    *,
+    thread_id: str,
+    cache_key: str,
+    ctx: TurnContext,
+    provider: Provider,
+    curator_provider: Provider | None,
+    user_message: str,
+    max_memory_lines: int,
+) -> CuratedPromptParts:
+    cached = _cached_curator_lines(thread_id, cache_key)
+    if cached is not None:
+        _log.debug(
+            "curation cache hit %s",
+            kv(thread_id=thread_id, lines=len(cached)),
+        )
+        return CuratedPromptParts(cached, success=True)
+
+    _log.debug(
+        "curation cache miss %s",
+        kv(thread_id=thread_id, index_lines=len(ctx.memory_index)),
+    )
+    parts = await run_context_curator(
+        ctx=ctx,
+        provider=provider,
+        curator_model=curator_model_id(ctx),
+        user_message=user_message,
+        max_memory_lines=max_memory_lines,
+        curator_provider=curator_provider,
+    )
+    if parts.success:
+        _store_curator_cache(thread_id, cache_key, parts.memory_lines)
+    return parts
+
+
+async def prepare_memory_for_prompt(
+    *,
+    ctx: TurnContext,
+    user_message: str,
+    provider: Provider,
+    curator_provider: Provider | None,
+) -> MemoryPromptSelection:
+    """Select memory lines for the system prompt.
+
+    Default path: recent window. When the full index is token-heavy, optionally
+    call the LLM curator; on curator failure, fall back to the window.
+    """
+    total = list(ctx.memory_index)
+    if not curation_enabled_from_env() or not ctx.context_curation_enabled:
+        return _selection_from_lines(total, total, use_custom_lines=False)
+
+    window_n = memory_window_lines_from_env()
+    token_n = memory_token_threshold_from_env()
+    tokens = memory_index_token_estimate(total)
+    exceeds_window = len(total) > window_n
+    token_heavy = tokens > token_n
+
+    if not exceeds_window and not token_heavy:
+        return _selection_from_lines(total, total, use_custom_lines=False)
+
+    window = memory_window_slice(total, window_n)
+    if not token_heavy:
+        _log.debug(
+            "curation window %s",
+            kv(
+                thread_id=ctx.thread_id,
+                injected=len(window),
+                total=len(total),
+                tokens=tokens,
+            ),
+        )
+        return _selection_from_lines(window, total, use_custom_lines=len(window) < len(total))
+
+    parts = await _run_curator_cached(
+        thread_id=ctx.thread_id,
+        cache_key=_curator_cache_key(total, user_message),
+        ctx=ctx,
+        provider=provider,
+        curator_provider=curator_provider,
+        user_message=user_message,
+        max_memory_lines=window_n,
+    )
+    if parts.success:
+        injected = list(parts.memory_lines)
+        _log.info(
+            "curation curator %s",
+            kv(
+                thread_id=ctx.thread_id,
+                injected=len(injected),
+                total=len(total),
+                tokens=tokens,
+            ),
+        )
+        return _selection_from_lines(injected, total, use_custom_lines=True)
+
+    _log.warning(
+        "curation curator failed; using window %s",
+        kv(
+            thread_id=ctx.thread_id,
+            window=len(window),
+            total=len(total),
+            tokens=tokens,
+        ),
+    )
+    return _selection_from_lines(window, total, use_custom_lines=len(window) < len(total))

monkeybot/core/context/tool_output_policy.py

@@ -0,0 +1,270 @@
+"""Per-tool output shaping budgets (command_allowlist.yaml ``tool_output`` section)."""
+
+from __future__ import annotations
+
+import os
+import re
+from collections.abc import Iterable
+from dataclasses import dataclass
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Literal
+
+import yaml
+
+from monkeybot.core.tools.inspector import CommandTierConfigError
+
+ContentTypeHint = Literal["json", "logs", "code", "prose", "auto"]
+_SUPPORTED_CONTENT_TYPES = frozenset({"json", "logs", "code", "prose", "auto"})
+_MIN_OUTPUT_LINES = 5
+
+
+@dataclass(frozen=True)
+class ToolOutputBudget:
+    """Operator-configured caps and shaper hints for one tool name."""
+
+    content_type: ContentTypeHint | None = None
+    max_output_lines: int | None = None
+    max_array_items: int | None = None
+    keep_patterns: tuple[str, ...] = ()
+    collapse_repeated: bool = False
+
+
+# Sensible defaults when YAML omits ``tool_output`` entries.
+_MCP_DEFAULT_TOOL_BUDGET = ToolOutputBudget(
+    content_type="auto",
+    max_output_lines=120,
+    max_array_items=40,
+)
+
+_known_mcp_tools: frozenset[str] = frozenset()
+
+
+_BUILTIN_TOOL_BUDGETS: dict[str, ToolOutputBudget] = {
+    "run_command": ToolOutputBudget(
+        content_type="logs",
+        max_output_lines=400,
+        collapse_repeated=True,
+        keep_patterns=(
+            r"(?i)\berror\b",
+            r"(?i)\bfatal\b",
+            r"(?i)traceback",
+            r"(?i)\bexception\b",
+        ),
+    ),
+    "web_search": ToolOutputBudget(
+        content_type="json",
+        max_array_items=30,
+    ),
+}
+
+
+def _resolve_policy_path(path: Path | None) -> Path | None:
+    if path is not None:
+        return path if path.is_file() else None
+    raw = os.environ.get("COMMAND_ALLOWLIST_CONFIG", "").strip()
+    if not raw:
+        return None
+    p = Path(raw).expanduser()
+    return p if p.is_file() else None
+
+
+def _parse_tool_output_entry(path: Path, tool_name: str, raw: Any) -> ToolOutputBudget:
+    if not isinstance(raw, dict):
+        raise CommandTierConfigError(
+            path, f"tool_output.{tool_name} must be a mapping"
+        )
+    content_type: ContentTypeHint | None = None
+    ct_raw = raw.get("content_type")
+    if ct_raw is not None:
+        if not isinstance(ct_raw, str) or ct_raw.strip().lower() not in _SUPPORTED_CONTENT_TYPES:
+            raise CommandTierConfigError(
+                path,
+                f"tool_output.{tool_name}.content_type must be one of "
+                f"{sorted(_SUPPORTED_CONTENT_TYPES)}",
+            )
+        content_type = ct_raw.strip().lower()  # type: ignore[assignment]
+
+    max_output_lines: int | None = None
+    if "max_output_lines" in raw:
+        val = raw["max_output_lines"]
+        if not isinstance(val, int) or val < 1:
+            raise CommandTierConfigError(
+                path, f"tool_output.{tool_name}.max_output_lines must be a positive integer"
+            )
+        max_output_lines = val
+
+    max_array_items: int | None = None
+    if "max_array_items" in raw:
+        val = raw["max_array_items"]
+        if not isinstance(val, int) or val < 1:
+            raise CommandTierConfigError(
+                path, f"tool_output.{tool_name}.max_array_items must be a positive integer"
+            )
+        max_array_items = val
+
+    keep_patterns: tuple[str, ...] = ()
+    kp_raw = raw.get("keep_patterns")
+    if kp_raw is not None:
+        if not isinstance(kp_raw, list):
+            raise CommandTierConfigError(
+                path, f"tool_output.{tool_name}.keep_patterns must be a list of regex strings"
+            )
+        compiled: list[str] = []
+        for i, item in enumerate(kp_raw):
+            if not isinstance(item, str) or not item.strip():
+                raise CommandTierConfigError(
+                    path, f"tool_output.{tool_name}.keep_patterns[{i}] must be a non-empty string"
+                )
+            try:
+                re.compile(item)
+            except re.error as exc:
+                raise CommandTierConfigError(
+                    path,
+                    f"tool_output.{tool_name}.keep_patterns[{i}]: invalid regex: {exc}",
+                ) from exc
+            compiled.append(item)
+        keep_patterns = tuple(compiled)
+
+    collapse_repeated = bool(raw.get("collapse_repeated", False))
+    return ToolOutputBudget(
+        content_type=content_type,
+        max_output_lines=max_output_lines,
+        max_array_items=max_array_items,
+        keep_patterns=keep_patterns,
+        collapse_repeated=collapse_repeated,
+    )
+
+
+def parse_tool_output_section(path: Path, data: dict[str, Any]) -> dict[str, ToolOutputBudget]:
+    """Parse optional ``tool_output`` mapping from command allowlist YAML."""
+    raw = data.get("tool_output")
+    if raw is None:
+        return dict(_BUILTIN_TOOL_BUDGETS)
+    if not isinstance(raw, dict):
+        raise CommandTierConfigError(path, "'tool_output' must be a mapping of tool names")
+
+    merged = dict(_BUILTIN_TOOL_BUDGETS)
+    for tool_name, entry in raw.items():
+        if not isinstance(tool_name, str) or not tool_name.strip():
+            raise CommandTierConfigError(path, "tool_output keys must be non-empty tool names")
+        parsed = _parse_tool_output_entry(path, tool_name.strip(), entry)
+        base = merged.get(tool_name.strip())
+        if base is not None:
+            merged[tool_name.strip()] = ToolOutputBudget(
+                content_type=(
+                    parsed.content_type if "content_type" in entry else base.content_type
+                ),
+                max_output_lines=(
+                    parsed.max_output_lines
+                    if "max_output_lines" in entry
+                    else base.max_output_lines
+                ),
+                max_array_items=(
+                    parsed.max_array_items if "max_array_items" in entry else base.max_array_items
+                ),
+                keep_patterns=(
+                    parsed.keep_patterns if "keep_patterns" in entry else base.keep_patterns
+                ),
+                collapse_repeated=(
+                    parsed.collapse_repeated
+                    if "collapse_repeated" in entry
+                    else base.collapse_repeated
+                ),
+            )
+        else:
+            merged[tool_name.strip()] = parsed
+    return merged
+
+
+def load_tool_output_policies(path: Path | None = None) -> dict[str, ToolOutputBudget]:
+    """Load merged built-in + YAML per-tool output budgets."""
+    resolved = _resolve_policy_path(path)
+    if resolved is None:
+        return dict(_BUILTIN_TOOL_BUDGETS)
+    try:
+        data = yaml.safe_load(resolved.read_bytes())
+    except yaml.YAMLError as exc:
+        raise CommandTierConfigError(resolved, f"invalid YAML: {exc}") from exc
+    if data is None:
+        return dict(_BUILTIN_TOOL_BUDGETS)
+    if not isinstance(data, dict):
+        raise CommandTierConfigError(resolved, "root must be a mapping")
+    return parse_tool_output_section(resolved, data)
+
+
+@lru_cache(maxsize=1)
+def cached_tool_output_policies() -> dict[str, ToolOutputBudget]:
+    """Process-wide cache of tool output policies (invalidated only on restart)."""
+    return load_tool_output_policies()
+
+
+def register_mcp_tool_names(names: Iterable[str]) -> None:
+    """Record prefixed MCP tool names for output-budget resolution (updated on MCP connect)."""
+    global _known_mcp_tools
+    added = frozenset(n for n in names if n)
+    if not added:
+        return
+    _known_mcp_tools = _known_mcp_tools | added
+
+
+def unregister_mcp_server_tools(server_name: str) -> None:
+    """Drop every ``server__*`` tool registered for ``server_name`` (MCP disconnect)."""
+    global _known_mcp_tools
+    if not server_name:
+        return
+    prefix = f"{server_name}__"
+    _known_mcp_tools = frozenset(n for n in _known_mcp_tools if not n.startswith(prefix))
+
+
+def reset_mcp_tool_registry_for_tests() -> None:
+    """Clear the MCP tool registry (tests only)."""
+    global _known_mcp_tools
+    _known_mcp_tools = frozenset()
+
+
+def resolve_tool_budget(tool_name: str) -> ToolOutputBudget | None:
+    """Return configured budget for ``tool_name``, or MCP default for registered MCP tools."""
+    merged = cached_tool_output_policies()
+    if tool_name in merged:
+        return merged[tool_name]
+    if tool_name in _known_mcp_tools:
+        return merged.get("*") or _MCP_DEFAULT_TOOL_BUDGET
+    return None
+
+
+def validate_tool_output_budgets(
+    policies: dict[str, ToolOutputBudget],
+) -> list[str]:
+    """Return human-readable warnings for implausibly tight operator caps."""
+    warnings: list[str] = []
+    for name, budget in policies.items():
+        if budget.max_output_lines is not None and budget.max_output_lines < _MIN_OUTPUT_LINES:
+            warnings.append(
+                f"tool_output.{name}.max_output_lines={budget.max_output_lines} is very small "
+                f"(<{_MIN_OUTPUT_LINES}); tool output may lose critical signal"
+            )
+        if budget.max_array_items is not None and budget.max_array_items < 3:
+            warnings.append(
+                f"tool_output.{name}.max_array_items={budget.max_array_items} is very small; "
+                "search/API results may be unusable"
+            )
+    return warnings
+
+
+def reset_tool_output_policy_cache_for_tests() -> None:
+    cached_tool_output_policies.cache_clear()
+    reset_mcp_tool_registry_for_tests()
+
+
+__all__ = [
+    "ToolOutputBudget",
+    "load_tool_output_policies",
+    "parse_tool_output_section",
+    "register_mcp_tool_names",
+    "resolve_tool_budget",
+    "reset_mcp_tool_registry_for_tests",
+    "reset_tool_output_policy_cache_for_tests",
+    "unregister_mcp_server_tools",
+    "validate_tool_output_budgets",
+]