polymath-agent 0.4.0__py3-none-any.whl

polymath/subagents.py

@@ -0,0 +1,486 @@
+"""
+Named post-processing subagents for polymath.
+
+A subagent is a lightweight mini-pipeline that runs *after* the main pipeline
+output is ready.  Each subagent has a focused system prompt and optional
+iterative loop (e.g. verify runs tests, fixes, re-runs).
+
+Built-in subagents
+──────────────────
+  code-simplifier   Simplify / refactor after main work
+  verify            Detect test runner, run tests, iterate on failures (max 3×)
+  security-scan     Flag security issues (injection, hardcoded creds, etc.)
+  test-writer       Generate test cases for the produced code
+  docs-writer       Generate or update inline documentation
+  brainstorm-critic Challenge assumptions, surface edge cases
+  perf-reviewer     Spot performance hotspots in code output
+  diff-explainer    Plain-language explanation of what changed and why
+
+Custom subagents
+────────────────
+  Place .md files in .polymath/subagents/ (project) or ~/.polymath/subagents/ (global).
+  Frontmatter fields: name, description, max_iterations, requires_tool, trigger_on (comma list)
+  Body: the system prompt.
+"""
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable
+
+from polymath.config import CONFIG_DIR, TaskType
+
+
+@dataclass
+class SubagentConfig:
+    name: str
+    description: str
+    system_prompt: str
+    max_iterations: int = 1
+    requires_tool: str = ""   # skip subagent if this tool is unavailable
+    trigger_on: list[TaskType] = field(default_factory=list)
+    source: str = "builtin"
+
+
+@dataclass
+class SubagentResult:
+    name: str
+    output: str
+    iterations: int
+    success: bool
+
+
+# ── Built-in subagent definitions ─────────────────────────────────────────────
+
+BUILTIN_SUBAGENTS: dict[str, SubagentConfig] = {
+
+    "code-simplifier": SubagentConfig(
+        name="code-simplifier",
+        description="Simplify and clean code output after the main pipeline.",
+        system_prompt=(
+            "You are a code simplification expert. "
+            "Review the provided code output and:\n"
+            "1. Remove unnecessary complexity and dead code\n"
+            "2. Simplify logic without changing behaviour\n"
+            "3. Improve naming clarity\n"
+            "4. Output only the simplified code — no commentary."
+        ),
+        trigger_on=[TaskType.CODE],
+    ),
+
+    "verify": SubagentConfig(
+        name="verify",
+        description="Run the project's test suite and fix failures (up to 3 iterations).",
+        system_prompt=(
+            "You are a debugging agent. Tests have failed. "
+            "Analyse the test output, identify root causes, and provide a corrected "
+            "implementation that fixes all failures. "
+            "Output only the corrected code — no explanation."
+        ),
+        max_iterations=3,
+        requires_tool="run_shell_command",
+        trigger_on=[TaskType.CODE],
+    ),
+
+    "security-scan": SubagentConfig(
+        name="security-scan",
+        description="Scan output for security issues: injection, hardcoded secrets, OWASP top-10.",
+        system_prompt=(
+            "You are a security reviewer. Scan the following code for:\n"
+            "- Injection vulnerabilities (SQL, command, XSS)\n"
+            "- Hardcoded secrets or credentials\n"
+            "- Insecure deserialization\n"
+            "- Missing input validation\n"
+            "- Broken authentication patterns\n\n"
+            "Format: VERDICT: CLEAN | ISSUES_FOUND\n"
+            "ISSUES:\n- <issue> at <location>\n"
+            "SEVERITY: LOW | MEDIUM | HIGH for each issue."
+        ),
+        trigger_on=[TaskType.CODE],
+    ),
+
+    "test-writer": SubagentConfig(
+        name="test-writer",
+        description="Generate comprehensive test cases for the produced code.",
+        system_prompt=(
+            "You are a test engineering expert. "
+            "Given the following code, write comprehensive tests covering:\n"
+            "1. Happy path\n"
+            "2. Edge cases and boundary values\n"
+            "3. Error conditions\n\n"
+            "Match the test framework already used in the project if detectable. "
+            "Output only the test code."
+        ),
+        trigger_on=[TaskType.CODE],
+    ),
+
+    "docs-writer": SubagentConfig(
+        name="docs-writer",
+        description="Generate or update inline documentation for the output.",
+        system_prompt=(
+            "You are a technical writer. "
+            "Add clear, concise docstrings/comments to the following code. "
+            "Document: purpose, parameters, return values, side effects, and "
+            "any non-obvious logic. Do not change the code itself. "
+            "Output the fully-documented version."
+        ),
+        trigger_on=[TaskType.CODE],
+    ),
+
+    "brainstorm-critic": SubagentConfig(
+        name="brainstorm-critic",
+        description="Challenge assumptions and surface edge cases in any output.",
+        system_prompt=(
+            "You are a rigorous critic. Review the following output and:\n"
+            "1. Identify unstated assumptions that may not hold\n"
+            "2. Surface edge cases or failure modes not addressed\n"
+            "3. Point out logical inconsistencies\n"
+            "4. Suggest alternative approaches worth considering\n\n"
+            "Be constructive but direct. Format as a bulleted list."
+        ),
+        trigger_on=[TaskType.GENERAL, TaskType.ANALYSIS, TaskType.RESEARCH],
+    ),
+
+    "perf-reviewer": SubagentConfig(
+        name="perf-reviewer",
+        description="Identify performance hotspots and inefficiencies in code output.",
+        system_prompt=(
+            "You are a performance engineering expert. "
+            "Review the following code for:\n"
+            "1. O(n²) or worse loops that could be optimised\n"
+            "2. Unnecessary memory allocations\n"
+            "3. N+1 query patterns\n"
+            "4. Blocking I/O in async contexts\n"
+            "5. Missing caching opportunities\n\n"
+            "Format: list each issue with location and suggested fix."
+        ),
+        trigger_on=[TaskType.CODE],
+    ),
+
+    "diff-explainer": SubagentConfig(
+        name="diff-explainer",
+        description="Explain what changed and why in plain language (great for PRs).",
+        system_prompt=(
+            "You are a senior developer writing a PR description. "
+            "Given the following code or diff, explain in plain language:\n"
+            "1. What changed\n"
+            "2. Why it changed (the motivation)\n"
+            "3. Any risks or things reviewers should pay attention to\n\n"
+            "Keep it concise — 3–5 bullet points."
+        ),
+        trigger_on=[TaskType.CODE],
+    ),
+
+    "brain-gap": SubagentConfig(
+        name="brain-gap",
+        description=(
+            "After a session, propose what should be added to the project "
+            "brain (rules.md / decisions.md) so the same mistake doesn't "
+            "repeat. Boris-style failure-driven brain growth."
+        ),
+        system_prompt=(
+            "You are inspecting a finished session to find what the project "
+            "brain was MISSING that would have prevented mistakes or saved "
+            "the model time. Look for:\n"
+            "1. Corrections the user made to the model (e.g. 'no, do X instead')\n"
+            "2. Repeated clarification rounds that suggest missing context\n"
+            "3. Wrong assumptions the model made that a rule would fix\n"
+            "4. Decisions the user made that future sessions should respect\n\n"
+            "Output format — propose 0-5 entries, each tagged with the bucket:\n"
+            "  [rules] one-sentence imperative rule (e.g. 'never commit to main')\n"
+            "  [decisions] one-line decision + brief why\n"
+            "  [glossary] term — definition\n\n"
+            "If nothing should be added, output the single word: NONE.\n"
+            "Be conservative — only propose entries that are durably useful "
+            "across sessions, not session-specific facts."
+        ),
+        trigger_on=[TaskType.GENERAL, TaskType.CODE, TaskType.ANALYSIS, TaskType.RESEARCH],
+    ),
+}
+
+
+# ── File parser ───────────────────────────────────────────────────────────────
+
+def _parse_subagent_file(path: Path, source: str) -> SubagentConfig | None:
+    try:
+        text = path.read_text()
+    except Exception:
+        return None
+
+    # JSON format
+    if path.suffix == ".json":
+        try:
+            data = json.loads(text)
+            trigger_raw = data.get("trigger_on", [])
+            trigger_on = _parse_trigger_on(trigger_raw)
+            return SubagentConfig(
+                name=data.get("name", path.stem),
+                description=data.get("description", ""),
+                system_prompt=data.get("system_prompt", ""),
+                max_iterations=int(data.get("max_iterations", 1)),
+                requires_tool=data.get("requires_tool", ""),
+                trigger_on=trigger_on,
+                source=source,
+            )
+        except Exception:
+            return None
+
+    # Markdown with frontmatter
+    parts = re.split(r"^---\s*$", text, maxsplit=2, flags=re.MULTILINE)
+    if len(parts) < 3:
+        return SubagentConfig(
+            name=path.stem, description="", system_prompt=text.strip(), source=source
+        )
+
+    _, front_raw, body = parts
+    meta: dict[str, str] = {}
+    for line in front_raw.splitlines():
+        if ":" in line:
+            k, _, v = line.partition(":")
+            meta[k.strip()] = v.strip()
+
+    trigger_on = _parse_trigger_on(
+        [t.strip() for t in meta.get("trigger_on", "").split(",") if t.strip()]
+    )
+    return SubagentConfig(
+        name=meta.get("name", path.stem),
+        description=meta.get("description", ""),
+        system_prompt=body.strip(),
+        max_iterations=int(meta.get("max_iterations", 1)),
+        requires_tool=meta.get("requires_tool", ""),
+        trigger_on=trigger_on,
+        source=source,
+    )
+
+
+def _parse_trigger_on(raw: list[str]) -> list[TaskType]:
+    mapping = {t.value: t for t in TaskType}
+    result = []
+    for item in raw:
+        tt = mapping.get(item.lower())
+        if tt:
+            result.append(tt)
+    return result
+
+
+# ── Loaders ───────────────────────────────────────────────────────────────────
+
+def _load_from_dir(directory: Path | None, source: str) -> dict[str, SubagentConfig]:
+    if directory is None or not directory.exists():
+        return {}
+    result = {}
+    for f in sorted(directory.glob("*.md")) + sorted(directory.glob("*.json")):
+        cfg = _parse_subagent_file(f, source)
+        if cfg:
+            result[cfg.name] = cfg
+    return result
+
+
+def get_subagents(project_subagents_dir: Path | None = None) -> dict[str, SubagentConfig]:
+    """Merge: builtins < global < project (project wins on name conflicts)."""
+    merged: dict[str, SubagentConfig] = {}
+    merged.update(BUILTIN_SUBAGENTS)
+    merged.update(_load_from_dir(CONFIG_DIR / "subagents", "global"))
+    merged.update(_load_from_dir(project_subagents_dir, "project"))
+    return merged
+
+
+# ── Test runner detection ─────────────────────────────────────────────────────
+
+def _detect_test_runner(cwd: str = ".") -> tuple[str, str] | None:
+    """
+    Inspect workspace for test runner indicators.
+    Returns (runner_name, command) or None.
+    """
+    base = Path(cwd)
+
+    # pytest
+    for indicator in ("pytest.ini", "conftest.py", "setup.cfg"):
+        if (base / indicator).exists():
+            return ("pytest", "pytest -x --tb=short")
+    pyproject = base / "pyproject.toml"
+    if pyproject.exists() and "[tool.pytest" in pyproject.read_text():
+        return ("pytest", "pytest -x --tb=short")
+
+    # jest / vitest
+    pkg = base / "package.json"
+    if pkg.exists():
+        try:
+            data = json.loads(pkg.read_text())
+            scripts = data.get("scripts", {})
+            if "test" in scripts:
+                runner = "vitest" if "vitest" in scripts["test"] else "jest"
+                return (runner, "npm test -- --run")
+        except Exception:
+            pass
+
+    # cargo
+    if (base / "Cargo.toml").exists():
+        return ("cargo", "cargo test")
+
+    # go test
+    if list(base.glob("*.go")):
+        return ("go test", "go test ./...")
+
+    return None
+
+
+def _run_tests(command: str) -> tuple[str, int]:
+    """Run the detected test command and return its output and exit status.
+
+    Every test runner reports failure through its exit status, so that is
+    the signal to use. run_shell_command discards it, which is why this
+    keeps its own runner rather than calling the tool.
+
+    No shell, like everywhere else: the command is split into an argument
+    vector. That is also why the detected commands carry no redirection —
+    a "2>&1" would reach the program as a literal argument, and pytest
+    would report it as a missing test path on every run.
+    """
+    import shlex
+    import subprocess
+
+    try:
+        argv = shlex.split(command)
+    except ValueError as e:
+        return (f"Error: invalid test command: {e}", 1)
+    try:
+        result = subprocess.run(
+            argv, shell=False, capture_output=True, text=True, timeout=600,
+        )
+    except FileNotFoundError:
+        return (f"Error: test runner not found: {argv[0]}", 127)
+    except subprocess.TimeoutExpired:
+        return ("Error: the test run timed out after 600s.", 1)
+    except Exception as e:  # pragma: no cover - defensive
+        return (f"Error: {e}", 1)
+
+    output = result.stdout
+    if result.stderr:
+        output += ("\n" if output else "") + result.stderr
+    return (output, result.returncode)
+
+
+def _has_test_failures(output: str) -> bool:
+    """Fallback for callers holding only the text. Prefer the exit status:
+    every test runner reports failure through it, and these patterns match
+    any output mentioning a failure, including a test named for one."""
+    failure_patterns = [
+        r"FAILED",
+        r"failed",
+        r"ERROR",
+        r"\berror\[E",         # Rust
+        r"FAIL\b",
+        r"Test failed",
+        r"AssertionError",
+        r"not ok",             # TAP
+    ]
+    for pat in failure_patterns:
+        if re.search(pat, output):
+            return True
+    return False
+
+
+# ── Runner ────────────────────────────────────────────────────────────────────
+
+async def run_subagent(
+    config: SubagentConfig,
+    pipeline_output: str,
+    primary_adapter: Any,
+    primary_model: Any,
+    session_id: str,
+    on_token: Callable[[str], None] | None = None,
+) -> SubagentResult:
+    """
+    Execute a subagent against the pipeline output.
+    Single-pass for max_iterations=1; iterative test-fix loop for verify.
+    """
+    from polymath.adapters.base import Message
+    from polymath.tools import run_shell_command
+
+    _emit = on_token or (lambda t: None)
+
+    # ── Verify: iterative test-fix loop ───────────────────────────────────────
+    if config.name == "verify" or config.max_iterations > 1:
+        runner = _detect_test_runner()
+        if runner is None:
+            return SubagentResult(
+                name=config.name, output="No test runner detected — skipping verify.",
+                iterations=0, success=True,
+            )
+        runner_name, run_cmd = runner
+        current_code = pipeline_output
+        for iteration in range(1, config.max_iterations + 1):
+            test_output, exit_code = _run_tests(run_cmd)
+            if exit_code == 0:
+                return SubagentResult(
+                    name=config.name,
+                    output=f"All tests passed ({runner_name}). {iteration} iteration(s).",
+                    iterations=iteration,
+                    success=True,
+                )
+            # Ask model to fix
+            fix_prompt = (
+                f"Tests failed ({runner_name}):\n\n{test_output}\n\n"
+                f"Here is the current code:\n\n{current_code}\n\n"
+                f"Fix all test failures. Output only the corrected code."
+            )
+            resp = await primary_adapter.complete(
+                messages=[Message(role="user", content=fix_prompt)],
+                model_id=primary_model.id,
+                system=config.system_prompt,
+                temperature=0.3,
+            )
+            current_code = resp if isinstance(resp, str) else resp.content
+
+        # After max iterations — tests still failing
+        _, final_exit = _run_tests(run_cmd)
+        return SubagentResult(
+            name=config.name,
+            output=current_code,
+            iterations=config.max_iterations,
+            success=final_exit == 0,
+        )
+
+    # ── Single-pass subagent ──────────────────────────────────────────────────
+    prompt = f"{pipeline_output}"
+    output = ""
+    async for token in primary_adapter.stream(
+        messages=[Message(role="user", content=prompt)],
+        model_id=primary_model.id,
+        system=config.system_prompt,
+    ):
+        if isinstance(token, str):
+            output += token
+            _emit(token)
+
+    return SubagentResult(name=config.name, output=output, iterations=1, success=True)
+
+
+# ── Auto-selection ────────────────────────────────────────────────────────────
+
+def select_auto_subagents(
+    task_type: TaskType,
+    flags: dict[str, bool],
+    available: dict[str, SubagentConfig],
+) -> list[SubagentConfig]:
+    """
+    Decide which subagents to run.
+    flags: {"verify": bool, "simplify": bool, ...}
+    Explicit flags take priority; trigger_on handles auto-selection.
+    """
+    selected: list[SubagentConfig] = []
+    seen: set[str] = set()
+
+    # Explicit flags first (in a stable order)
+    explicit_order = ["verify", "simplify", "security-scan", "test-writer",
+                      "docs-writer", "brainstorm-critic", "perf-reviewer",
+                      "diff-explainer", "brain-gap"]
+    for name in explicit_order:
+        if flags.get(name) and name in available:
+            selected.append(available[name])
+            seen.add(name)
+
+    return selected