millforge 0.1.0__py3-none-any.whl

millforge/_forge/guardrails/error_tracker.py

@@ -0,0 +1,71 @@
+"""Error budget tracking — consecutive retries and tool errors."""
+
+from __future__ import annotations
+
+
+class ErrorTracker:
+    """Tracks consecutive retry and tool error counts against limits.
+
+    Stateful — instantiate per session/task.
+
+    Args:
+        max_retries: Consecutive formatting/validation failures before
+            exhaustion.
+        max_tool_errors: Consecutive tool execution errors before
+            exhaustion. Soft errors (ToolResolutionError equivalent)
+            do not count.
+    """
+
+    def __init__(self, max_retries: int = 3, max_tool_errors: int = 2) -> None:
+        self.max_retries = max_retries
+        self.max_tool_errors = max_tool_errors
+        self._consecutive_retries = 0
+        self._consecutive_tool_errors = 0
+
+    def record_retry(self) -> None:
+        """Record a validation failure (TextResponse or unknown tool)."""
+        self._consecutive_retries += 1
+
+    def reset_retries(self) -> None:
+        """Reset retry counter (call on successful validation)."""
+        self._consecutive_retries = 0
+
+    def record_result(self, success: bool, is_soft_error: bool = False) -> None:
+        """Record a tool execution result.
+
+        Args:
+            success: True if the tool executed without error.
+            is_soft_error: True if the error is a resolution/soft error
+                that should not count toward the error budget (e.g.,
+                ToolResolutionError). Ignored when success is True.
+        """
+        if success:
+            # Individual success doesn't reset — only a fully clean batch does.
+            # Call reset_errors() after a batch with zero errors.
+            return
+        if not is_soft_error:
+            self._consecutive_tool_errors += 1
+
+    def reset_errors(self) -> None:
+        """Reset tool error counter (call after a fully clean batch)."""
+        self._consecutive_tool_errors = 0
+
+    @property
+    def retries_exhausted(self) -> bool:
+        """True if consecutive retries exceed the limit."""
+        return self._consecutive_retries > self.max_retries
+
+    @property
+    def tool_errors_exhausted(self) -> bool:
+        """True if consecutive tool errors exceed the limit."""
+        return self._consecutive_tool_errors > self.max_tool_errors
+
+    @property
+    def consecutive_retries(self) -> int:
+        """Current consecutive retry count."""
+        return self._consecutive_retries
+
+    @property
+    def consecutive_tool_errors(self) -> int:
+        """Current consecutive tool error count."""
+        return self._consecutive_tool_errors

millforge/_forge/guardrails/guardrails.py

@@ -0,0 +1,194 @@
+"""Guardrails -- bundled middleware for foreign orchestration loops.
+
+Two-method API that wraps ResponseValidator, StepEnforcer, and ErrorTracker:
+
+    result = guardrails.check(response)   # before execution
+    done   = guardrails.record(["tool"])  # after execution
+
+For granular control, use the individual components directly.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Literal
+
+from millforge._forge.core.workflow import LLMResponse, ToolCall
+from millforge._forge.guardrails.error_tracker import ErrorTracker
+from millforge._forge.guardrails.nudge import (
+    TOOL_CHANNEL_KINDS,
+    TOOL_ERROR_KINDS,
+    Nudge,
+)
+from millforge._forge.guardrails.response_validator import ResponseValidator
+from millforge._forge.guardrails.step_enforcer import StepEnforcer
+
+
+@dataclass(frozen=True)
+class CheckResult:
+    """Result of checking an LLM response against all guardrails.
+
+    Attributes:
+        action: What the caller should do next.
+            "execute"      -- tool_calls are safe to run.
+            "retry"        -- model produced unusable output (bad format / bare
+                              text); inject nudge as a user message and re-prompt.
+            "tool_error"   -- model called a tool incorrectly (unknown name or
+                              malformed args); inject nudge as a tool result
+                              (nudge.role == "tool") and re-prompt.
+            "step_blocked" -- model tried to skip required steps; inject nudge.
+            "fatal"        -- error budget exhausted; stop the workflow.
+        tool_calls: Validated tool calls (only set when action == "execute").
+        nudge: Corrective message to inject (set when action is "retry",
+            "tool_error", or "step_blocked"). Emit it with its own role:
+            ``messages.append({"role": nudge.role, "content": nudge.content})``.
+        reason: Human-readable explanation (only set when action == "fatal").
+    """
+
+    action: Literal["execute", "retry", "tool_error", "step_blocked", "fatal"]
+    tool_calls: list[ToolCall] | None = None
+    nudge: Nudge | None = None
+    reason: str | None = None
+
+
+class Guardrails:
+    """Bundled guardrail middleware for foreign orchestration loops.
+
+    Wraps ResponseValidator, StepEnforcer, and ErrorTracker into a
+    two-method API. Use ``check()`` after each LLM response and
+    ``record()`` after executing tools.
+
+    Args:
+        tool_names: Valid tool names for this workflow.
+        required_steps: Tools that must be called before the terminal tool.
+            Defaults to no required steps.
+        terminal_tool: The tool(s) that can end the workflow. Accepts a
+            single name or a frozenset of names.
+        max_retries: Consecutive bad responses before ``check()`` returns
+            ``"fatal"``. Default 3.
+        max_tool_errors: Consecutive tool execution failures before
+            exhaustion. Default 2.
+        rescue_enabled: Attempt to parse tool calls from plain text
+            responses. Default True.
+        max_premature_attempts: Premature terminal attempts before
+            ``check()`` returns ``"fatal"``. Default 3.
+        retry_nudge: Custom nudge for bare text responses. Pass a callable
+            ``(raw_response) -> str`` for dynamic nudges. If None, uses
+            the default.
+    """
+
+    def __init__(
+        self,
+        tool_names: list[str],
+        terminal_tool: str | frozenset[str],
+        required_steps: list[str] | None = None,
+        max_retries: int = 3,
+        max_tool_errors: int = 2,
+        rescue_enabled: bool = True,
+        max_premature_attempts: int = 3,
+        retry_nudge: Callable[[str], str] | None = None,
+    ) -> None:
+        self._validator = ResponseValidator(
+            tool_names=tool_names,
+            rescue_enabled=rescue_enabled,
+            retry_nudge_fn=retry_nudge,
+        )
+        if isinstance(terminal_tool, str):
+            terminal_tools = frozenset([terminal_tool])
+        else:
+            terminal_tools = terminal_tool
+        self._enforcer = StepEnforcer(
+            required_steps=required_steps or [],
+            terminal_tools=terminal_tools,
+            max_premature_attempts=max_premature_attempts,
+        )
+        self._errors = ErrorTracker(
+            max_retries=max_retries,
+            max_tool_errors=max_tool_errors,
+        )
+
+    def check(
+        self,
+        response: LLMResponse,
+    ) -> CheckResult:
+        """Check an LLM response against all guardrails.
+
+        Call this after each LLM response, before executing any tools.
+
+        Args:
+            response: The LLM response -- either a TextResponse or a
+                list of ToolCall objects.
+
+        Returns:
+            CheckResult indicating what the caller should do next.
+        """
+        # Checkpoint 1: Is this response usable?
+        validation = self._validator.validate(response)
+
+        if validation.needs_retry:
+            nudge = validation.nudge
+            kind = nudge.kind if nudge is not None else ""
+            # Budget: malformed args drain the tool-error budget; everything
+            # else (bare text, unknown tool) drains the retry budget — matching
+            # run_inference so all three integration modes account identically.
+            if kind in TOOL_ERROR_KINDS:
+                self._errors.record_result(success=False)
+                if self._errors.tool_errors_exhausted:
+                    return CheckResult(
+                        action="fatal",
+                        reason="too many consecutive tool-argument errors",
+                    )
+            else:
+                self._errors.record_retry()
+                if self._errors.retries_exhausted:
+                    return CheckResult(
+                        action="fatal",
+                        reason="too many consecutive bad responses",
+                    )
+            # Channel: tool-call faults (unknown tool, malformed args) ride the
+            # tool-result channel; bare-text failures stay a user-role retry.
+            action = "tool_error" if kind in TOOL_CHANNEL_KINDS else "retry"
+            return CheckResult(action=action, nudge=nudge)
+
+        self._errors.reset_retries()
+
+        # Checkpoint 2: Is the model skipping required steps?
+        step_check = self._enforcer.check(validation.tool_calls)
+
+        if step_check.needs_nudge:
+            if self._enforcer.premature_exhausted:
+                return CheckResult(
+                    action="fatal",
+                    reason="model repeatedly skipped required steps",
+                )
+            return CheckResult(action="step_blocked", nudge=step_check.nudge)
+
+        return CheckResult(action="execute", tool_calls=validation.tool_calls)
+
+    def record(self, executed: list[str | tuple[str, dict]]) -> bool:
+        """Record which tools were successfully executed.
+
+        Call this after executing tools to keep the middleware in sync.
+
+        Args:
+            executed: Names of tools that succeeded, or (name, args) tuples
+                for prerequisite tracking.
+
+        Returns:
+            True if the terminal tool was reached and all required
+            steps are satisfied (workflow is done).
+        """
+        for entry in executed:
+            if isinstance(entry, tuple):
+                name, args = entry
+                self._enforcer.record(name, args)
+            else:
+                self._enforcer.record(entry)
+        self._errors.reset_errors()
+        self._enforcer.reset_premature()
+        return self._enforcer.is_satisfied() and any(
+            (entry if isinstance(entry, str) else entry[0])
+            in self._enforcer.terminal_tools
+            for entry in executed
+        )

millforge/_forge/guardrails/nudge.py

@@ -0,0 +1,47 @@
+"""Lightweight nudge message returned by guardrail components."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Nudge:
+    """A message to inject into conversation history.
+
+    Returned by guardrail components when the model needs correction.
+    The consumer maps this to their framework's message format::
+
+        # OpenAI-style
+        messages.append({"role": nudge.role, "content": nudge.content})
+
+        # LangChain
+        messages.append(HumanMessage(content=nudge.content))
+
+    Attributes:
+        role: Message role for injection ("user", "system", or "tool").
+        content: The nudge text.
+        kind: Identifies what generated the nudge ("retry", "unknown_tool",
+            "step"). Useful for logging/metrics and for WorkflowRunner to
+            map back to MessageType for compaction prioritization.
+        tier: Escalation level for step nudges (0 = N/A, 1-3 = escalating).
+    """
+
+    role: str
+    content: str
+    kind: str
+    tier: int = 0
+
+
+# Nudge kinds that drain the tool-error BUDGET (record_result, max_tool_errors)
+# rather than the retry budget. Malformed args are conceptually "tool called
+# with bad inputs" — same family as a runtime FileNotFoundError. Shared by
+# run_inference (core) and the Guardrails facade so both account identically.
+TOOL_ERROR_KINDS: frozenset[str] = frozenset({"tool_arg_validation"})
+
+# Nudge kinds emitted on the tool-result CHANNEL (role="tool") — the model
+# called a tool incorrectly (bad name or bad args), so the correction rides the
+# canonical tool channel rather than a user nudge. Maps to the Guardrails
+# facade's action="tool_error". Superset of TOOL_ERROR_KINDS: an unknown-tool
+# call rides the tool channel but only drains the retry budget.
+TOOL_CHANNEL_KINDS: frozenset[str] = frozenset({"unknown_tool", "tool_arg_validation"})

millforge/_forge/guardrails/response_validator.py

@@ -0,0 +1,119 @@
+"""Response validation — rescue, retry, and unknown-tool nudges."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from millforge._forge.core.workflow import LLMResponse, TextResponse, ToolCall
+from millforge._forge.guardrails.nudge import Nudge
+from millforge._forge.prompts.nudges import (
+    retry_nudge,
+    tool_arg_validation_nudge,
+    unknown_tool_nudge,
+)
+from millforge._forge.prompts.templates import rescue_tool_call
+
+
+@dataclass
+class ValidationResult:
+    """Result of validating an LLM response.
+
+    Exactly one of ``tool_calls`` or ``nudge`` is set:
+    - If ``needs_retry`` is False: ``tool_calls`` contains validated tool calls.
+    - If ``needs_retry`` is True: ``nudge`` contains the message to inject.
+    """
+
+    tool_calls: list[ToolCall] | None
+    nudge: Nudge | None
+    needs_retry: bool
+
+
+class ResponseValidator:
+    """Validates LLM responses: rescues tool calls from text, checks tool names.
+
+    Stateless — safe to reuse across turns and sessions.
+
+    Args:
+        tool_names: Valid tool names for this workflow.
+        rescue_enabled: If True, attempt to parse tool calls from TextResponse
+            before generating a retry nudge.
+        retry_nudge_fn: Custom nudge function for bare text responses. Takes
+            the raw response text and returns the nudge message. If None,
+            uses the default retry nudge from ``millforge._forge.prompts.nudges``.
+    """
+
+    def __init__(
+        self,
+        tool_names: list[str],
+        rescue_enabled: bool = True,
+        retry_nudge_fn: Callable[[str], str] | None = None,
+    ) -> None:
+        self.tool_names = tool_names
+        self.rescue_enabled = rescue_enabled
+        self._retry_nudge_fn = retry_nudge_fn or retry_nudge
+
+    def validate(
+        self,
+        response: LLMResponse,
+    ) -> ValidationResult:
+        """Validate an LLM response.
+
+        Args:
+            response: Either a TextResponse or a list of ToolCall objects.
+
+        Returns:
+            ValidationResult with tool_calls on success, or a Nudge on failure.
+        """
+        # TextResponse: rescue, then retry nudge
+        if isinstance(response, TextResponse):
+            if self.rescue_enabled:
+                rescued = rescue_tool_call(response.content, self.tool_names)
+                if rescued:
+                    return ValidationResult(
+                        tool_calls=rescued, nudge=None, needs_retry=False
+                    )
+            return ValidationResult(
+                tool_calls=None,
+                nudge=Nudge(
+                    role="user",
+                    content=self._retry_nudge_fn(response.content),
+                    kind="retry",
+                ),
+                needs_retry=True,
+            )
+
+        # list[ToolCall]: check for unknown tools first (cheap, no point
+        # validating args of a hallucinated tool).
+        tool_calls = response
+        unknown = [tc for tc in tool_calls if tc.tool not in self.tool_names]
+        if unknown:
+            return ValidationResult(
+                tool_calls=None,
+                nudge=Nudge(
+                    role="tool",
+                    content=unknown_tool_nudge(unknown[0].tool, self.tool_names),
+                    kind="unknown_tool",
+                ),
+                needs_retry=True,
+            )
+
+        # Args-shape check. ToolCall no longer enforces dict-args at
+        # construction (see workflow.py); the structural check lives here so
+        # malformed args ride the tool-error channel via inference.py instead
+        # of crashing the client parser.
+        bad_args = [tc for tc in tool_calls if not isinstance(tc.args, dict)]
+        if bad_args:
+            return ValidationResult(
+                tool_calls=None,
+                nudge=Nudge(
+                    role="tool",
+                    content=tool_arg_validation_nudge(
+                        bad_args[0].tool, bad_args[0].args
+                    ),
+                    kind="tool_arg_validation",
+                ),
+                needs_retry=True,
+            )
+
+        return ValidationResult(tool_calls=tool_calls, nudge=None, needs_retry=False)

millforge/_forge/guardrails/step_enforcer.py

@@ -0,0 +1,183 @@
+"""Step enforcement — required step tracking, premature terminal nudges,
+and prerequisite enforcement."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from millforge._forge.core.steps import StepTracker
+from millforge._forge.core.workflow import ToolCall
+from millforge._forge.guardrails.nudge import Nudge
+from millforge._forge.prompts.nudges import prerequisite_nudge, step_nudge
+
+
+@dataclass
+class StepCheck:
+    """Result of checking tool calls against step requirements.
+
+    If ``needs_nudge`` is True, ``nudge`` contains the message to inject.
+    """
+
+    nudge: Nudge | None
+    needs_nudge: bool
+
+
+class StepEnforcer:
+    """Tracks required steps and enforces them with escalating nudges.
+
+    Also enforces tool prerequisites — conditional dependencies between tools.
+
+    Stateful — instantiate per session/task.
+
+    Args:
+        required_steps: Tool names that must be called before the terminal tool.
+        terminal_tools: The tools that can end the workflow.
+        tool_prerequisites: Map of tool name to its ToolDef.prerequisites list.
+        max_premature_attempts: How many premature terminal attempts before
+            the enforcer signals exhaustion (via StepCheck or raising).
+        max_prereq_violations: How many consecutive prerequisite violations
+            before the enforcer signals exhaustion.
+    """
+
+    def __init__(
+        self,
+        required_steps: list[str],
+        terminal_tools: frozenset[str],
+        tool_prerequisites: dict[str, list[str | dict[str, str]]] | None = None,
+        max_premature_attempts: int = 3,
+        max_prereq_violations: int = 2,
+    ) -> None:
+        self._tracker = StepTracker(required_steps=required_steps)
+        self.terminal_tools = terminal_tools
+        self._tool_prerequisites = tool_prerequisites or {}
+        self.max_premature_attempts = max_premature_attempts
+        self.max_prereq_violations = max_prereq_violations
+        self._premature_attempts = 0
+        self._consecutive_prereq_violations = 0
+
+    def check(self, tool_calls: list[ToolCall]) -> StepCheck:
+        """Check whether tool calls include a premature terminal call.
+
+        If a terminal tool is in the batch and required steps aren't
+        satisfied, returns a StepCheck with an escalating nudge. The
+        escalation tier increments on each premature attempt (1=polite,
+        2=direct, 3=aggressive).
+
+        Args:
+            tool_calls: The tool calls the model wants to execute.
+
+        Returns:
+            StepCheck with nudge if premature, or no nudge if clear to proceed.
+        """
+        has_terminal = any(tc.tool in self.terminal_tools for tc in tool_calls)
+
+        if has_terminal and not self._tracker.is_satisfied():
+            self._premature_attempts += 1
+            tier = min(self._premature_attempts, 3)
+            # Find which terminal tool was attempted for the nudge message
+            attempted = next(
+                tc.tool for tc in tool_calls if tc.tool in self.terminal_tools
+            )
+            return StepCheck(
+                nudge=Nudge(
+                    role="user",
+                    content=step_nudge(
+                        attempted,
+                        self._tracker.pending(),
+                        tier=tier,
+                    ),
+                    kind="step",
+                    tier=tier,
+                ),
+                needs_nudge=True,
+            )
+
+        return StepCheck(nudge=None, needs_nudge=False)
+
+    def check_prerequisites(self, tool_calls: list[ToolCall]) -> StepCheck:
+        """Check whether any tool call has unsatisfied prerequisites.
+
+        Evaluates against pre-batch state. Any violation in the batch blocks
+        the entire batch (whole-batch blocking).
+
+        Args:
+            tool_calls: The tool calls the model wants to execute.
+
+        Returns:
+            StepCheck with nudge if any prereq is unsatisfied.
+        """
+        for tc in tool_calls:
+            prereqs = self._tool_prerequisites.get(tc.tool)
+            if not prereqs:
+                continue
+            result = self._tracker.check_prerequisites(tc.tool, tc.args, prereqs)
+            if not result.satisfied:
+                self._consecutive_prereq_violations += 1
+                return StepCheck(
+                    nudge=Nudge(
+                        role="user",
+                        content=prerequisite_nudge(
+                            tc.tool,
+                            result.missing,
+                        ),
+                        kind="prerequisite",
+                    ),
+                    needs_nudge=True,
+                )
+
+        return StepCheck(nudge=None, needs_nudge=False)
+
+    def record(self, tool_name: str, args: dict[str, Any] | None = None) -> None:
+        """Record a successful tool execution."""
+        self._tracker.record(tool_name, args)
+
+    def is_satisfied(self) -> bool:
+        """True if all required steps have been completed."""
+        return self._tracker.is_satisfied()
+
+    def pending(self) -> list[str]:
+        """Return required steps not yet completed."""
+        return self._tracker.pending()
+
+    def terminal_reached(self, tool_calls: list[ToolCall]) -> bool:
+        """True if a terminal tool is in the batch and steps are satisfied."""
+        has_terminal = any(tc.tool in self.terminal_tools for tc in tool_calls)
+        return has_terminal and self._tracker.is_satisfied()
+
+    @property
+    def premature_attempts(self) -> int:
+        """Number of premature terminal attempts so far."""
+        return self._premature_attempts
+
+    @property
+    def premature_exhausted(self) -> bool:
+        """True if premature attempts exceed the limit."""
+        return self._premature_attempts > self.max_premature_attempts
+
+    @property
+    def prereq_violations(self) -> int:
+        """Number of consecutive prerequisite violations."""
+        return self._consecutive_prereq_violations
+
+    @property
+    def prereq_exhausted(self) -> bool:
+        """True if consecutive prereq violations exceed the limit."""
+        return self._consecutive_prereq_violations > self.max_prereq_violations
+
+    def reset_premature(self) -> None:
+        """Reset premature attempt counter (call after a clean batch)."""
+        self._premature_attempts = 0
+
+    def reset_prereq_violations(self) -> None:
+        """Reset consecutive prereq violation counter (call after a clean batch)."""
+        self._consecutive_prereq_violations = 0
+
+    @property
+    def completed_steps(self) -> dict[str, None]:
+        """Steps completed so far (for diagnostics / error reporting)."""
+        return self._tracker.completed_steps
+
+    def summary_hint(self) -> str:
+        """Human-readable hint for context compaction."""
+        return self._tracker.summary_hint()

millforge/_forge/prompts/__init__.py

@@ -0,0 +1,16 @@
+"""Private Forge prompt helpers."""
+
+from millforge._forge.prompts.nudges import retry_nudge, step_nudge
+from millforge._forge.prompts.templates import (
+    build_tool_prompt,
+    extract_tool_call,
+    rescue_tool_call,
+)
+
+__all__ = [
+    "build_tool_prompt",
+    "extract_tool_call",
+    "rescue_tool_call",
+    "retry_nudge",
+    "step_nudge",
+]