polymath-agent 0.4.0__py3-none-any.whl

polymath/tools.py

@@ -0,0 +1,333 @@
+"""
+Tool (function) support for polymath models.
+Provides a registry of tools that models can call to interact with the environment.
+"""
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import subprocess
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional
+
+from polymath.config import CONFIG_DIR
+
+
+# ── Scope ─────────────────────────────────────────────────────────────────────
+#
+# Reading is not free. The model can put whatever it reads into its output,
+# and that output goes to the next model, the transcript, and the brain. So a
+# read outside the project is an escalation and has to be asked about, and the
+# directory holding the API keys is refused however the call arrives.
+
+
+#: Markers that identify a real project root, in preference order.
+_PROJECT_MARKERS = (".polymath", ".git")
+
+#: Paths that are refused wherever they sit, whatever the project root is.
+#: Position cannot launder these: a project root that happens to contain
+#: your ssh keys does not make reading them ordinary. Relative to home
+#: unless absolute.
+_DENIED_HOME_PATHS = (
+    ".ssh", ".aws", ".gnupg", ".config/gh", ".config/gcloud",
+    ".kube", ".docker/config.json", ".netrc", ".npmrc", ".pypirc",
+    ".git-credentials", ".bash_history", ".zsh_history",
+)
+
+#: File names that carry secrets wherever they appear, including inside
+#: the project root.
+_DENIED_NAMES = (".env", ".netrc", ".git-credentials")
+_DENIED_PREFIXES = ("id_rsa", "id_ed25519", "id_ecdsa", "id_dsa")
+
+
+def workspace_root() -> Path | None:
+    """The project the user is working in, or None if there isn't one.
+
+    Anchored to a real project marker rather than to the current directory.
+    Polymath is a REPL you launch from anywhere, and launching it from home
+    used to make the whole home directory the project, so ssh keys and .env
+    files counted as ordinary in-project reads.
+
+    None means no free-read zone: every read is asked about. That is the
+    right default for "I started this in a directory that is not a project".
+    """
+    start = Path.cwd().resolve()
+    for directory in [start, *start.parents]:
+        if directory == Path.home():
+            # Home is never a project root, even with a dotfiles repo in it.
+            break
+        for marker in _PROJECT_MARKERS:
+            if (directory / marker).exists():
+                return directory
+    return None
+
+
+def is_denied(path: str) -> bool:
+    """True for paths nothing the model does should reach, wherever they are."""
+    try:
+        target = resolve_path(path)
+    except (OSError, ValueError):
+        return True
+    if target.name in _DENIED_NAMES:
+        return True
+    if any(target.name.startswith(prefix) for prefix in _DENIED_PREFIXES):
+        return True
+    home = Path.home().resolve()
+    for relative in _DENIED_HOME_PATHS:
+        candidate = (home / relative)
+        try:
+            if target == candidate or target.is_relative_to(candidate):
+                return True
+        except (OSError, ValueError):
+            continue
+    return False
+
+
+def resolve_path(path: str) -> Path:
+    """Absolute, symlink-resolved. Resolving matters: a link inside the
+    workspace can point anywhere, and a string prefix test would not notice."""
+    return Path(path).expanduser().resolve()
+
+
+def is_inside_workspace(path: str) -> bool:
+    """True only when there is a project root and the path resolves inside it."""
+    root = workspace_root()
+    if root is None:
+        return False
+    try:
+        return resolve_path(path).is_relative_to(root)
+    except (OSError, ValueError):
+        return False
+
+
+def is_protected(path: str) -> bool:
+    """True for paths nothing the model does has any business touching.
+
+    ~/.polymath holds the API keys for every provider. Reading it once puts
+    them in the transcript, which is then sent to the next model.
+    """
+    try:
+        if resolve_path(path).is_relative_to(CONFIG_DIR.resolve()):
+            return True
+    except (OSError, ValueError):
+        return False
+    return is_denied(path)
+
+
+#: Arguments naming a filesystem path, per tool.
+_PATH_ARGUMENT = {
+    "read_file": "file_path",
+    "write_file": "file_path",
+    "list_directory": "path",
+}
+
+
+def permission_scope(tool_name: str, arguments: Dict[str, Any] | None) -> str:
+    """Why this particular call needs asking about, or "" if it does not.
+
+    This is per-call, not per-tool: reading inside the project is ordinary,
+    reading your ssh key is not, and both are `read_file`.
+    """
+    argument = _PATH_ARGUMENT.get(tool_name)
+    if not argument:
+        return ""
+    path = (arguments or {}).get(argument)
+    if not isinstance(path, str) or not path:
+        return ""
+    if is_protected(path):
+        return "a credentials path, which the model is never given"
+    if not is_inside_workspace(path):
+        if workspace_root() is None:
+            return "outside any project — polymath was not started in one"
+        return "outside the current project"
+    return ""
+
+
+@dataclass
+class Tool:
+    name: str
+    description: str
+    parameters: Dict[str, Any]  # JSON schema for arguments
+    callable: Callable[..., Any]
+    requires_permission: bool = False
+
+    def needs_permission(self, arguments: Dict[str, Any] | None = None) -> bool:
+        """Whether this call must be approved. Tools that always need consent
+        say so; the rest need it only when the call leaves the project."""
+        if self.requires_permission:
+            return True
+        return bool(permission_scope(self.name, arguments))
+
+    def to_json_schema(self) -> Dict[str, Any]:
+        """Convert to standard JSON schema format."""
+        return {
+            "name": self.name,
+            "description": self.description,
+            "parameters": self.parameters,
+        }
+
+
+class ToolRegistry:
+    def __init__(self):
+        self._tools: Dict[str, Tool] = {}
+
+    def register(self, tool: Tool):
+        self._tools[tool.name] = tool
+
+    def get(self, name: str) -> Optional[Tool]:
+        return self._tools.get(name)
+
+    def list_tools(self) -> List[Tool]:
+        return list(self._tools.values())
+
+    async def execute(self, name: str, args: Dict[str, Any]) -> str:
+        """Execute tool and return result as string."""
+        tool = self.get(name)
+        if not tool:
+            return f"Error: Tool '{name}' not found."
+
+        try:
+            # Handle sync call for now; can be expanded for async if needed
+            result = tool.callable(**args)
+            if result is None:
+                return "Success"
+            if isinstance(result, (dict, list)):
+                return json.dumps(result, indent=2)
+            return str(result)
+        except Exception as e:
+            return f"Error executing '{name}': {str(e)}"
+
+
+# ── Built-in Tools ────────────────────────────────────────────────────────────
+
+def list_directory(path: str = ".") -> List[str]:
+    """Lists files and directories in a given path."""
+    if is_protected(path):
+        return [f"Error: refused — {path} is inside polymath's config directory."]
+    return os.listdir(path)
+
+
+def read_file(file_path: str) -> str:
+    """Reads the content of a file."""
+    if is_protected(file_path):
+        return (
+            f"Error: refused — {file_path} is inside polymath's config "
+            "directory, which holds your API keys."
+        )
+    with open(file_path, "r") as f:
+        return f.read()
+
+
+def write_file(file_path: str, content: str) -> str:
+    """Writes content to a file (restricted to current workspace).
+
+    The path is resolved before the test, so a symlink inside the project
+    that points elsewhere is refused. The previous version compared the
+    textual absolute path, which normalises ".." but follows no links, so
+    `project/link.txt -> /etc/passwd` passed the check.
+    """
+    if is_protected(file_path):
+        return (
+            f"Error: refused — {file_path} is inside polymath's config "
+            "directory, which holds your API keys."
+        )
+    root = workspace_root()
+    target = resolve_path(file_path)
+    if root is None:
+        return (
+            "Error: write_file needs a project. polymath was started outside "
+            "any .polymath folder or git repository, so there is nowhere it "
+            "is safe to write."
+        )
+    if target != root and not target.is_relative_to(root):
+        return (
+            f"Error: write_file is restricted to current workspace ({root}). "
+            f"Refused: {target}"
+        )
+    target.parent.mkdir(parents=True, exist_ok=True)
+    target.write_text(content)
+    return f"File written to {file_path}"
+
+
+def run_shell_command(command: str) -> str:
+    """Runs a shell command and returns stdout + stderr."""
+    import shlex
+    try:
+        args = shlex.split(command)
+    except ValueError as e:
+        return f"Error: Invalid command syntax: {e}"
+    try:
+        result = subprocess.run(
+            args,
+            shell=False,
+            capture_output=True,
+            text=True,
+            timeout=30
+        )
+        output = result.stdout
+        if result.stderr:
+            output += f"\nErrors:\n{result.stderr}"
+        return output
+    except Exception as e:
+        return str(e)
+
+
+# ── Global Registry ───────────────────────────────────────────────────────────
+
+registry = ToolRegistry()
+
+registry.register(Tool(
+    name="list_directory",
+    description="List files and directories in a given path.",
+    parameters={
+        "type": "object",
+        "properties": {
+            "path": {"type": "string", "description": "The directory path (default: '.')"}
+        },
+    },
+    callable=list_directory
+))
+
+registry.register(Tool(
+    name="read_file",
+    description="Read the contents of a file.",
+    parameters={
+        "type": "object",
+        "properties": {
+            "file_path": {"type": "string", "description": "The path to the file to read."}
+        },
+        "required": ["file_path"]
+    },
+    callable=read_file
+))
+
+registry.register(Tool(
+    name="write_file",
+    description="Write content to a file (creates parent directories if needed).",
+    parameters={
+        "type": "object",
+        "properties": {
+            "file_path": {"type": "string", "description": "The path to write the file to."},
+            "content": {"type": "string", "description": "The content to write."}
+        },
+        "required": ["file_path", "content"]
+    },
+    callable=write_file,
+    requires_permission=True
+))
+
+registry.register(Tool(
+    name="run_shell_command",
+    description="Execute a shell command in the local environment.",
+    parameters={
+        "type": "object",
+        "properties": {
+            "command": {"type": "string", "description": "The command to run."}
+        },
+        "required": ["command"]
+    },
+    callable=run_shell_command,
+    requires_permission=True
+))

polymath/ui_state.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from collections import deque
+from dataclasses import dataclass
+
+from rich.panel import Panel
+from rich.table import Table
+
+from polymath.domain import Project
+
+
+def prompt_toolbar(
+    project: Project,
+    queued_inputs: deque[str | None],
+    active: bool,
+    session_id: str,
+) -> str:
+    queue_preview = [q for q in list(queued_inputs)[-2:] if q]
+    queue_text = " | ".join(queue_preview)[:120] if queue_preview else "empty"
+    repo_name = project.repo_root.name if project.repo_root else "no-repo"
+    branch = project.repo_branch or "-"
+    dirty = "dirty" if project.repo_dirty else "clean"
+    mode = "busy" if active else "idle"
+    source = "shared" if project.shared else ("local" if project.name else "none")
+    return (
+        f" session {session_id} | project {project.name or '-'} ({source}) | "
+        f"git {repo_name}:{branch}:{dirty} | queue {len(queued_inputs)} [{queue_text}] | {mode} "
+    )
+
+
+def render_queue_table(queued_inputs: list[str]) -> Table:
+    table = Table(title="Queued Commands", show_header=True, header_style="bold cyan")
+    table.add_column("#", style="dim", width=4)
+    table.add_column("Command", style="white")
+    for idx, item in enumerate(queued_inputs, 1):
+        table.add_row(str(idx), item)
+    return table
+
+
+@dataclass
+class RunStatus:
+    phase: str
+    model: str
+    note: str = ""
+    reviewer: str = ""
+    queue_count: int = 0
+    context_line: str = ""
+
+
+def render_run_status(status: RunStatus) -> Panel:
+    lines = [
+        f"[bold cyan]≡ Running Agent[/bold cyan]",
+        f"[dim]phase:[/dim] {status.phase}    [dim]model:[/dim] {status.model}    [dim]queue:[/dim] {status.queue_count}",
+    ]
+    if status.reviewer:
+        lines.append(f"[dim]reviewer:[/dim] {status.reviewer}")
+    if status.context_line:
+        lines.append(f"[dim]{status.context_line}[/dim]")
+    if status.note:
+        lines.append(status.note)
+    return Panel("\n".join(lines), border_style="cyan")
+
+
+@dataclass
+class CompletionSummary:
+    session_id: str
+    primary_model: str
+    reviewer_model: str = ""
+    token_line: str = ""
+    queue_count: int = 0
+
+
+def render_completion_summary(summary: CompletionSummary) -> Panel:
+    lines = [
+        "[bold green]✓ Agent Complete[/bold green]",
+        f"[dim]session:[/dim] {summary.session_id}",
+        f"[dim]primary:[/dim] {summary.primary_model}",
+    ]
+    if summary.reviewer_model:
+        lines.append(f"[dim]reviewer:[/dim] {summary.reviewer_model}")
+    if summary.token_line:
+        lines.append(f"[dim]{summary.token_line}[/dim]")
+    lines.append(f"[dim]queue remaining:[/dim] {summary.queue_count}")
+    return Panel("\n".join(lines), border_style="green")

polymath/workspace.py

@@ -0,0 +1,66 @@
+"""
+Workspace awareness for polymath.
+Scans the current directory to provide context to the models.
+"""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+
+def scan_workspace(root_dir: str = ".") -> Dict[str, Any]:
+    """Scans the workspace for structure and key files."""
+    root = Path(root_dir).resolve()
+    
+    # 1. Directory Structure (shallow)
+    structure = []
+    try:
+        for item in root.iterdir():
+            if item.name.startswith(".") or item.name == "__pycache__":
+                continue
+            if item.is_dir():
+                structure.append(f"  [DIR]  {item.name}/")
+            else:
+                structure.append(f"  [FILE] {item.name}")
+    except Exception:
+        pass
+
+    # 2. Key Files
+    key_files = ["README.md", "pyproject.toml", "package.json", "requirements.txt", "setup.py"]
+    file_contents = {}
+    for kf in key_files:
+        p = root / kf
+        if p.exists() and p.is_file():
+            try:
+                # Read first 2KB of each key file
+                with open(p, "r") as f:
+                    file_contents[kf] = f.read(2048)
+            except Exception:
+                pass
+
+    return {
+        "root": str(root),
+        "structure": structure,
+        "key_files": file_contents
+    }
+
+
+def get_workspace_context(root_dir: str = ".") -> str:
+    """Formats workspace scan into a string for system prompts."""
+    data = scan_workspace(root_dir)
+    
+    lines = [
+        f"CURRENT WORKSPACE: {data['root']}",
+        "DIRECTORY STRUCTURE:",
+    ]
+    lines.extend(data["structure"][:30]) # Limit to 30 items
+    
+    if data["key_files"]:
+        lines.append("\nKEY FILE CONTENTS:")
+        for name, content in data["key_files"].items():
+            lines.append(f"--- {name} ---")
+            lines.append(content)
+            lines.append("-" * (len(name) + 8))
+            
+    return "\n".join(lines)