zwarm 2.3.5__py3-none-any.whl

zwarm/core/compact.py

@@ -0,0 +1,329 @@
+"""
+Message compaction for context window management.
+
+Safely prunes old messages while preserving:
+- System prompt and initial user task
+- Tool call/response pairs (never orphaned)
+- Recent conversation context
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def _get_attr(obj: Any, key: str, default: Any = None) -> Any:
+    """Get attribute from dict or object (handles both Pydantic models and dicts)."""
+    if isinstance(obj, dict):
+        return obj.get(key, default)
+    return getattr(obj, key, default)
+
+
+@dataclass
+class CompactionResult:
+    """Result of a compaction operation."""
+
+    messages: list[dict[str, Any]]
+    removed_count: int
+    original_count: int
+    preserved_reason: str | None = None
+
+    @property
+    def was_compacted(self) -> bool:
+        return self.removed_count > 0
+
+
+def estimate_tokens(messages: list[Any]) -> int:
+    """
+    Rough token estimate for messages.
+
+    Uses ~4 chars per token as a simple heuristic.
+    This is intentionally conservative.
+    Handles both dict messages and Pydantic model messages.
+    """
+    total_chars = 0
+    for msg in messages:
+        content = _get_attr(msg, "content", "")
+        if isinstance(content, str):
+            total_chars += len(content)
+        elif isinstance(content, list):
+            # Anthropic-style content blocks
+            for block in content:
+                if isinstance(block, dict):
+                    total_chars += len(str(block.get("text", "")))
+                    total_chars += len(str(block.get("input", "")))
+                elif isinstance(block, str):
+                    total_chars += len(block)
+                else:
+                    # Pydantic model block
+                    total_chars += len(str(_get_attr(block, "text", "")))
+                    total_chars += len(str(_get_attr(block, "input", "")))
+
+        # Tool calls add tokens too
+        tool_calls = _get_attr(msg, "tool_calls", []) or []
+        for tc in tool_calls:
+            func = _get_attr(tc, "function", {}) or {}
+            args = _get_attr(func, "arguments", "") if isinstance(func, dict) else getattr(func, "arguments", "")
+            total_chars += len(str(args))
+
+    return total_chars // 4
+
+
+def find_tool_groups(messages: list[Any]) -> list[tuple[int, int]]:
+    """
+    Find message index ranges that form tool call groups.
+
+    A tool call group is:
+    - An assistant message with tool_calls
+    - All following tool/user response messages until the next assistant message
+
+    This handles both OpenAI format (role="tool") and Anthropic format
+    (role="user" with tool_result content).
+    Also handles Pydantic model messages.
+
+    Returns list of (start_idx, end_idx) tuples (inclusive).
+    """
+    groups = []
+    i = 0
+
+    while i < len(messages):
+        msg = messages[i]
+
+        # Check for tool calls in assistant message
+        has_tool_calls = False
+
+        # OpenAI format: tool_calls field
+        if _get_attr(msg, "role") == "assistant" and _get_attr(msg, "tool_calls"):
+            has_tool_calls = True
+
+        # Anthropic format: content blocks with type="tool_use"
+        if _get_attr(msg, "role") == "assistant":
+            content = _get_attr(msg, "content", [])
+            if isinstance(content, list):
+                for block in content:
+                    block_type = _get_attr(block, "type", None)
+                    if block_type == "tool_use":
+                        has_tool_calls = True
+                        break
+
+        if has_tool_calls:
+            start = i
+            j = i + 1
+
+            # Find all following tool responses
+            while j < len(messages):
+                next_msg = messages[j]
+                role = _get_attr(next_msg, "role", "")
+
+                # OpenAI format: tool role
+                if role == "tool":
+                    j += 1
+                    continue
+
+                # Anthropic format: user message with tool_result
+                if role == "user":
+                    content = _get_attr(next_msg, "content", [])
+                    if isinstance(content, list):
+                        has_tool_result = any(
+                            _get_attr(b, "type", None) == "tool_result"
+                            for b in content
+                        )
+                        if has_tool_result:
+                            j += 1
+                            continue
+
+                # Not a tool response, stop here
+                break
+
+            groups.append((start, j - 1))
+            i = j
+        else:
+            i += 1
+
+    return groups
+
+
+def compact_messages(
+    messages: list[Any],
+    keep_first_n: int = 2,
+    keep_last_n: int = 10,
+    max_tokens: int | None = None,
+    target_token_pct: float = 0.7,
+) -> CompactionResult:
+    """
+    Compact message history by removing old messages (LRU-style).
+
+    Preserves:
+    - First N messages (system prompt, user task)
+    - Last N messages (recent context)
+    - Tool call/response pairs are NEVER split
+
+    Args:
+        messages: The message list to compact
+        keep_first_n: Number of messages to always keep at the start
+        keep_last_n: Number of messages to always keep at the end
+        max_tokens: If set, compact when estimated tokens exceed this
+        target_token_pct: Target percentage of max_tokens after compaction
+
+    Returns:
+        CompactionResult with the compacted messages and stats
+    """
+    original_count = len(messages)
+
+    # Nothing to compact if we have few messages
+    if len(messages) <= keep_first_n + keep_last_n:
+        return CompactionResult(
+            messages=messages,
+            removed_count=0,
+            original_count=original_count,
+            preserved_reason="Too few messages to compact",
+        )
+
+    # Check if compaction is needed based on tokens
+    if max_tokens:
+        current_tokens = estimate_tokens(messages)
+        if current_tokens < max_tokens:
+            return CompactionResult(
+                messages=messages,
+                removed_count=0,
+                original_count=original_count,
+                preserved_reason=f"Under token limit ({current_tokens}/{max_tokens})",
+            )
+
+    # Find tool call groups (these must stay together)
+    tool_groups = find_tool_groups(messages)
+
+    # Build a set of "protected" indices (in tool groups)
+    protected_indices: set[int] = set()
+    for start, end in tool_groups:
+        for idx in range(start, end + 1):
+            protected_indices.add(idx)
+
+    # Determine which messages are in the "middle" (candidates for removal)
+    # Middle = not in first N, not in last N
+    middle_start = keep_first_n
+    middle_end = len(messages) - keep_last_n
+
+    if middle_start >= middle_end:
+        return CompactionResult(
+            messages=messages,
+            removed_count=0,
+            original_count=original_count,
+            preserved_reason="No middle messages to remove",
+        )
+
+    # Find removable message ranges in the middle
+    # We remove from the oldest (lowest index) first
+    removable_ranges: list[tuple[int, int]] = []
+    i = middle_start
+
+    while i < middle_end:
+        # Check if this index is in a tool group
+        in_group = False
+        for start, end in tool_groups:
+            if start <= i <= end:
+                # This message is part of a tool group
+                # Check if the ENTIRE group is in the middle
+                if start >= middle_start and end < middle_end:
+                    # Entire group is removable as a unit
+                    removable_ranges.append((start, end))
+                    i = end + 1
+                    in_group = True
+                    break
+                else:
+                    # Group spans protected region, skip it entirely
+                    i = end + 1
+                    in_group = True
+                    break
+
+        if not in_group:
+            # Single message, can be removed individually
+            removable_ranges.append((i, i))
+            i += 1
+
+    # Deduplicate and sort ranges
+    removable_ranges = sorted(set(removable_ranges), key=lambda x: x[0])
+
+    if not removable_ranges:
+        return CompactionResult(
+            messages=messages,
+            removed_count=0,
+            original_count=original_count,
+            preserved_reason="All middle messages are in protected tool groups",
+        )
+
+    # Determine how many to remove
+    # Start by removing the oldest half of removable ranges
+    if max_tokens:
+        # Token-based: remove until under target
+        target_tokens = int(max_tokens * target_token_pct)
+        indices_to_remove: set[int] = set()
+
+        for start, end in removable_ranges:
+            for idx in range(start, end + 1):
+                indices_to_remove.add(idx)
+
+            # Check if we've removed enough
+            remaining = [m for i, m in enumerate(messages) if i not in indices_to_remove]
+            if estimate_tokens(remaining) <= target_tokens:
+                break
+    else:
+        # Count-based: remove oldest half of middle
+        total_removable = sum(end - start + 1 for start, end in removable_ranges)
+        target_remove = total_removable // 2
+
+        indices_to_remove = set()
+        removed = 0
+
+        for start, end in removable_ranges:
+            if removed >= target_remove:
+                break
+            for idx in range(start, end + 1):
+                indices_to_remove.add(idx)
+                removed += 1
+
+    # Build new message list
+    new_messages = [m for i, m in enumerate(messages) if i not in indices_to_remove]
+
+    # Add a compaction marker so the model knows history was truncated
+    if indices_to_remove and len(new_messages) > keep_first_n:
+        # Insert marker after the preserved first messages
+        marker = {
+            "role": "system",
+            "content": (
+                f"[Context compacted: {len(indices_to_remove)} older messages removed "
+                f"to manage context window. Conversation continues below.]"
+            ),
+        }
+        new_messages.insert(keep_first_n, marker)
+
+    logger.info(
+        f"Compacted messages: {original_count} -> {len(new_messages)} "
+        f"(removed {len(indices_to_remove)})"
+    )
+
+    return CompactionResult(
+        messages=new_messages,
+        removed_count=len(indices_to_remove),
+        original_count=original_count,
+    )
+
+
+def should_compact(
+    messages: list[Any],
+    max_tokens: int,
+    threshold_pct: float = 0.85,
+) -> bool:
+    """
+    Check if messages should be compacted.
+
+    Returns True if estimated tokens exceed threshold percentage of max.
+    Handles both dict messages and Pydantic model messages.
+    """
+    current = estimate_tokens(messages)
+    threshold = int(max_tokens * threshold_pct)
+    return current >= threshold

zwarm/core/config.py

@@ -0,0 +1,344 @@
+"""
+Configuration system for zwarm.
+
+Supports:
+- config.toml for user settings (weave project, defaults)
+- .env for environment variables
+- Composable YAML configs with inheritance (extends:)
+- CLI overrides via --set key=value
+"""
+
+from __future__ import annotations
+
+import os
+import tomllib
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import yaml
+from dotenv import load_dotenv
+
+
+@dataclass
+class WeaveConfig:
+    """Weave integration settings."""
+
+    project: str | None = None
+    enabled: bool = True
+
+
+@dataclass
+class ExecutorConfig:
+    """Configuration for an executor (coding agent)."""
+
+    adapter: str = "codex_mcp"  # codex_mcp | codex_exec | claude_code
+    model: str | None = None
+    sandbox: str = "workspace-write"  # read-only | workspace-write | danger-full-access
+    timeout: int = 3600
+    reasoning_effort: str | None = "high"  # low | medium | high (default to high for compatibility)
+
+
+@dataclass
+class CompactionConfig:
+    """Configuration for context window compaction."""
+
+    enabled: bool = True
+    max_tokens: int = 100000  # Trigger compaction when estimated tokens exceed this
+    threshold_pct: float = 0.85  # Compact when at this % of max_tokens
+    target_pct: float = 0.7  # Target this % after compaction
+    keep_first_n: int = 2  # Always keep first N messages (system + task)
+    keep_last_n: int = 10  # Always keep last N messages (recent context)
+
+
+@dataclass
+class OrchestratorConfig:
+    """Configuration for the orchestrator."""
+
+    lm: str = "gpt-5-mini"
+    prompt: str | None = None  # path to prompt yaml
+    tools: list[str] = field(default_factory=lambda: ["delegate", "converse", "check_session", "end_session", "bash"])
+    max_steps: int = 50
+    parallel_delegations: int = 4
+    sync_first: bool = True  # prefer sync mode by default
+    compaction: CompactionConfig = field(default_factory=CompactionConfig)
+
+    # Directory restrictions for agent delegations
+    # None = only working_dir allowed (most restrictive, default)
+    # ["*"] = any directory allowed (dangerous)
+    # ["/path/a", "/path/b"] = only these directories allowed
+    allowed_dirs: list[str] | None = None
+
+
+@dataclass
+class WatcherConfigItem:
+    """Configuration for a single watcher."""
+
+    name: str
+    enabled: bool = True
+    config: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class WatchersConfig:
+    """Configuration for watchers."""
+
+    enabled: bool = True
+    watchers: list[WatcherConfigItem] = field(default_factory=lambda: [
+        WatcherConfigItem(name="progress"),
+        WatcherConfigItem(name="budget"),
+        WatcherConfigItem(name="delegation_reminder"),
+    ])
+    # Role for watcher nudge messages: "user" | "assistant" | "system"
+    # "user" (default) - Appears as if user sent the message, strong nudge
+    # "assistant" - Appears as previous assistant thought, softer nudge
+    # "system" - Appears as system instruction, authoritative
+    message_role: str = "user"
+
+
+@dataclass
+class ZwarmConfig:
+    """Root configuration for zwarm."""
+
+    weave: WeaveConfig = field(default_factory=WeaveConfig)
+    executor: ExecutorConfig = field(default_factory=ExecutorConfig)
+    orchestrator: OrchestratorConfig = field(default_factory=OrchestratorConfig)
+    watchers: WatchersConfig = field(default_factory=WatchersConfig)
+    state_dir: str = ".zwarm"
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> ZwarmConfig:
+        """Create config from dictionary."""
+        weave_data = data.get("weave", {})
+        executor_data = data.get("executor", {})
+        orchestrator_data = data.get("orchestrator", {})
+        watchers_data = data.get("watchers", {})
+
+        # Parse compaction config from orchestrator
+        compaction_data = orchestrator_data.pop("compaction", {}) if orchestrator_data else {}
+        compaction_config = CompactionConfig(**compaction_data) if compaction_data else CompactionConfig()
+
+        # Parse watchers config - handle both list shorthand and dict format
+        if isinstance(watchers_data, list):
+            # Shorthand: watchers: [progress, budget, scope]
+            watchers_config = WatchersConfig(
+                enabled=True,
+                watchers=[
+                    WatcherConfigItem(name=w) if isinstance(w, str) else WatcherConfigItem(**w)
+                    for w in watchers_data
+                ],
+            )
+        else:
+            # Full format: watchers: {enabled: true, watchers: [...], message_role: "user"}
+            watchers_config = WatchersConfig(
+                enabled=watchers_data.get("enabled", True),
+                watchers=[
+                    WatcherConfigItem(name=w) if isinstance(w, str) else WatcherConfigItem(**w)
+                    for w in watchers_data.get("watchers", [])
+                ] or WatchersConfig().watchers,
+                message_role=watchers_data.get("message_role", "user"),
+            )
+
+        # Build orchestrator config with nested compaction
+        if orchestrator_data:
+            orchestrator_config = OrchestratorConfig(**orchestrator_data, compaction=compaction_config)
+        else:
+            orchestrator_config = OrchestratorConfig(compaction=compaction_config)
+
+        return cls(
+            weave=WeaveConfig(**weave_data) if weave_data else WeaveConfig(),
+            executor=ExecutorConfig(**executor_data) if executor_data else ExecutorConfig(),
+            orchestrator=orchestrator_config,
+            watchers=watchers_config,
+            state_dir=data.get("state_dir", ".zwarm"),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "weave": {
+                "project": self.weave.project,
+                "enabled": self.weave.enabled,
+            },
+            "executor": {
+                "adapter": self.executor.adapter,
+                "model": self.executor.model,
+                "sandbox": self.executor.sandbox,
+                "timeout": self.executor.timeout,
+                "reasoning_effort": self.executor.reasoning_effort,
+            },
+            "orchestrator": {
+                "lm": self.orchestrator.lm,
+                "prompt": self.orchestrator.prompt,
+                "tools": self.orchestrator.tools,
+                "max_steps": self.orchestrator.max_steps,
+                "parallel_delegations": self.orchestrator.parallel_delegations,
+                "sync_first": self.orchestrator.sync_first,
+                "compaction": {
+                    "enabled": self.orchestrator.compaction.enabled,
+                    "max_tokens": self.orchestrator.compaction.max_tokens,
+                    "threshold_pct": self.orchestrator.compaction.threshold_pct,
+                    "target_pct": self.orchestrator.compaction.target_pct,
+                    "keep_first_n": self.orchestrator.compaction.keep_first_n,
+                    "keep_last_n": self.orchestrator.compaction.keep_last_n,
+                },
+            },
+            "watchers": {
+                "enabled": self.watchers.enabled,
+                "watchers": [
+                    {"name": w.name, "enabled": w.enabled, "config": w.config}
+                    for w in self.watchers.watchers
+                ],
+                "message_role": self.watchers.message_role,
+            },
+            "state_dir": self.state_dir,
+        }
+
+
+def load_env(path: Path | None = None) -> None:
+    """Load .env file if it exists."""
+    if path is None:
+        path = Path.cwd() / ".env"
+    if path.exists():
+        load_dotenv(path)
+
+
+def load_toml_config(path: Path | None = None) -> dict[str, Any]:
+    """
+    Load config.toml file.
+
+    Search order:
+    1. Explicit path (if provided)
+    2. .zwarm/config.toml (new standard location)
+    3. config.toml (legacy location for backwards compat)
+    """
+    if path is None:
+        # Try new location first
+        new_path = Path.cwd() / ".zwarm" / "config.toml"
+        legacy_path = Path.cwd() / "config.toml"
+        if new_path.exists():
+            path = new_path
+        elif legacy_path.exists():
+            path = legacy_path
+        else:
+            return {}
+    if not path.exists():
+        return {}
+    with open(path, "rb") as f:
+        return tomllib.load(f)
+
+
+def load_yaml_config(path: Path) -> dict[str, Any]:
+    """
+    Load YAML config with inheritance support.
+
+    Supports 'extends: path/to/base.yaml' for composition.
+    """
+    if not path.exists():
+        raise FileNotFoundError(f"Config not found: {path}")
+
+    with open(path) as f:
+        data = yaml.safe_load(f) or {}
+
+    # Handle inheritance
+    extends = data.pop("extends", None)
+    if extends:
+        base_path = (path.parent / extends).resolve()
+        base_data = load_yaml_config(base_path)
+        # Deep merge: data overrides base
+        data = deep_merge(base_data, data)
+
+    return data
+
+
+def deep_merge(base: dict, override: dict) -> dict:
+    """Deep merge two dicts, with override taking precedence."""
+    result = base.copy()
+    for key, value in override.items():
+        if key in result and isinstance(result[key], dict) and isinstance(value, dict):
+            result[key] = deep_merge(result[key], value)
+        else:
+            result[key] = value
+    return result
+
+
+def apply_overrides(config: dict[str, Any], overrides: list[str]) -> dict[str, Any]:
+    """
+    Apply CLI overrides in format 'key.path=value'.
+
+    Example: 'orchestrator.lm=claude-sonnet' sets config['orchestrator']['lm'] = 'claude-sonnet'
+    """
+    result = config.copy()
+    for override in overrides:
+        if "=" not in override:
+            continue
+        key_path, value = override.split("=", 1)
+        keys = key_path.split(".")
+
+        # Parse value (try int, float, bool, then string)
+        parsed_value: Any = value
+        if value.lower() == "true":
+            parsed_value = True
+        elif value.lower() == "false":
+            parsed_value = False
+        else:
+            try:
+                parsed_value = int(value)
+            except ValueError:
+                try:
+                    parsed_value = float(value)
+                except ValueError:
+                    pass  # Keep as string
+
+        # Navigate and set
+        target = result
+        for key in keys[:-1]:
+            if key not in target:
+                target[key] = {}
+            target = target[key]
+        target[keys[-1]] = parsed_value
+
+    return result
+
+
+def load_config(
+    config_path: Path | None = None,
+    toml_path: Path | None = None,
+    env_path: Path | None = None,
+    overrides: list[str] | None = None,
+) -> ZwarmConfig:
+    """
+    Load configuration with full precedence chain:
+    1. Defaults (in dataclasses)
+    2. config.toml (user settings)
+    3. YAML config file (if provided)
+    4. CLI overrides (--set key=value)
+    5. Environment variables (for secrets)
+    """
+    # Load .env first (for secrets)
+    load_env(env_path)
+
+    # Start with defaults
+    config_dict: dict[str, Any] = {}
+
+    # Layer in config.toml
+    toml_config = load_toml_config(toml_path)
+    if toml_config:
+        config_dict = deep_merge(config_dict, toml_config)
+
+    # Layer in YAML config
+    if config_path and config_path.exists():
+        yaml_config = load_yaml_config(config_path)
+        config_dict = deep_merge(config_dict, yaml_config)
+
+    # Apply CLI overrides
+    if overrides:
+        config_dict = apply_overrides(config_dict, overrides)
+
+    # Apply environment variables for weave
+    if os.getenv("WEAVE_PROJECT"):
+        if "weave" not in config_dict:
+            config_dict["weave"] = {}
+        config_dict["weave"]["project"] = os.getenv("WEAVE_PROJECT")
+
+    return ZwarmConfig.from_dict(config_dict)