co-agent-coder 0.1.0__py3-none-any.whl

coagent/llm/base.py

@@ -0,0 +1,235 @@
+"""The model interface (DESIGN §4).
+
+`generate(messages) -> text` behind one interface, so the loop never knows which
+backend it is talking to. That is what keeps the loop testable on a laptop with
+a scripted backend while production runs a 4-bit Qwen on a T4.
+
+Two rules this module enforces because a 7B model cannot be trusted with them:
+
+* **No silent truncation.** A request over the token budget is rejected before
+  generation (PRD D2), counted with the model's own tokenizer.
+* **No inference egress.** Nothing in this layer opens a socket. The only
+  outbound traffic in the product is the model download and git (FR-MDL-7).
+"""
+
+from __future__ import annotations
+
+import enum
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass, field
+from typing import Any
+
+__all__ = [
+    "BudgetError",
+    "FakeBackend",
+    "GenerationConfig",
+    "GenerationResult",
+    "LLMBackend",
+    "Message",
+    "StopReason",
+    "TokenBudget",
+    "Usage",
+]
+
+
+class BudgetError(RuntimeError):
+    """The request exceeds the context budget. Raised before generation."""
+
+
+class StopReason(str, enum.Enum):
+    STOP = "stop"  # model emitted EOS
+    LENGTH = "length"  # hit max_new_tokens (may be truncated — say so)
+    CANCELLED = "cancelled"  # user pressed stop
+    ERROR = "error"
+
+
+@dataclass(slots=True)
+class Message:
+    role: str  # "system" | "user" | "assistant"
+    content: str
+
+    def to_dict(self) -> dict[str, str]:
+        return {"role": self.role, "content": self.content}
+
+
+@dataclass(slots=True)
+class GenerationConfig:
+    max_new_tokens: int = 1024
+    temperature: float = 0.2
+    top_p: float = 0.95
+    seed: int = 0
+    stop: tuple[str, ...] = ()
+
+
+@dataclass(slots=True)
+class Usage:
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+
+    @property
+    def total(self) -> int:
+        return self.prompt_tokens + self.completion_tokens
+
+
+@dataclass(slots=True)
+class GenerationResult:
+    text: str
+    stop_reason: StopReason = StopReason.STOP
+    usage: Usage = field(default_factory=Usage)
+    model: str = ""
+
+    @property
+    def truncated(self) -> bool:
+        return self.stop_reason is StopReason.LENGTH
+
+
+@dataclass(slots=True)
+class TokenBudget:
+    """Context arithmetic against a real tokenizer (PRD D2, FR-CTX-6).
+
+    `count` is injected so this type works with HF tokenizers, a fake counter in
+    tests, and a character heuristic only as an explicit last resort.
+    """
+
+    max_context: int
+    count: Callable[[str], int]
+    reserve_output: int = 1500
+
+    @property
+    def input_limit(self) -> int:
+        return max(0, self.max_context - self.reserve_output)
+
+    def count_messages(self, messages: Iterable[Message]) -> int:
+        # Chat templates add role markers and separators; ~4 tokens per turn is
+        # the usual overhead and is included so the estimate is not optimistic.
+        return sum(self.count(m.content) for m in messages) + 4 * len(list(messages))
+
+    def check(self, messages: list[Message], *, max_new_tokens: int) -> int:
+        """Return the input token count, or raise `BudgetError`."""
+        used = self.count_messages(messages)
+        cap = self.max_context - min(max_new_tokens, self.reserve_output)
+        if used > cap:
+            raise BudgetError(
+                f"request is {used} tokens but the budget allows {cap} "
+                f"(max_context={self.max_context}, reserved output={min(max_new_tokens, self.reserve_output)}). "
+                "Evict observations or shorten the retrieval set — the request was not sent."
+            )
+        return used
+
+
+class LLMBackend(ABC):
+    """What the loop requires of any backend."""
+
+    name: str = ""
+
+    @property
+    @abstractmethod
+    def model_id(self) -> str:
+        """Human-readable `model@quant` used in trailers and the status line."""
+
+    @abstractmethod
+    def count_tokens(self, text: str) -> int:
+        """Tokenizer truth for budget arithmetic (FR-CTX-6)."""
+
+    @abstractmethod
+    def generate(
+        self,
+        messages: list[Message],
+        config: GenerationConfig | None = None,
+        *,
+        on_token: Callable[[str], None] | None = None,
+        cancel: Callable[[], bool] | None = None,
+    ) -> GenerationResult:
+        """Generate a completion. `on_token` streams; `cancel` polls a stop flag."""
+
+    def load(self) -> None:  # noqa: B027 - optional hook, not every backend needs it
+        """Load weights. Backends that need no load (tests) do nothing."""
+
+    def unload(self) -> None:  # noqa: B027 - optional hook
+        """Release the GPU. Required before `verify.gpu_commands` (DESIGN §4.6)."""
+
+    def close(self) -> None:
+        self.unload()
+
+
+class FakeBackend(LLMBackend):
+    """A scripted backend for tests and for `--dry-run`.
+
+    Produces the queued responses in order; when the queue runs dry it repeats
+    the last one, which is what makes loop-detection tests deterministic.
+    """
+
+    def __init__(
+        self,
+        responses: list[str] | None = None,
+        *,
+        model_id: str = "fake/coder@none",
+        tokens_per_char: float = 0.25,
+    ) -> None:
+        self._responses = list(responses or ['```json\n{"tool": "done", "args": {}}\n```'])
+        self._index = 0
+        self._model_id = model_id
+        self._tpc = tokens_per_char
+        self.loaded = False
+        self.calls: list[dict[str, Any]] = []
+
+    @property
+    def model_id(self) -> str:
+        return self._model_id
+
+    def count_tokens(self, text: str) -> int:
+        return max(1, int(len(text) * self._tpc))
+
+    def load(self) -> None:
+        self.loaded = True
+
+    def unload(self) -> None:
+        self.loaded = False
+
+    def queue(self, *responses: str) -> None:
+        self._responses.extend(responses)
+
+    def generate(
+        self,
+        messages: list[Message],
+        config: GenerationConfig | None = None,
+        *,
+        on_token: Callable[[str], None] | None = None,
+        cancel: Callable[[], bool] | None = None,
+    ) -> GenerationResult:
+        cfg = config or GenerationConfig()
+        index = min(self._index, len(self._responses) - 1)
+        text = self._responses[index]
+        self._index += 1
+        self.calls.append(
+            {
+                "messages": [m.to_dict() for m in messages],
+                "max_new_tokens": cfg.max_new_tokens,
+                "temperature": cfg.temperature,
+            }
+        )
+
+        if on_token is not None:
+            # Stream in small chunks so a UI test sees multiple updates.
+            for i in range(0, len(text), 16):
+                if cancel is not None and cancel():
+                    emitted = text[:i]
+                    return GenerationResult(
+                        text=emitted,
+                        stop_reason=StopReason.CANCELLED,
+                        usage=Usage(self.count_tokens("".join(m.content for m in messages)), self.count_tokens(emitted)),
+                        model=self._model_id,
+                    )
+                on_token(text[i : i + 16])
+
+        stop = StopReason.LENGTH if len(text) > cfg.max_new_tokens else StopReason.STOP
+        return GenerationResult(
+            text=text,
+            stop_reason=stop,
+            usage=Usage(
+                prompt_tokens=self.count_tokens("".join(m.content for m in messages)),
+                completion_tokens=self.count_tokens(text),
+            ),
+            model=self._model_id,
+        )

coagent/llm/grammar.py

@@ -0,0 +1,306 @@
+"""Tolerant extraction of the tool-call envelope (DESIGN §4.3).
+
+A 7B model emits malformed JSON perhaps 3–8% of the time, and a malformed
+envelope costs a full turn. This module is the reliability core: it never
+raises on bad model output, always returns a diagnosis the loop can feed back.
+
+Escalation ladder, cheapest first:
+
+1. Exactly one ```` ```json ```` fence — the contract we ask for.
+2. Any fence containing `{...}` — models sometimes label it ```` ``` ````.
+3. A bare top-level `{...}` object (no fence at all).
+4. JSON5-ish repair: single quotes, trailing commas, unquoted keys.
+5. Failure with a **corrective message** naming what was wrong, so the loop's
+   single re-prompt is targeted rather than "try again".
+
+Only `{"tool": ..., "args": {...}}` is accepted. An invented tool name is a
+schema error the registry catches, not a parse error.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass
+from typing import Any
+
+__all__ = [
+    "Envelope",
+    "EnvelopeError",
+    "parse_envelope",
+    "extract_json_candidates",
+    "repair_json",
+    "strip_fences",
+    "looks_like_final_answer",
+]
+
+_FENCE_RE = re.compile(r"```(?:json|json5|JSON)?\s*\n?(.*?)```", re.DOTALL)
+_TRAILING_COMMA_RE = re.compile(r",(\s*[}\]])")
+_UNQUOTED_KEY_RE = re.compile(r"([{,]\s*)([A-Za-z_][A-Za-z0-9_\-]*)(\s*:)")
+_SINGLE_QUOTE_RE = re.compile(r"'([^'\\]*(?:\\.[^'\\]*)*)'")
+
+
+class EnvelopeError(ValueError):
+    """No usable envelope. `corrective` is fed back to the model."""
+
+    def __init__(self, message: str, *, corrective: str, raw: str = "") -> None:
+        super().__init__(message)
+        self.corrective = corrective
+        self.raw = raw
+
+
+@dataclass(slots=True)
+class Envelope:
+    tool: str
+    args: dict[str, Any]
+    raw: str = ""
+    repaired: bool = False
+
+    def to_dict(self) -> dict[str, Any]:
+        return {"tool": self.tool, "args": self.args}
+
+
+def strip_fences(text: str) -> list[str]:
+    """Every fenced block, plus the whole text as a last candidate."""
+    blocks = [m.group(1).strip() for m in _FENCE_RE.finditer(text)]
+    blocks.append(text.strip())
+    return blocks
+
+
+def _balanced_objects(text: str) -> list[str]:
+    """Extract balanced `{...}` runs, ignoring braces inside strings.
+
+    Needed because a model often writes prose around a bare object:
+    `Here is the call: {"tool": ...} and I will now...`
+    """
+    found: list[str] = []
+    depth = 0
+    start = -1
+    in_string = False
+    escape = False
+    for i, ch in enumerate(text):
+        if in_string:
+            if escape:
+                escape = False
+            elif ch == "\\":
+                escape = True
+            elif ch == '"':
+                in_string = False
+            continue
+        if ch == '"':
+            in_string = True
+        elif ch == "{":
+            if depth == 0:
+                start = i
+            depth += 1
+        elif ch == "}" and depth > 0:
+            depth -= 1
+            if depth == 0 and start >= 0:
+                found.append(text[start : i + 1])
+                start = -1
+    return found
+
+
+def extract_json_candidates(text: str) -> list[str]:
+    """Ordered candidates, most likely first, deduplicated."""
+    seen: set[str] = set()
+    out: list[str] = []
+    for block in strip_fences(text):
+        for candidate in (block, *_balanced_objects(block)):
+            candidate = candidate.strip()
+            if candidate and candidate not in seen:
+                seen.add(candidate)
+                out.append(candidate)
+    return out
+
+
+
+def _escape_controls(text: str) -> str:
+    r"""Escape literal control characters that appear inside JSON strings.
+
+    A small model writing code into a ``replace`` field emits a real newline
+    instead of ``\n``, which makes the block invalid JSON. That is the most
+    common envelope failure and it is always recoverable: inside a JSON string
+    a raw control character is never legal.
+    """
+    mapping = {"\n": "\\n", "\r": "\\r", "\t": "\\t"}
+    out = []
+    in_string = False
+    escaped = False
+    for ch in text:
+        if in_string:
+            if escaped:
+                escaped = False
+            elif ch == "\\":
+                escaped = True
+            elif ch == '"':
+                in_string = False
+            elif ch in mapping:
+                out.append(mapping[ch])
+                continue
+        elif ch == '"':
+            in_string = True
+        out.append(ch)
+    return "".join(out)
+
+
+def _truncated(text: str) -> bool:
+    """True when strings/braces are unbalanced, i.e. the block was cut off."""
+    depth = 0
+    in_string = False
+    escaped = False
+    for ch in text:
+        if in_string:
+            if escaped:
+                escaped = False
+            elif ch == "\\":
+                escaped = True
+            elif ch == '"':
+                in_string = False
+        elif ch == '"':
+            in_string = True
+        elif ch in "{[":
+            depth += 1
+        elif ch in "}]":
+            depth -= 1
+    return depth > 0 or in_string
+
+
+def repair_json(text: str) -> str:
+    """Best-effort JSON5-ish repair for the most common model slips."""
+    fixed = _escape_controls(text)
+    # Single-quoted strings → double-quoted, without touching apostrophes in prose.
+    fixed = _SINGLE_QUOTE_RE.sub(lambda m: '"' + m.group(1).replace('"', '\\"') + '"', fixed)
+    # Unquoted keys: {tool: "read_file"} → {"tool": "read_file"}
+    fixed = _UNQUOTED_KEY_RE.sub(lambda m: f'{m.group(1)}"{m.group(2)}"{m.group(3)}', fixed)
+    # Trailing commas before a closing brace/bracket.
+    fixed = _TRAILING_COMMA_RE.sub(r"\1", fixed)
+    return fixed
+
+
+def _try_load(candidate: str) -> tuple[dict[str, Any] | None, bool]:
+    """Return (object, repaired?) for one candidate."""
+    for attempt, payload in ((False, candidate), (True, repair_json(candidate))):
+        try:
+            value = json.loads(payload)
+        except (json.JSONDecodeError, ValueError):
+            continue
+        if isinstance(value, dict):
+            return value, attempt
+    return None, False
+
+
+def parse_envelope(text: str, *, allowed_tools: set[str] | None = None) -> Envelope:
+    """Extract `{"tool": ..., "args": {...}}` from model output.
+
+    Raises `EnvelopeError` with a `corrective` message — never a bare
+    `JSONDecodeError` — because the loop feeds that text straight back.
+    """
+    if not text or not text.strip():
+        raise EnvelopeError(
+            "the response was empty",
+            corrective=(
+                "Your last message was empty. Reply with exactly one fenced JSON block:\n"
+                '```json\n{"tool": "<tool name>", "args": {...}}\n```'
+            ),
+        )
+
+    saw_json = False
+    saw_unparseable = ""
+    for candidate in extract_json_candidates(text):
+        value, repaired = _try_load(candidate)
+        if value is None:
+            # Only a brace-bearing candidate is a real parse attempt; plain prose
+            # reaching here means the model never tried to call a tool.
+            if "{" in candidate and not saw_unparseable:
+                saw_unparseable = candidate
+            continue
+        saw_json = True
+        if "tool" not in value:
+            continue
+        tool = value.get("tool")
+        if not isinstance(tool, str):
+            raise EnvelopeError(
+                "`tool` is not a string",
+                corrective=f'`tool` must be a string tool name, got {type(tool).__name__}.',
+                raw=text,
+            )
+        if allowed_tools is not None and tool not in allowed_tools:
+            raise EnvelopeError(
+                f"unknown tool {tool!r}",
+                corrective=(
+                    f"There is no tool named {tool!r}. Choose one of: "
+                    f"{', '.join(sorted(allowed_tools))}."
+                ),
+                raw=text,
+            )
+        args = value.get("args", {})
+        if args is None:
+            args = {}
+        if not isinstance(args, dict):
+            raise EnvelopeError(
+                "`args` is not an object",
+                corrective=(
+                    f'`args` must be an object such as {{"path": "src/a.py"}}, got '
+                    f"{type(args).__name__}."
+                ),
+                raw=text,
+            )
+        return Envelope(tool=tool, args=args, raw=text, repaired=repaired)
+
+    if saw_json:
+        raise EnvelopeError(
+            "JSON found but no `tool` key",
+            corrective=(
+                "Your JSON block is missing the `tool` key. Reply with exactly:\n"
+                '```json\n{"tool": "<tool name>", "args": {...}}\n```'
+            ),
+            raw=text,
+        )
+    if saw_unparseable and _truncated(saw_unparseable):
+        raise EnvelopeError(
+            "the tool call was cut off",
+            corrective=(
+                "Your last JSON block was cut off before it finished, most likely "
+                "because the output limit was reached. Keep the edit small: send a "
+                "shorter `replace` block, or split the change into two edits. Reply "
+                'with one complete fenced JSON block:\n```json\n{"tool": "<tool name>", '
+                '"args": {...}}\n```'
+            ),
+            raw=text,
+        )
+    if saw_unparseable:
+        raise EnvelopeError(
+            "the JSON was malformed",
+            corrective=(
+                "Your JSON block could not be parsed. Common causes: an unescaped "
+                "quote or backslash inside a string, a missing comma, or a trailing "
+                'comma. If you are writing code into `replace`, escape newlines as '
+                '\\n and quotes as \\". Reply with one complete fenced JSON block.'
+            ),
+            raw=text,
+        )
+    raise EnvelopeError(
+        "no JSON object in the response",
+        corrective=(
+            "I could not find a JSON tool call in your last message. Do not describe "
+            "the action in prose; reply with exactly one fenced JSON block:\n"
+            '```json\n{"tool": "<tool name>", "args": {...}}\n```'
+        ),
+        raw=text,
+    )
+
+
+_FINAL_MARKERS = ("final_answer", "task_complete", "done")
+
+
+def looks_like_final_answer(text: str) -> bool:
+    """True when the model is answering in prose rather than calling a tool.
+
+    Used by the loop to decide between "prose: treat as the report" and "re-prompt
+    for an envelope" — never to guess at an action.
+    """
+    lowered = text.lower()
+    if any(marker in lowered for marker in _FINAL_MARKERS):
+        return True
+    return not any("{" in candidate for candidate in extract_json_candidates(text))

coagent/llm/registry.py

@@ -0,0 +1,147 @@
+"""Model registry and the T4 preflight (PRD D1, FR-MDL-1).
+
+v1 supports exactly two configurations: 7B at NF4/12k, and 3B at NF4/8k as the
+only fallback. Nothing here imports torch at module scope, so `coagent doctor`
+and the test suite run on a laptop.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+__all__ = [
+    "MODELS",
+    "ModelSpec",
+    "PreflightResult",
+    "get_model",
+    "preflight_gpu",
+    "T4_MIN_VRAM_GB",
+    "T4_COMPUTE_CAPABILITY",
+]
+
+T4_MIN_VRAM_GB = 14.0
+T4_COMPUTE_CAPABILITY = (7, 5)
+
+
+@dataclass(frozen=True, slots=True)
+class ModelSpec:
+    key: str
+    name: str
+    quant: str
+    max_context: int
+    vram_gb: float
+    notes: str = ""
+
+    @property
+    def trailer_id(self) -> str:
+        """The `Coagent-Model` value: `Qwen2.5-Coder-7B-Instruct@nf4`."""
+        short = self.name.rsplit("/", 1)[-1]
+        return f"{short}@{self.quant}"
+
+
+MODELS: dict[str, ModelSpec] = {
+    "qwen-coder-7b": ModelSpec(
+        key="qwen-coder-7b",
+        name="Qwen/Qwen2.5-Coder-7B-Instruct",
+        quant="nf4",
+        max_context=12288,
+        vram_gb=6.5,
+        notes="the only validated default on a free T4",
+    ),
+    "qwen-coder-3b": ModelSpec(
+        key="qwen-coder-3b",
+        name="Qwen/Qwen2.5-Coder-3B-Instruct",
+        quant="nf4",
+        max_context=8192,
+        vram_gb=3.0,
+        notes="OOM fallback only; never a bigger model (PRD D6)",
+    ),
+}
+
+
+def get_model(key: str) -> ModelSpec:
+    try:
+        return MODELS[key]
+    except KeyError:
+        raise KeyError(f"unknown model {key!r}; expected one of {sorted(MODELS)}") from None
+
+
+def get_model_by_name(name: str) -> ModelSpec | None:
+    for spec in MODELS.values():
+        if spec.name == name:
+            return spec
+    return None
+
+
+@dataclass(slots=True)
+class PreflightResult:
+    ok: bool
+    device: str = "none"
+    total_vram_gb: float = 0.0
+    compute_capability: tuple[int, int] = (0, 0)
+    reason: str = ""
+    override_allowed: bool = True
+
+    def describe(self) -> str:
+        if self.device == "none":
+            return "no CUDA device detected"
+        cap = f"{self.compute_capability[0]}.{self.compute_capability[1]}"
+        return f"{self.device} · {self.total_vram_gb:.0f} GB · CC {cap}"
+
+
+def preflight_gpu(*, device_index: int = 0) -> PreflightResult:
+    """Check the free-tier T4 envelope without loading any weights.
+
+    Returns a result rather than raising: the caller decides whether to refuse,
+    and the message is what the user reads (FR-BOOT-2).
+    """
+    try:
+        import torch
+    except ImportError:
+        return PreflightResult(
+            ok=False,
+            reason=(
+                "torch is not installed. coAgent needs the GPU stack: "
+                'pip install "coagent[llm]"'
+            ),
+        )
+
+    if not torch.cuda.is_available():
+        return PreflightResult(
+            ok=False,
+            reason=(
+                "no CUDA device is available. coAgent runs on a free Colab T4; "
+                "select Runtime → Change runtime type → T4 GPU. CPU-only is out of "
+                "scope (PRD D-1)."
+            ),
+        )
+
+    props = torch.cuda.get_device_properties(device_index)
+    name = props.name.replace("NVIDIA ", "").strip()
+    vram = props.total_memory / 1024**3
+    cap = (props.major, props.minor)
+
+    problems: list[str] = []
+    if "T4" not in name.upper():
+        problems.append(
+            f"the GPU is {name}, not a T4. coAgent is validated only on the free-tier T4 "
+            "(PRD D-1). Pass --allow-unsupported to run anyway, unsupported and unvalidated."
+        )
+    if cap < T4_COMPUTE_CAPABILITY:
+        problems.append(
+            f"compute capability {cap[0]}.{cap[1]} is below 7.5; fp16 and SDPA are the only "
+            "safe path and this device does not guarantee them."
+        )
+    if vram < T4_MIN_VRAM_GB:
+        problems.append(
+            f"only {vram:.1f} GB of VRAM is visible; the 7B at NF4 needs ~14 GB free to be safe. "
+            "The 3B fallback at 8k needs less — pass --allow-unsupported to try it."
+        )
+
+    return PreflightResult(
+        ok=not problems,
+        device=name,
+        total_vram_gb=vram,
+        compute_capability=cap,
+        reason=" ".join(problems),
+    )