pdo-agent 2.0.0__py3-none-any.whl

pdo/mcp.py

@@ -0,0 +1,264 @@
+"""Minimal Model Context Protocol (MCP) client.
+
+Connects to MCP servers over the stdio transport (newline-delimited JSON-RPC
+2.0), lists their tools, and exposes each as a PDO :class:`~pdo.tools.base.Tool`
+so the agent can use them transparently. Implemented synchronously with the
+standard library only — no extra dependencies, consistent with PDO's v1 design.
+
+Servers are declared in ``<PDO_HOME>/mcp.json`` using the widely-used format::
+
+    {
+      "mcpServers": {
+        "filesystem": {
+          "command": "npx",
+          "args": ["-y", "@modelcontextprotocol/server-filesystem", "/path"]
+        }
+      }
+    }
+"""
+from __future__ import annotations
+
+import json
+import logging
+import os
+import re
+import select
+import subprocess
+import time
+from pathlib import Path
+from typing import Any
+
+from . import __version__
+from .tools.base import Tool, truncate
+from .tools.registry import ToolRegistry
+
+logger = logging.getLogger(__name__)
+
+PROTOCOL_VERSION = "2024-11-05"
+_REQUEST_TIMEOUT = 30
+_CALL_TIMEOUT = 120
+
+# Servers started this process, for introspection (e.g. the /mcp command).
+_ACTIVE_SERVERS: list[MCPServer] = []
+
+
+class MCPError(RuntimeError):
+    """Raised when an MCP server errors or can't be reached."""
+
+
+def _sanitize(name: str) -> str:
+    """Make a tool name safe for the model API (^[A-Za-z0-9_-]{1,64}$)."""
+    return re.sub(r"[^A-Za-z0-9_-]", "_", name)[:64]
+
+
+class MCPServer:
+    """A single MCP server subprocess, spoken to over stdio JSON-RPC."""
+
+    def __init__(
+        self, name: str, command: str, args: list[str], env: dict[str, str] | None = None
+    ) -> None:
+        self.name = name
+        self.command = command
+        self.args = list(args)
+        self.env = dict(env or {})
+        self._proc: subprocess.Popen | None = None
+        self._id = 0
+        self._tools: list[dict[str, Any]] = []
+
+    # --- lifecycle ---------------------------------------------------------- #
+    def start(self) -> None:
+        """Spawn the server, perform the MCP handshake, and cache its tools."""
+        self._proc = subprocess.Popen(  # noqa: S603 — launching a user-configured server
+            [self.command, *self.args],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.DEVNULL,
+            text=True,
+            bufsize=1,
+            env={**os.environ, **self.env},
+        )
+        self._request(
+            "initialize",
+            {
+                "protocolVersion": PROTOCOL_VERSION,
+                "capabilities": {},
+                "clientInfo": {"name": "pdo", "version": __version__},
+            },
+        )
+        self._notify("notifications/initialized")
+        self._tools = self._request("tools/list").get("tools", [])
+
+    def stop(self) -> None:
+        if self._proc is None:
+            return
+        try:
+            self._proc.terminate()
+            self._proc.wait(timeout=5)
+        except Exception:  # noqa: BLE001 — force-kill if it won't exit cleanly
+            self._proc.kill()
+        finally:
+            self._proc = None
+
+    # --- public API --------------------------------------------------------- #
+    def list_tools(self) -> list[dict[str, Any]]:
+        return list(self._tools)
+
+    def call_tool(self, name: str, arguments: dict[str, Any]) -> str:
+        result = self._request(
+            "tools/call", {"name": name, "arguments": arguments}, timeout=_CALL_TIMEOUT
+        )
+        parts = []
+        for item in result.get("content", []):
+            if item.get("type") == "text":
+                parts.append(item.get("text", ""))
+            else:
+                parts.append(json.dumps(item))
+        text = "\n".join(parts) if parts else "(no content)"
+        if result.get("isError"):
+            return f"Error from MCP tool: {truncate(text)}"
+        return truncate(text)
+
+    # --- JSON-RPC plumbing -------------------------------------------------- #
+    def _send(self, message: dict[str, Any]) -> None:
+        if self._proc is None or self._proc.stdin is None:
+            raise MCPError(f"MCP server {self.name!r} is not running")
+        self._proc.stdin.write(json.dumps(message) + "\n")
+        self._proc.stdin.flush()
+
+    def _notify(self, method: str, params: dict[str, Any] | None = None) -> None:
+        self._send({"jsonrpc": "2.0", "method": method, "params": params or {}})
+
+    def _request(
+        self, method: str, params: dict[str, Any] | None = None, timeout: int = _REQUEST_TIMEOUT
+    ) -> dict[str, Any]:
+        self._id += 1
+        request_id = self._id
+        self._send({"jsonrpc": "2.0", "id": request_id, "method": method, "params": params or {}})
+
+        deadline = time.time() + timeout
+        while time.time() < deadline:
+            line = self._read_line(deadline - time.time())
+            if not line:
+                continue
+            try:
+                message = json.loads(line)
+            except json.JSONDecodeError:
+                continue
+            if message.get("id") == request_id:
+                if "error" in message:
+                    raise MCPError(message["error"].get("message", "MCP error"))
+                return message.get("result", {})
+            # Otherwise it's a notification or an unrelated message: ignore it.
+        raise MCPError(f"timed out waiting for {method!r} from {self.name!r}")
+
+    def _read_line(self, timeout: float) -> str | None:
+        if self._proc is None or self._proc.stdout is None:
+            return None
+        try:
+            ready, _, _ = select.select([self._proc.stdout], [], [], max(0.0, timeout))
+        except (OSError, ValueError):
+            # select may not work on this platform's pipes; fall back to blocking.
+            return self._proc.stdout.readline() or None
+        if not ready:
+            return None
+        return self._proc.stdout.readline() or None
+
+
+class MCPTool(Tool):
+    """Adapts a remote MCP tool to PDO's Tool interface."""
+
+    def __init__(self, server: MCPServer, mcp_name: str, description: str, schema: dict | None):
+        self.name = _sanitize(f"mcp__{server.name}__{mcp_name}")
+        self.description = description or f"MCP tool {mcp_name} from {server.name}."
+        self.parameters = schema or {"type": "object", "properties": {}}
+        self._server = server
+        self._mcp_name = mcp_name
+
+    def run(self, **kwargs: Any) -> str:
+        try:
+            return self._server.call_tool(self._mcp_name, self._clean_args(kwargs))
+        except MCPError as exc:
+            return f"Error: {exc}"
+
+    def _clean_args(self, kwargs: dict[str, Any]) -> dict[str, Any]:
+        """Drop blank optional arguments.
+
+        Smaller models often pad optional parameters with junk like ``" "``,
+        ``"/"`` or ``[]``, which some MCP servers reject (e.g. Canva fails on a
+        meaningless ``brand_kit_id``). For any argument not required by the tool's
+        schema, we drop None, empty collections, and strings with no alphanumeric
+        characters (a real id/value always contains letters or digits).
+        """
+        required = set(self.parameters.get("required", []) or [])
+        cleaned: dict[str, Any] = {}
+        for key, value in kwargs.items():
+            if key in required:
+                cleaned[key] = value
+                continue
+            if value is None:
+                continue
+            if isinstance(value, str) and not any(ch.isalnum() for ch in value):
+                continue
+            if isinstance(value, (list, dict)) and len(value) == 0:
+                continue
+            cleaned[key] = value
+        return cleaned
+
+
+def load_mcp_config(path: Path) -> dict[str, dict]:
+    """Read the ``mcpServers`` mapping from ``path`` (empty if missing/invalid)."""
+    if not path.exists():
+        return {}
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+    except (json.JSONDecodeError, OSError) as exc:
+        logger.warning("Could not read MCP config %s: %s", path, exc)
+        return {}
+    servers = data.get("mcpServers", data) if isinstance(data, dict) else {}
+    return servers if isinstance(servers, dict) else {}
+
+
+def start_servers(
+    registry: ToolRegistry, servers_config: dict[str, dict]
+) -> tuple[list[MCPServer], list[tuple[str, int, str | None]]]:
+    """Start each configured server and register its tools.
+
+    Returns the started servers and a per-server ``(name, tool_count, error)``
+    summary. A server that fails to start is skipped, not fatal.
+    """
+    started: list[MCPServer] = []
+    summary: list[tuple[str, int, str | None]] = []
+
+    for name, spec in servers_config.items():
+        command = spec.get("command")
+        if not command:
+            summary.append((name, 0, "no 'command' specified"))
+            continue
+        server = MCPServer(name, command, spec.get("args", []), spec.get("env", {}))
+        try:
+            server.start()
+        except Exception as exc:  # noqa: BLE001 — a bad server must not crash PDO
+            logger.warning("MCP server %r failed to start: %s", name, exc)
+            summary.append((name, 0, str(exc)))
+            continue
+
+        count = 0
+        for spec_tool in server.list_tools():
+            tool = MCPTool(
+                server,
+                spec_tool["name"],
+                spec_tool.get("description", ""),
+                spec_tool.get("inputSchema"),
+            )
+            if not registry.has(tool.name):
+                registry.register(tool)
+                count += 1
+        started.append(server)
+        _ACTIVE_SERVERS.append(server)
+        summary.append((name, count, None))
+
+    return started, summary
+
+
+def active_servers() -> list[MCPServer]:
+    return list(_ACTIVE_SERVERS)

pdo/prompts/system.md

@@ -0,0 +1,46 @@
+You are **PDO** (Python Do) — a terminal-first AI agent. Your motto is **Think. Plan. Do.**
+
+You are not a chatbot. You help the user accomplish real work on their machine by
+reasoning about the goal, planning the steps, deciding whether tools are needed,
+executing them safely, reviewing the result, and replying clearly.
+
+## How you work
+
+1. **Understand** the user's goal before acting. Ask a brief clarifying question
+   only when genuinely blocked — otherwise proceed with sensible defaults.
+2. **Think and plan.** For multi-step tasks, work through the steps in order.
+3. **Decide if tools are needed.** Use a tool only when it makes the answer
+   better or is required to complete the task. Plain conversation, explanations,
+   and questions you can answer directly should stay a plain reply with **no
+   tool calls**.
+4. **Execute safely.** Prefer the smallest action that accomplishes the step.
+   Inspect before you change: read a file before editing it, list a directory
+   before assuming its contents.
+5. **Review** what came back. If a tool returned an error, explain it and adjust
+   rather than blindly retrying.
+6. **Respond clearly** in concise Markdown. Show commands and code in fenced
+   blocks. Summarise what you did and what (if anything) the user should do next.
+
+## Tools
+
+You have native function/tool calling. Available tools include reading, writing,
+and appending files; listing and creating directories; running shell commands;
+and saving, searching, and deleting long-term memories. Call tools by name with
+JSON arguments — never describe a tool call in prose instead of making it.
+
+## Safety
+
+- Destructive or privileged commands (`rm`, `sudo`, `shutdown`, `reboot`, disk
+  operations, recursive deletes) require explicit user confirmation, which the
+  tool layer enforces. Do not try to bypass it.
+- Writes are sandboxed to the working directory by default; writing elsewhere or
+  overwriting a file will prompt the user.
+- When you are about to do something irreversible, say so plainly first.
+
+## Style
+
+- Be direct and practical. Lead with the answer or the result.
+- Don't narrate your internal reasoning at length; show the outcome.
+- Use the user's working directory as the default location for new files.
+- When you remember something durable about the user or project, use the memory
+  tools so it persists across sessions.

pdo/providers.py

@@ -0,0 +1,86 @@
+"""Known LLM providers for the ``/models`` picker.
+
+All three providers expose an OpenAI-compatible chat API, so PDO reaches them
+through the same :class:`pdo.llm.OpenAIClient` — only the ``base_url`` and the
+API key differ. Adding a provider is just another entry in ``PROVIDERS``.
+
+The model lists are a curated starting point, not an exhaustive catalogue; the
+picker always offers a "custom model id" option so any model the provider
+supports can be used.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+
+@dataclass(frozen=True)
+class Provider:
+    """Connection details and suggested models for one provider."""
+
+    key: str  # short internal id, e.g. "openai"
+    label: str  # human label shown in the menu
+    env_key: str  # environment variable holding the API key
+    base_url: str | None  # OpenAI-compatible endpoint (None = OpenAI default)
+    models: list[str] = field(default_factory=list)
+
+
+PROVIDERS: dict[str, Provider] = {
+    "openai": Provider(
+        key="openai",
+        label="OpenAI",
+        env_key="OPENAI_API_KEY",
+        base_url=None,
+        models=[
+            "gpt-4.1-mini",
+            "gpt-4.1",
+            "gpt-4.1-nano",
+            "gpt-4o",
+            "gpt-4o-mini",
+            "o4-mini",
+        ],
+    ),
+    "anthropic": Provider(
+        key="anthropic",
+        label="Anthropic (Claude)",
+        env_key="ANTHROPIC_API_KEY",
+        # Anthropic's OpenAI-compatibility endpoint.
+        base_url="https://api.anthropic.com/v1/",
+        models=[
+            "claude-sonnet-4-5",
+            "claude-opus-4-1",
+            "claude-3-7-sonnet-latest",
+            "claude-3-5-sonnet-latest",
+            "claude-3-5-haiku-latest",
+        ],
+    ),
+    "openrouter": Provider(
+        key="openrouter",
+        label="OpenRouter",
+        env_key="OPENROUTER_API_KEY",
+        base_url="https://openrouter.ai/api/v1",
+        models=[
+            "openai/gpt-4.1-mini",
+            "anthropic/claude-3.7-sonnet",
+            "anthropic/claude-3.5-sonnet",
+            "google/gemini-flash-1.5",
+            "meta-llama/llama-3.3-70b-instruct",
+            "deepseek/deepseek-chat",
+        ],
+    ),
+    "ollama": Provider(
+        key="ollama",
+        label="Ollama (local)",
+        env_key="OLLAMA_API_KEY",  # unused; a local server needs no key
+        # OpenAI-compatible local endpoint. Override with OLLAMA_BASE_URL.
+        base_url="http://localhost:11434/v1",
+        models=[
+            "llama3.2",
+            "llama3.1",
+            "qwen2.5",
+            "mistral",
+            "gemma2",
+            "phi4",
+            "deepseek-r1",
+        ],
+    ),
+}

pdo/rag.py

@@ -0,0 +1,191 @@
+"""Lexical codebase retrieval (BM25).
+
+``/index`` builds a chunk index of the working directory; the
+``codebase_search`` tool ranks chunks against a query with BM25 and returns the
+best snippets with ``path:line`` references.
+
+We use lexical retrieval rather than embeddings on purpose: it needs no
+embeddings endpoint (OpenRouter doesn't offer one), no API key, and no new
+dependencies — and BM25 over code identifiers is strong in practice. Embedding
+support can be layered on later without changing the tool interface.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+import math
+import re
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from .config import get_home_dir
+
+logger = logging.getLogger(__name__)
+
+# Files worth indexing (source + docs + config).
+_INDEX_EXTS = {
+    ".py", ".md", ".txt", ".rst", ".toml", ".ini", ".cfg", ".json", ".yml",
+    ".yaml", ".js", ".jsx", ".ts", ".tsx", ".html", ".css", ".sh", ".zsh",
+    ".swift", ".rs", ".go", ".java", ".kt", ".c", ".h", ".cpp", ".hpp", ".sql",
+}
+_SKIP_DIRS = {
+    ".git", ".venv", "venv", "node_modules", "__pycache__", ".ruff_cache",
+    ".pytest_cache", "dist", "build", ".mypy_cache", ".idea", ".vscode",
+}
+_MAX_FILE_BYTES = 200_000
+_CHUNK_LINES = 40
+_CHUNK_OVERLAP = 10
+_MAX_CHUNKS = 20_000
+
+# BM25 constants (standard defaults).
+_K1 = 1.5
+_B = 0.75
+
+_WORD = re.compile(r"[A-Za-z0-9_]+")
+# Split camelCase and acronym boundaries: fooBar -> foo|Bar, HTTPServer -> HTTP|Server.
+_CAMEL = re.compile(r"(?<=[a-z0-9])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])")
+
+
+def tokenize(text: str) -> list[str]:
+    """Split text into lowercase tokens, including snake/camelCase sub-words."""
+    tokens: list[str] = []
+    for word in _WORD.findall(text):
+        lowered = word.lower()
+        tokens.append(lowered)
+        parts = [p for chunk in _CAMEL.split(word) for p in chunk.split("_") if p]
+        if len(parts) > 1:
+            tokens.extend(p.lower() for p in parts)
+    return tokens
+
+
+@dataclass
+class Chunk:
+    path: str  # relative to the index root
+    start: int  # 1-based first line
+    end: int  # 1-based last line
+    text: str
+
+
+@dataclass
+class Index:
+    root: str
+    built: float
+    chunks: list[Chunk] = field(default_factory=list)
+
+
+def _index_path(root: Path) -> Path:
+    digest = hashlib.sha1(str(root.resolve()).encode()).hexdigest()[:16]
+    directory = get_home_dir() / "index"
+    directory.mkdir(parents=True, exist_ok=True)
+    return directory / f"{digest}.json"
+
+
+def _iter_files(root: Path):
+    for path in sorted(root.rglob("*")):
+        if not path.is_file() or path.suffix.lower() not in _INDEX_EXTS:
+            continue
+        if any(part in _SKIP_DIRS for part in path.parts):
+            continue
+        try:
+            if path.stat().st_size > _MAX_FILE_BYTES:
+                continue
+        except OSError:
+            continue
+        yield path
+
+
+def build_index(root: Path) -> Index:
+    """Chunk every indexable file under ``root`` and persist the index."""
+    root = root.resolve()
+    index = Index(root=str(root), built=time.time())
+    for file in _iter_files(root):
+        try:
+            lines = file.read_text("utf-8", "ignore").splitlines()
+        except OSError:
+            continue
+        rel = str(file.relative_to(root))
+        step = _CHUNK_LINES - _CHUNK_OVERLAP
+        for start in range(0, max(len(lines), 1), step):
+            block = lines[start : start + _CHUNK_LINES]
+            if not any(line.strip() for line in block):
+                continue
+            index.chunks.append(
+                Chunk(path=rel, start=start + 1, end=start + len(block), text="\n".join(block))
+            )
+            if len(index.chunks) >= _MAX_CHUNKS:
+                logger.warning("Index chunk cap reached (%d); stopping", _MAX_CHUNKS)
+                save_index(index)
+                return index
+            if start + _CHUNK_LINES >= len(lines):
+                break
+    save_index(index)
+    return index
+
+
+def save_index(index: Index) -> None:
+    payload = {
+        "root": index.root,
+        "built": index.built,
+        "chunks": [vars(chunk) for chunk in index.chunks],
+    }
+    _index_path(Path(index.root)).write_text(json.dumps(payload), encoding="utf-8")
+
+
+def load_index(root: Path) -> Index | None:
+    path = _index_path(root)
+    if not path.exists():
+        return None
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+        return Index(
+            root=data["root"],
+            built=data["built"],
+            chunks=[Chunk(**chunk) for chunk in data["chunks"]],
+        )
+    except (json.JSONDecodeError, KeyError, TypeError, OSError) as exc:
+        logger.warning("Could not load index %s: %s", path, exc)
+        return None
+
+
+@dataclass
+class SearchResult:
+    chunk: Chunk
+    score: float
+
+
+def search(index: Index, query: str, top_k: int = 5) -> list[SearchResult]:
+    """Rank chunks against ``query`` with BM25 and return the top ``top_k``."""
+    query_tokens = set(tokenize(query))
+    if not query_tokens or not index.chunks:
+        return []
+
+    token_lists = [tokenize(chunk.text) for chunk in index.chunks]
+    n = len(token_lists)
+    avg_len = sum(len(t) for t in token_lists) / n
+
+    # Document frequency per query token.
+    df = dict.fromkeys(query_tokens, 0)
+    for tokens in token_lists:
+        present = query_tokens.intersection(tokens)
+        for token in present:
+            df[token] += 1
+
+    results: list[SearchResult] = []
+    for chunk, tokens in zip(index.chunks, token_lists, strict=True):
+        if not tokens:
+            continue
+        score = 0.0
+        length_norm = _K1 * (1 - _B + _B * len(tokens) / avg_len)
+        for token in query_tokens:
+            tf = tokens.count(token)
+            if tf == 0 or df[token] == 0:
+                continue
+            idf = math.log(1 + (n - df[token] + 0.5) / (df[token] + 0.5))
+            score += idf * (tf * (_K1 + 1)) / (tf + length_norm)
+        if score > 0:
+            results.append(SearchResult(chunk=chunk, score=score))
+
+    results.sort(key=lambda r: r.score, reverse=True)
+    return results[:top_k]