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

coagent/llm/transformers_backend.py

@@ -0,0 +1,326 @@
+"""The local `transformers` backend (DESIGN §4.2, §4.5).
+
+Pinned to what a Turing T4 can actually do:
+
+* **NF4 4-bit weights, fp16 compute.** T4 (SM 7.5) has no bf16.
+* **`attn_implementation="sdpa"`.** FlashAttention-2 needs Ampere+.
+* **No `torch.compile`.** Compile time exceeds most sessions' savings.
+* **`PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` set before torch import.**
+
+Everything torch-related is imported inside methods so `import coagent` stays
+light and the test suite runs without a GPU.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import gc
+import os
+import threading
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+from .base import (
+    GenerationConfig,
+    GenerationResult,
+    LLMBackend,
+    Message,
+    StopReason,
+    Usage,
+)
+from .registry import MODELS, ModelSpec, get_model
+
+__all__ = ["TransformersBackend", "OutOfMemoryLadder"]
+
+
+def _torch():
+    """Resolve torch at call time.
+
+    Module-level accessors keep `import coagent` torch-free while still letting a
+    test substitute a stub — patching a local `import` is impossible.
+    """
+    import torch
+
+    return torch
+
+
+def _transformers():
+    """Resolve transformers at call time (same reason as `_torch`)."""
+    import transformers
+
+    return transformers
+
+
+@dataclass(slots=True)
+class OutOfMemoryLadder:
+    """The OOM response order from DESIGN §4.5.
+
+    The session must never be lost to an OOM, and the model must never grow.
+    """
+
+    step: int = 0
+    context_reduction: float = 0.25
+    max_new_tokens_reduction: float = 0.5
+
+    def next_action(self) -> str:
+        """`shrink_context`, `halve_output`, `fallback_model`, or `stop`."""
+        return ("shrink_context", "halve_output", "fallback_model", "stop")[min(self.step, 3)]
+
+    def advance(self) -> str:
+        action = self.next_action()
+        self.step += 1
+        return action
+
+    @property
+    def exhausted(self) -> bool:
+        """True once every recovery action has been tried (3 actions, then stop)."""
+        return self.step >= 3
+
+
+class TransformersBackend(LLMBackend):
+    """A 4-bit causal LM on the local GPU."""
+
+    name = "transformers"
+
+    def __init__(
+        self,
+        model_key: str = "qwen-coder-7b",
+        *,
+        spec: ModelSpec | None = None,
+        device_map: str | dict[str, Any] = "auto",
+        trust_remote_code: bool = False,
+        max_context: int | None = None,
+    ) -> None:
+        # Set the allocator flag before torch ever initialises CUDA.
+        os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+        self.spec = spec or get_model(model_key)
+        self._device_map = device_map
+        self._trust = trust_remote_code
+        self._max_context_override = max_context
+        self._model: Any = None
+        self._tokenizer: Any = None
+        self._tokenizer_errors: set[str] = set()
+        self._lock = threading.Lock()
+
+    # ------------------------------------------------------------- lifecycle
+
+    @property
+    def model_id(self) -> str:
+        return self.spec.trailer_id
+
+    @property
+    def max_context(self) -> int:
+        return self._max_context_override or self.spec.max_context
+
+    @property
+    def loaded(self) -> bool:
+        return self._model is not None
+
+    def load(self) -> None:
+        if self._model is not None:
+            return
+        torch = _torch()
+        transformers = _transformers()
+        auto_model = transformers.AutoModelForCausalLM
+        auto_tokenizer = transformers.AutoTokenizer
+
+        quant = transformers.BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            # T4 has no bf16. fp16 compute is the only correct choice.
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+        )
+
+        self._tokenizer = auto_tokenizer.from_pretrained(
+            self.spec.name, trust_remote_code=self._trust
+        )
+        self._model = auto_model.from_pretrained(
+            self.spec.name,
+            quantization_config=quant,
+            device_map=self._device_map,
+            attn_implementation="sdpa",  # flash-attn 2 needs Ampere+
+            torch_dtype=torch.float16,
+            trust_remote_code=self._trust,
+        )
+        self._model.eval()
+
+    def unload(self) -> None:
+        """Release VRAM. Required before `verify.gpu_commands` (DESIGN §4.6)."""
+        if self._model is None:
+            return
+        self._model = None
+        self._tokenizer = None
+        gc.collect()
+        try:
+            torch = _torch()
+
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                torch.cuda.ipc_collect()
+        except Exception:  # noqa: BLE001 - unloading must never raise
+            pass
+
+    # ------------------------------------------------------------- tokenizer
+
+    def count_tokens(self, text: str) -> int:
+        """Tokenizer truth for the budget (FR-CTX-6).
+
+        Falls back to a conservative character estimate only if the tokenizer
+        cannot be loaded at all, and records that it did so.
+        """
+        if not text:
+            return 0
+        if self._tokenizer is None:
+            try:
+                self._tokenizer = _transformers().AutoTokenizer.from_pretrained(
+                    self.spec.name, trust_remote_code=self._trust
+                )
+            except Exception:  # noqa: BLE001 - no network / no weights
+                if "estimate" not in self._tokenizer_errors:
+                    self._tokenizer_errors.add("estimate")
+                # ~4 chars/token is the usual ratio; over-estimating is the safe bias.
+                return max(1, len(text) // 3)
+        return len(self._tokenizer.encode(text, add_special_tokens=False))
+
+    def using_estimated_tokens(self) -> bool:
+        return "estimate" in self._tokenizer_errors
+
+    # ------------------------------------------------------------ generation
+
+    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()
+        self.load()
+        torch = _torch()
+        transformers = _transformers()
+
+        assert self._tokenizer is not None
+        prompt = self._tokenizer.apply_chat_template(
+            [m.to_dict() for m in messages], tokenize=False, add_generation_prompt=True
+        )
+        inputs = self._tokenizer(prompt, return_tensors="pt").to(self._model.device)
+
+        if cfg.seed:
+            torch.manual_seed(cfg.seed)
+
+        streamer = transformers.TextIteratorStreamer(
+            self._tokenizer, skip_prompt=True, skip_special_tokens=True
+        )
+
+        stop_criteria = _make_stop_criteria(cancel) if cancel is not None else None
+        stopping = transformers.StoppingCriteriaList([stop_criteria]) if stop_criteria else None
+
+        gen_kwargs: dict[str, Any] = {
+            "input_ids": inputs["input_ids"],
+            "attention_mask": inputs.get("attention_mask"),
+            "max_new_tokens": cfg.max_new_tokens,
+            "do_sample": cfg.temperature > 0,
+            "temperature": max(cfg.temperature, 1e-5),
+            "top_p": cfg.top_p,
+            "repetition_penalty": 1.05,
+            "pad_token_id": self._tokenizer.eos_token_id,
+            "streamer": streamer,
+        }
+        if stopping is not None:
+            gen_kwargs["stopping_criteria"] = stopping
+
+        error_holder: list[BaseException] = []
+
+        def _run() -> None:
+            try:
+                with torch.inference_mode():
+                    self._model.generate(**gen_kwargs)
+            except BaseException as exc:  # noqa: BLE001 - surfaced below
+                error_holder.append(exc)
+                with contextlib.suppress(Exception):
+                    streamer.end()
+
+        worker = threading.Thread(target=_run, daemon=True)
+        worker.start()
+
+        pieces: list[str] = []
+        try:
+            for chunk in streamer:
+                if chunk:
+                    pieces.append(chunk)
+                    if on_token is not None:
+                        on_token(chunk)
+        finally:
+            worker.join(timeout=1.0)
+
+        if error_holder:
+            exc = error_holder[0]
+            if _is_cuda_oom(exc):
+                self.unload()
+                raise _oom_from(exc) from exc
+            raise exc
+
+        text = "".join(pieces)
+        cancelled = cancel is not None and cancel()
+        stop = StopReason.CANCELLED if cancelled else StopReason.STOP
+        if not cancelled and len(pieces) and cfg.max_new_tokens and self.count_tokens(text) >= cfg.max_new_tokens:
+            stop = StopReason.LENGTH
+
+        return GenerationResult(
+            text=text,
+            stop_reason=stop,
+            usage=Usage(
+                prompt_tokens=int(inputs["input_ids"].shape[-1]),
+                completion_tokens=self.count_tokens(text),
+            ),
+            model=self.model_id,
+        )
+
+    # ---------------------------------------------------------------- ladder
+
+    def apply_oom_action(self, action: str, *, context: int, max_new_tokens: int) -> tuple[int, int, str]:
+        """Return `(context, max_new_tokens, model_key)` for the next attempt."""
+        ladder = OutOfMemoryLadder()
+        if action == "shrink_context":
+            return int(context * (1 - ladder.context_reduction)), max_new_tokens, self.spec.key
+        if action == "halve_output":
+            return context, max(64, max_new_tokens // 2), self.spec.key
+        if action == "fallback_model":
+            fallback = MODELS["qwen-coder-3b"]
+            return fallback.max_context, max_new_tokens, fallback.key
+        return context, max_new_tokens, self.spec.key
+
+
+def _make_stop_criteria(cancel: Callable[[], bool]):
+    """Build a real HF `StoppingCriteria` that polls the cancel flag (FR-UX-8).
+
+    Imported lazily: this module must stay importable without transformers.
+    """
+    class _StopOnCancel(_transformers().StoppingCriteria):
+        def __call__(self, input_ids, scores, **kwargs) -> bool:  # noqa: ANN001, ARG002
+            return bool(cancel())
+
+    return _StopOnCancel()
+
+
+def _is_cuda_oom(exc: BaseException) -> bool:
+    name = type(exc).__name__.lower()
+    text = str(exc).lower()
+    return (
+        "outofmemory" in name
+        or "out of memory" in text
+        # A wrapped torch OOM often surfaces only as the class name in the message.
+        or "outofmemoryerror" in text
+        or "cuda error: out of memory" in text
+    )
+
+
+def _oom_from(exc: BaseException) -> RuntimeError:
+    return RuntimeError(
+        "CUDA is out of memory. coAgent unloaded the model and will retry with a "
+        "smaller context; if that fails it halves the output length, then falls "
+        "back to the 3B model. Your session and its checkpoints are intact."
+    )

coagent/safety/__init__.py

@@ -0,0 +1,24 @@
+"""Safety: redaction, denylist, and the approval matrix (DESIGN §9)."""
+
+from __future__ import annotations
+
+from .denylist import DEFAULT_DENIED, Denial, Denylist, deny_reasons
+from .policy import Decision, Policy, RiskClass, SafetyClass, Verdict
+from .redact import MASK, Redactor, contains_secret, redact, redact_obj
+
+__all__ = [
+    "DEFAULT_DENIED",
+    "MASK",
+    "Decision",
+    "Denial",
+    "Denylist",
+    "Policy",
+    "Redactor",
+    "RiskClass",
+    "SafetyClass",
+    "Verdict",
+    "contains_secret",
+    "deny_reasons",
+    "redact",
+    "redact_obj",
+]

coagent/safety/denylist.py

@@ -0,0 +1,146 @@
+"""Denylist of destructive commands (PRD FR-TOL-5, FR-GIT-9).
+
+This is a seatbelt, not a sandbox — the reality of a root-ish Colab shell.
+The real control is the approval matrix in `policy.py`. What this module does
+is evaluate a command *before* anything runs, deny rules before allow rules,
+and return an actionable reason.
+
+Matching strategy: a permissive tokenizer (shlex) so `git push --force` and
+`git push --force=true` and `git push -f` are all caught, but a commit message
+quoting `reset --hard` is not.
+"""
+
+from __future__ import annotations
+
+import shlex
+from dataclasses import dataclass
+
+__all__ = ["Denial", "Denylist", "DEFAULT_DENIED", "deny_reasons"]
+
+
+@dataclass(frozen=True, slots=True)
+class Denial:
+    rule: str
+    reason: str
+
+
+# Pattern: token sequence of a dangerous command. Tokens are matched in order
+# and may skip over "=" joins: `git push --force` and `git push --force=true`.
+_RULES: tuple[tuple[str, tuple[str, ...], str], ...] = (
+    ("git-force-push", ("git", "push", "--force*"), "force-push rewrites remote history (FR-GIT-9)."),
+    ("git-force-push-short", ("git", "push", "-f"), "force-push rewrites remote history (FR-GIT-9)."),
+    ("git-push-main", ("git", "push", "*", "main*"), "pushing main is never allowed (FR-GIT-8)."),
+    ("git-reset-hard", ("git", "reset", "--hard*"), "hard reset discards work (FR-GIT-9)."),
+    (
+        "git-branch-delete",
+        ("git", "branch", "-D*"),
+        "deleting branches is destructive (FR-GIT-9).",
+    ),
+    ("git-clean-fd", ("git", "clean", "-f*"), "git clean -f deletes untracked files."),
+    ("rm-root", ("rm", "-rf", "/"), "rm -rf / is never valid here (FR-TOL-5)."),
+    ("rm-home", ("rm", "-rf", "~"), "rm -rf ~ is never valid here."),
+    ("chmod-root", ("chmod", "-R", "777", "/"), "chmod -R 777 / destroys a system."),
+    ("chmod-home", ("chmod", "-R", "777", "~"), "chmod -R 777 ~ destroys a system."),
+    (
+        "pipe-to-shell",
+        ("*", "|", "bash"),
+        "piping to a shell executes untrusted input; requires approval at a higher level.",
+    ),
+    ("curl-pipe-sh", ("curl*", "|", "sh"), "curl | sh executes remote code (FR-TOL-5)."),
+    ("wget-pipe-sh", ("wget*", "|", "sh"), "wget | sh executes remote code (FR-TOL-5)."),
+    ("dd-disk", ("dd", "*", "of=/dev/*"), "dd to a device overwrites disks."),
+    ("fork-bomb", ("*", ":(){", "|", "*", "&"), "fork bombs violate Colab fair use."),
+    ("shutdown", ("shutdown", "*"), "shutting down the VM kills the session (FR-REL-2)."),
+    ("reboot", ("reboot", "*"), "rebooting the VM kills the session (FR-REL-2)."),
+    ("rm-var", ("rm", "-rf", "/var"), "deleting /var destroys the runtime."),
+    ("mkfs", ("mkfs*", "*"), "mkfs formats disks."),
+)
+
+
+def _tokens(cmd: str) -> list[str]:
+    try:
+        return shlex.split(cmd)
+    except ValueError:
+        # Unbalanced quotes — split on whitespace instead so a denial still
+        # applies; the command would fail to run anyway.
+        return cmd.split()
+
+
+def _token_matches(actual: str, rule: str) -> bool:
+    if rule.endswith("*") and len(rule) > 1:
+        prefix = rule[:-1]
+        if actual.startswith(prefix):
+            return True
+        # Flags are case-sensitive in git, but `-d` (merged-only delete) is the
+        # common lowercase slip; treat any `-d`/`-D` variant as destructive.
+        if prefix == "-D":
+            return actual.startswith("-d") or actual.startswith("-D")
+        return False
+    if rule == "*":
+        return True
+    return actual == rule
+
+
+DEFAULT_DENIED: tuple[str, ...] = tuple(rule[0] for rule in _RULES)
+
+
+def deny_reasons(cmd: str, rules: tuple[str, ...] = DEFAULT_DENIED) -> tuple[Denial, ...]:
+    """Return every denial that applies to `cmd`.
+
+    Deny rules are evaluated before allow rules (PRD FR-GIT-9): a force-push to
+    a coagent/* ref is still denied.
+    """
+    tokens = _tokens(cmd)
+    if not tokens:
+        return ()
+    reasons: list[Denial] = []
+    for rule_id, pattern, reason in _RULES:
+        if rule_id not in rules:
+            continue
+        if _match_pattern(tokens, pattern):
+            reasons.append(Denial(rule=rule_id, reason=reason))
+    return tuple(reasons)
+
+
+def _match_pattern(tokens: list[str], pattern: tuple[str, ...]) -> bool:
+    """Sequence match with wildcards.
+
+    * a trailing `*` on a rule token is a prefix match, so `--force=true`
+      matches the rule `--force*`;
+    * the rule token `"*"` matches exactly one token;
+    * the rule token `"|"` means the *rest* of the pattern must match
+      somewhere later in the token list (a shell pipe).
+    """
+    if not tokens or not pattern:
+        return False
+
+    def match_from(tok_i: int, pat_i: int) -> bool:
+        while pat_i < len(pattern):
+            rule = pattern[pat_i]
+            if rule == "|":
+                rest = pattern[pat_i + 1 :]
+                # A pipe: try every later token as the start of the rest.
+                for start in range(tok_i, len(tokens)):
+                    if match_from(start, pat_i + 1) if rest else True:
+                        return True
+                return False
+            if tok_i >= len(tokens):
+                return False
+            if not _token_matches(tokens[tok_i], rule):
+                return False
+            tok_i += 1
+            pat_i += 1
+        return True
+
+    return match_from(0, 0)
+
+
+class Denylist:
+    """Holds the configurable deny rules for one session."""
+
+    def __init__(self, rules: tuple[str, ...] | list[str] | None = None) -> None:
+        self.rules: tuple[str, ...] = tuple(rules) if rules is not None else DEFAULT_DENIED
+
+    def check(self, cmd: str) -> Denial | None:
+        reasons = deny_reasons(cmd, self.rules)
+        return reasons[0] if reasons else None