chuk-ai-session-manager 0.7.1__py3-none-any.whl → 0.8.1__py3-none-any.whl

chuk_ai_session_manager/guards/models.py

@@ -0,0 +1,498 @@
+# chuk_ai_session_manager/guards/models.py
+"""Pydantic models for tool state management.
+
+All state-related models in one place, fully type-safe.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from datetime import datetime
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+def classify_value_type(value: Any) -> "ValueType":
+    """Classify the type of a value for binding."""
+    if isinstance(value, (int, float)):
+        return ValueType.NUMBER
+    if isinstance(value, str):
+        try:
+            float(value)
+            return ValueType.NUMBER
+        except (ValueError, TypeError):
+            return ValueType.STRING
+    if isinstance(value, list):
+        return ValueType.LIST
+    if isinstance(value, dict):
+        return ValueType.OBJECT
+    return ValueType.UNKNOWN
+
+
+def compute_args_hash(arguments: dict[str, Any]) -> str:
+    """Compute a stable hash of tool arguments."""
+    args_str = json.dumps(arguments, sort_keys=True, default=str)
+    return hashlib.sha256(args_str.encode()).hexdigest()[:16]
+
+
+class ValueType(str, Enum):
+    """Types for bound values."""
+
+    NUMBER = "number"
+    STRING = "string"
+    JSON = "json"
+    LIST = "list"
+    OBJECT = "object"
+    UNKNOWN = "unknown"
+
+
+class ToolClassification:
+    """Central definitions for tool classification.
+
+    Guards and managers should use these definitions rather than
+    maintaining their own hardcoded sets.
+    """
+
+    # Discovery tools - search/list/get schemas (count against discovery budget)
+    DISCOVERY_TOOLS: frozenset[str] = frozenset(
+        {
+            "list_tools",
+            "search_tools",
+            "get_tool_schema",
+            "get_tool_schemas",
+        }
+    )
+
+    # Idempotent math tools - safe to call multiple times, exempt from per-tool limits
+    IDEMPOTENT_MATH_TOOLS: frozenset[str] = frozenset(
+        {
+            "add",
+            "subtract",
+            "multiply",
+            "divide",
+            "sqrt",
+            "pow",
+            "power",
+            "log",
+            "exp",
+            "sin",
+            "cos",
+            "tan",
+            "abs",
+            "floor",
+            "ceil",
+            "round",
+        }
+    )
+
+    # Parameterized tools - require computed input values (precondition guard)
+    PARAMETERIZED_TOOLS: frozenset[str] = frozenset(
+        {
+            "normal_cdf",
+            "normal_pdf",
+            "normal_sf",
+            "t_cdf",
+            "t_sf",
+            "t_test",
+            "chi_cdf",
+            "chi_sf",
+            "chi_square",
+        }
+    )
+
+    @classmethod
+    def is_discovery_tool(cls, tool_name: str) -> bool:
+        """Check if tool is a discovery tool."""
+        base = (
+            tool_name.split(".")[-1].lower() if "." in tool_name else tool_name.lower()
+        )
+        return base in cls.DISCOVERY_TOOLS
+
+    @classmethod
+    def is_idempotent_math_tool(cls, tool_name: str) -> bool:
+        """Check if tool is an idempotent math tool."""
+        base = (
+            tool_name.split(".")[-1].lower() if "." in tool_name else tool_name.lower()
+        )
+        return base in cls.IDEMPOTENT_MATH_TOOLS
+
+    @classmethod
+    def is_parameterized_tool(cls, tool_name: str) -> bool:
+        """Check if tool requires computed values."""
+        base = (
+            tool_name.split(".")[-1].lower() if "." in tool_name else tool_name.lower()
+        )
+        return base in cls.PARAMETERIZED_TOOLS
+
+
+class RuntimeMode(str, Enum):
+    """Runtime enforcement mode presets."""
+
+    SMOOTH = "smooth"  # Feels like ChatGPT/Claude UI - warn but allow
+    STRICT = "strict"  # Best for solver/math/physics - hard enforcement
+
+
+class EnforcementLevel(str, Enum):
+    """Enforcement level for guards and constraints."""
+
+    OFF = "off"  # No enforcement
+    WARN = "warn"  # Proceed but log warning
+    BLOCK = "block"  # Do not execute, return error
+
+
+class CacheScope(str, Enum):
+    """Scope for result caching."""
+
+    TURN = "turn"  # Cache per conversation turn
+    SESSION = "session"  # Cache for entire session
+
+
+class UnusedResultAction(str, Enum):
+    """Action to take for unused tool results."""
+
+    OFF = "off"  # No enforcement
+    WARN = "warn"  # Warn but continue
+    BLOCK_NEXT_TOOL = "block-next-tool"  # Block next tool call
+
+
+class ValueBinding(BaseModel):
+    """A bound value from a tool result with a stable ID.
+
+    Every tool result gets assigned a value ID (v1, v2, v3...) that can
+    be referenced in subsequent tool calls using $vN syntax.
+    """
+
+    id: str = Field(..., description="Value ID, e.g., 'v1', 'v2'")
+    tool_name: str = Field(..., description="Name of the tool that produced this value")
+    args_hash: str = Field(..., description="Hash of arguments for dedup")
+    raw_value: Any = Field(..., description="The raw value from the tool")
+    value_type: ValueType = Field(..., description="Classified type of the value")
+    timestamp: datetime = Field(default_factory=datetime.now)
+    aliases: list[str] = Field(default_factory=list, description="Model-provided names")
+    used: bool = Field(default=False, description="Has this value been referenced?")
+    used_in: list[str] = Field(
+        default_factory=list, description="Tool calls that used this"
+    )
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    @property
+    def typed_value(self) -> Any:
+        """Get the value with appropriate type coercion."""
+        if self.value_type == ValueType.NUMBER:
+            if isinstance(self.raw_value, (int, float)):
+                return float(self.raw_value)
+            try:
+                return float(self.raw_value)
+            except (ValueError, TypeError):
+                return self.raw_value
+        return self.raw_value
+
+    def format_for_model(self) -> str:
+        """Format this binding for display to the model."""
+        if self.value_type == ValueType.NUMBER:
+            val = self.typed_value
+            if isinstance(val, float):
+                if abs(val) < 0.0001 or abs(val) > 10000:
+                    formatted = f"{val:.6e}"
+                else:
+                    formatted = f"{val:.6f}"
+            else:
+                formatted = str(val)
+        elif self.value_type == ValueType.STRING:
+            raw_str = str(self.raw_value)
+            formatted = f'"{raw_str}"' if len(raw_str) < 50 else f'"{raw_str[:47]}..."'
+        elif self.value_type == ValueType.LIST:
+            lst = self.raw_value
+            if isinstance(lst, list):
+                if len(lst) == 0:
+                    formatted = "[]"
+                elif len(lst) <= 3:
+                    formatted = str(lst)[:60]
+                else:
+                    formatted = f"[{len(lst)} items]"
+            else:
+                formatted = str(lst)[:50]
+        elif self.value_type == ValueType.OBJECT:
+            obj = self.raw_value
+            if isinstance(obj, dict):
+                keys = list(obj.keys())
+                if len(keys) == 0:
+                    formatted = "{}"
+                elif len(keys) <= 3:
+                    formatted = f"{{keys: {keys}}}"
+                else:
+                    formatted = f"{{object with {len(keys)} keys}}"
+            else:
+                formatted = str(obj)[:50]
+        else:
+            formatted = str(self.raw_value)[:50]
+
+        alias_str = f" (aka {', '.join(self.aliases)})" if self.aliases else ""
+        return f"${self.id} = {formatted}{alias_str}  # from {self.tool_name}"
+
+
+class ReferenceCheckResult(BaseModel):
+    """Result of checking references in tool arguments."""
+
+    valid: bool = Field(..., description="Whether all references are valid")
+    missing_refs: list[str] = Field(
+        default_factory=list, description="References that don't exist"
+    )
+    resolved_refs: dict[str, Any] = Field(
+        default_factory=dict, description="ref -> resolved value"
+    )
+    message: str = Field(default="", description="Human-readable message")
+
+
+class PerToolCallStatus(BaseModel):
+    """Status of per-tool call tracking for anti-thrash."""
+
+    tool_name: str
+    call_count: int = Field(default=0)
+    max_calls: int = Field(
+        default=3, description="Default max before requiring justification"
+    )
+    requires_justification: bool = Field(default=False)
+
+
+class UngroundedCallResult(BaseModel):
+    """Result of checking if a tool call is ungrounded.
+
+    An ungrounded call is one where:
+    - Arguments contain numeric literals
+    - No $vN references are present
+    - Values exist that could have been referenced
+    """
+
+    is_ungrounded: bool = Field(default=False)
+    numeric_args: list[str] = Field(default_factory=list)
+    has_bindings: bool = Field(default=False)
+    message: str = Field(default="")
+
+
+class RunawayStatus(BaseModel):
+    """Status of runaway detection checks."""
+
+    should_stop: bool = Field(default=False)
+    reason: str | None = Field(default=None)
+    budget_exhausted: bool = Field(default=False)
+    degenerate_detected: bool = Field(default=False)
+    saturation_detected: bool = Field(default=False)
+    calls_remaining: int = Field(default=0)
+
+    @property
+    def message(self) -> str:
+        """Get user-friendly message about why we should stop."""
+        if self.budget_exhausted:
+            return f"Tool call budget exhausted ({self.calls_remaining} remaining). Use computed values to answer."
+        if self.degenerate_detected:
+            return "Degenerate output detected (0.0, 1.0, or repeating). Results have saturated."
+        if self.saturation_detected:
+            return (
+                "Numeric saturation detected. Values are at machine precision limits."
+            )
+        return self.reason or "Unknown stop reason"
+
+
+class CachedToolResult(BaseModel):
+    """A cached tool call result with metadata."""
+
+    tool_name: str
+    arguments: dict[str, Any]
+    result: Any
+    timestamp: datetime = Field(default_factory=datetime.now)
+    call_count: int = Field(default=1)
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    @property
+    def signature(self) -> str:
+        """Generate unique signature for this tool call."""
+        args_str = json.dumps(self.arguments, sort_keys=True, default=str)
+        return f"{self.tool_name}:{args_str}"
+
+    @property
+    def is_numeric(self) -> bool:
+        """Check if result is numeric."""
+        if isinstance(self.result, (int, float)):
+            return True
+        if isinstance(self.result, str):
+            try:
+                float(self.result)
+                return True
+            except (ValueError, TypeError):
+                return False
+        return False
+
+    @property
+    def numeric_value(self) -> float | None:
+        """Extract numeric value if available."""
+        if isinstance(self.result, (int, float)):
+            return float(self.result)
+        if isinstance(self.result, str):
+            try:
+                return float(self.result)
+            except (ValueError, TypeError):
+                return None
+        return None
+
+    def format_compact(self) -> str:
+        """Format for compact state display."""
+        if self.is_numeric:
+            val = self.numeric_value
+            if val is not None:
+                if abs(val) < 0.0001 or abs(val) > 10000:
+                    return f"{self.tool_name}({self._format_args()}) = {val:.6e}"
+                else:
+                    return f"{self.tool_name}({self._format_args()}) = {val:.6f}"
+        result_str = str(self.result)
+        if len(result_str) > 50:
+            result_str = result_str[:47] + "..."
+        return f"{self.tool_name}({self._format_args()}) = {result_str}"
+
+    def _format_args(self) -> str:
+        """Format arguments compactly."""
+        if not self.arguments:
+            return ""
+        if len(self.arguments) == 1:
+            val = list(self.arguments.values())[0]
+            if isinstance(val, (int, float)):
+                return str(val)
+        parts = []
+        for k, v in self.arguments.items():
+            if isinstance(v, (int, float)):
+                parts.append(f"{k}={v}")
+            elif isinstance(v, str) and len(v) < 20:
+                parts.append(f'{k}="{v}"')
+            else:
+                parts.append(f"{k}=...")
+        return ", ".join(parts)
+
+
+class NamedVariable(BaseModel):
+    """A named variable binding from tool results."""
+
+    name: str
+    value: float
+    units: str | None = Field(default=None)
+    source_tool: str | None = Field(default=None)
+    source_args: dict[str, Any] | None = Field(default=None)
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    def format_compact(self) -> str:
+        """Format for state display."""
+        if self.units:
+            return f"{self.name} = {self.value:.6f} {self.units}"
+        return f"{self.name} = {self.value:.6f}"
+
+
+class SoftBlockReason(str, Enum):
+    """Reasons why a tool call might be soft-blocked."""
+
+    UNGROUNDED_ARGS = "ungrounded_args"
+    MISSING_REFS = "missing_refs"
+    BUDGET_EXHAUSTED = "budget_exhausted"
+    PER_TOOL_LIMIT = "per_tool_limit"
+    MISSING_DEPENDENCY = "missing_dependency"
+
+
+class RepairAction(str, Enum):
+    """Actions the runtime can take to repair a soft-blocked call."""
+
+    REBIND_FROM_EXISTING = "rebind_from_existing"
+    COMPUTE_MISSING = "compute_missing"
+    REWRITE_CALL = "rewrite_call"
+    SYMBOLIC_FALLBACK = "symbolic_fallback"
+    ASK_USER = "ask_user"
+
+
+class SoftBlock(BaseModel):
+    """A soft block that can potentially be repaired."""
+
+    reason: SoftBlockReason
+    tool_name: str = Field(default="")
+    arguments: dict[str, Any] = Field(default_factory=dict)
+    message: str = Field(default="")
+    repair_attempts: int = Field(default=0)
+    max_repairs: int = Field(default=3)
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    @property
+    def can_repair(self) -> bool:
+        """Check if we can attempt another repair."""
+        return self.repair_attempts < self.max_repairs
+
+    @property
+    def next_repair_action(self) -> RepairAction:
+        """Get the next repair action to try."""
+        if self.reason == SoftBlockReason.UNGROUNDED_ARGS:
+            return RepairAction.REBIND_FROM_EXISTING
+        elif self.reason == SoftBlockReason.MISSING_REFS:
+            return RepairAction.COMPUTE_MISSING
+        elif self.reason == SoftBlockReason.MISSING_DEPENDENCY:
+            return RepairAction.COMPUTE_MISSING
+        return RepairAction.ASK_USER
+
+
+class RuntimeLimits(BaseModel):
+    """Configuration for tool runtime enforcement.
+
+    Two preset modes:
+    - SMOOTH: Feels like ChatGPT/Claude UI - warn but allow, auto-retry
+    - STRICT: Best for solver/math - hard enforcement, dataflow discipline
+    """
+
+    # Discovery controls (search_tools, list_tools, get_tool_schema)
+    discovery_budget: int = Field(default=5)
+
+    # Execution controls (call_tool)
+    execution_budget: int = Field(default=12)
+    tool_budget_total: int = Field(default=15)
+
+    # Per-tool caps (anti-thrash)
+    # Set to 0 or negative to disable per-tool limits
+    per_tool_cap: int = Field(default=0)
+
+    # Cache behavior
+    cache_scope: CacheScope = Field(default=CacheScope.TURN)
+
+    # Binding enforcement
+    require_bindings: EnforcementLevel = Field(default=EnforcementLevel.WARN)
+    ungrounded_grace_calls: int = Field(default=1)
+
+    # Unused result enforcement
+    unused_results: UnusedResultAction = Field(default=UnusedResultAction.WARN)
+
+    @classmethod
+    def smooth(cls) -> "RuntimeLimits":
+        """Preset for smooth UI-like experience."""
+        return cls(
+            discovery_budget=6,
+            execution_budget=15,
+            tool_budget_total=20,
+            per_tool_cap=0,  # Unlimited per-tool calls
+            cache_scope=CacheScope.TURN,
+            require_bindings=EnforcementLevel.WARN,
+            ungrounded_grace_calls=2,
+            unused_results=UnusedResultAction.WARN,
+        )
+
+    @classmethod
+    def strict(cls) -> "RuntimeLimits":
+        """Preset for strict dataflow enforcement."""
+        return cls(
+            discovery_budget=4,
+            execution_budget=10,
+            tool_budget_total=12,
+            per_tool_cap=0,  # Unlimited per-tool calls
+            cache_scope=CacheScope.TURN,
+            require_bindings=EnforcementLevel.BLOCK,
+            ungrounded_grace_calls=0,
+            unused_results=UnusedResultAction.WARN,
+        )

chuk_ai_session_manager/guards/ungrounded.py

@@ -0,0 +1,159 @@
+# chuk_ai_session_manager/guards/ungrounded.py
+"""Ungrounded call guard - detects missing $vN references.
+
+Catches when the model passes numeric literals that should have been
+the result of prior computation. Enforces dataflow discipline.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+# Import base classes from chuk-tool-processor
+from chuk_tool_processor.guards import BaseGuard, EnforcementLevel, GuardResult
+
+# Reference pattern: $v1, $v2, ${v1}, ${myalias}
+REFERENCE_PATTERN = re.compile(r"\$\{?([a-zA-Z_][a-zA-Z0-9_]*|v\d+)\}?")
+
+
+class UngroundedGuardConfig(BaseModel):
+    """Configuration for ungrounded call detection."""
+
+    # How many ungrounded calls before blocking
+    grace_calls: int = Field(default=1)
+
+    # Enforcement level
+    mode: EnforcementLevel = Field(default=EnforcementLevel.WARN)
+
+
+class UngroundedGuard(BaseGuard):
+    """Guard that detects ungrounded numeric arguments.
+
+    An ungrounded call is when:
+    - Arguments contain numeric literals (int or float)
+    - No $vN references exist in the arguments
+    - The model should have used a computed value instead
+
+    This catches the anti-pattern where a model passes a literal number
+    that should have been the result of a prior computation.
+    """
+
+    def __init__(
+        self,
+        config: UngroundedGuardConfig | None = None,
+        get_user_literals: Any = None,  # Callable[[], set[float]]
+        get_bindings: Any = None,  # Callable[[], dict]
+    ):
+        self.config = config or UngroundedGuardConfig()
+        self._get_user_literals = get_user_literals
+        self._get_bindings = get_bindings
+        self._ungrounded_count = 0
+
+    def check(
+        self,
+        tool_name: str,
+        arguments: dict[str, Any],
+    ) -> GuardResult:
+        """Check if tool call has ungrounded numeric arguments.
+
+        Args:
+            tool_name: Name of the tool being called
+            arguments: Arguments passed to the tool
+
+        Returns:
+            GuardResult - WARN or BLOCK if ungrounded
+        """
+        if self.config.mode == EnforcementLevel.OFF:
+            return self.allow()
+
+        # Get user-provided literals (these are allowed)
+        user_literals = self._get_user_literals() if self._get_user_literals else set()
+
+        # Find numeric arguments not from user
+        numeric_args = self._find_numeric_args(arguments, user_literals)
+        if not numeric_args:
+            return self.allow()
+
+        # Check if any $vN references exist
+        args_str = json.dumps(arguments, default=str)
+        has_refs = bool(REFERENCE_PATTERN.search(args_str))
+
+        if has_refs:
+            # Has references - not ungrounded
+            return self.allow()
+
+        # Ungrounded call detected
+        self._ungrounded_count += 1
+
+        # Get available bindings for helpful message
+        bindings = self._get_bindings() if self._get_bindings else {}
+        has_bindings = bool(bindings)
+
+        # Build message
+        arg_names = ", ".join(f"`{name}`" for name in numeric_args.keys())
+        if has_bindings:
+            available = [f"${bid}" for bid in bindings.keys()]
+            message = (
+                f"Ungrounded call: `{tool_name}` has numeric arguments ({arg_names}) "
+                f"but no $vN references. Available values: {', '.join(available)}. "
+                "Did you mean to use a computed value?"
+            )
+        else:
+            message = (
+                f"Ungrounded call: `{tool_name}` has numeric arguments ({arg_names}) "
+                "but no prior computations exist. Compute input values first."
+            )
+
+        # Check enforcement level
+        if (
+            self.config.mode == EnforcementLevel.BLOCK
+            and self._ungrounded_count > self.config.grace_calls
+        ):
+            return self.block(
+                reason=message,
+                ungrounded_count=self._ungrounded_count,
+                numeric_args=numeric_args,
+                has_bindings=has_bindings,
+            )
+        else:
+            return self.warn(
+                reason=message,
+                ungrounded_count=self._ungrounded_count,
+                grace_remaining=max(
+                    0, self.config.grace_calls - self._ungrounded_count + 1
+                ),
+                numeric_args=numeric_args,
+                has_bindings=has_bindings,
+            )
+
+    def reset(self) -> None:
+        """Reset for new prompt."""
+        self._ungrounded_count = 0
+
+    def _find_numeric_args(
+        self,
+        arguments: dict[str, Any],
+        user_literals: set[float],
+    ) -> dict[str, float]:
+        """Find numeric arguments not from user input."""
+        numeric = {}
+        for key, value in arguments.items():
+            if key == "tool_name":
+                continue
+            if isinstance(value, bool):
+                continue
+            if isinstance(value, (int, float)):
+                if float(value) not in user_literals:
+                    numeric[key] = value
+            elif isinstance(value, str):
+                try:
+                    num_val = float(value)
+                    if num_val not in user_literals:
+                        numeric[key] = num_val
+                except (ValueError, TypeError):
+                    pass
+        return numeric