continualcode 0.1.0__py3-none-any.whl

continualcode/__init__.py

@@ -0,0 +1,17 @@
+"""
+ContinualCode - Human-in-the-loop coding agent with online learning.
+
+A Claude Code-style TUI that learns from your approvals via prompt distillation.
+"""
+
+__version__ = "0.1.0"
+
+from continualcode.session import ContextDistillSession
+from continualcode.tools import TOOL_SPECS, READONLY_TOOLS, execute_tool
+
+__all__ = [
+    "ContextDistillSession",
+    "TOOL_SPECS",
+    "READONLY_TOOLS",
+    "execute_tool",
+]

continualcode/cli.py

@@ -0,0 +1,67 @@
+#!/usr/bin/env python3
+"""
+CLI entry point for ContinualCode.
+
+Usage:
+    continualcode                    # Run the TUI
+    continualcode --help             # Show help
+    continualcode --version          # Show version
+"""
+
+import argparse
+import sys
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        prog="continualcode",
+        description="Human-in-the-loop coding agent with online learning via prompt distillation",
+    )
+    parser.add_argument(
+        "--version",
+        action="version",
+        version="%(prog)s 0.1.0",
+    )
+    parser.add_argument(
+        "--model",
+        default=None,
+        help="Override MODEL environment variable",
+    )
+    parser.add_argument(
+        "--policy",
+        default=None,
+        help="Path to policy file (default: ./policy_memory.md)",
+    )
+    parser.add_argument(
+        "--checkpoint",
+        default=None,
+        help="Load checkpoint on startup",
+    )
+    parser.add_argument(
+        "--no-training",
+        action="store_true",
+        help="Disable training (inference only)",
+    )
+
+    args = parser.parse_args()
+
+    # Set environment variables from CLI args
+    import os
+    if args.model:
+        os.environ["MODEL"] = args.model
+    if args.policy:
+        os.environ["POLICY_PATH"] = args.policy
+    if args.checkpoint:
+        os.environ["LOAD_CHECKPOINT"] = args.checkpoint
+    if args.no_training:
+        os.environ["ENABLE_TRAINING"] = "0"
+
+    # Import and run the TUI
+    from continualcode.tui import TinkerCodeApp
+    app = TinkerCodeApp()
+    app.run()
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

continualcode/session.py

@@ -0,0 +1,357 @@
+#!/usr/bin/env python3
+"""
+Context (prompt) distillation session for a coding agent.
+
+We maintain two prefixes:
+  - Teacher prefix: includes a long policy prompt (e.g. repo/user rules).
+  - Student prefix: excludes that policy prompt.
+
+Generation uses the teacher prefix (better behavior now).
+
+Two training modes:
+  - off_policy: cross-entropy on the approved assistant message with the student prefix
+               (classic "prompt distillation").
+  - on_policy: sample a rollout from the student prefix, then compute teacher logprobs
+               (teacher prefix) on the student's exact tokens and update via importance
+               sampling / KL advantages (on-policy distillation).
+
+This is "prompt distillation" / "context distillation":
+  teacher(p + x) -> y
+  train student(x) -> y
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any
+
+import tinker
+import torch
+from tinker import types
+from tinker.types.tensor_data import TensorData
+from tinker_cookbook.renderers import get_renderer
+from tinker_cookbook.renderers.base import TrainOnWhat
+from tinker_cookbook.supervised.common import datum_from_model_input_weights
+from tinker_cookbook.tokenizer_utils import get_tokenizer
+
+from continualcode.tools import TOOL_SPECS
+
+
+class ContextDistillSession:
+    """Sampling + online prompt distillation training from approved tool calls."""
+
+    def __init__(
+        self,
+        model: str,
+        tinker_url: str | None = None,
+        *,
+        enable_training: bool = True,
+        lora_rank: int = 32,
+        learning_rate: float = 1e-5,
+        max_tokens: int = 4096,
+        temperature: float = 0.7,
+        policy_path: str = "./policy_memory.md",
+        policy_tag: str = "policy_memory",
+        distill_mode: str = "on_policy",  # "on_policy" or "off_policy"
+        # On-policy distillation config
+        train_max_tokens: int = 1024,
+        train_temperature: float = 1.0,
+        kl_coef: float = 1.0,
+    ):
+        self.model = model
+        self.learning_rate = learning_rate
+        self.max_tokens = max_tokens
+        self.temperature = temperature
+        self.train_steps = 0
+
+        self.policy_path = Path(policy_path)
+        self.policy_tag = policy_tag
+        self.distill_mode = distill_mode
+        self.train_max_tokens = train_max_tokens
+        self.train_temperature = train_temperature
+        self.kl_coef = kl_coef
+
+        # Setup tokenizer and renderer
+        self.tokenizer = get_tokenizer(model)
+        self.renderer = get_renderer("qwen3_instruct", tokenizer=self.tokenizer)
+        self.service = tinker.ServiceClient(base_url=tinker_url)
+
+        # Conversation state (no prefix)
+        self.messages: list[dict[str, Any]] = []
+
+        # Build prefixes
+        self._base_system_prompt = f"You are a helpful coding assistant. cwd: {os.getcwd()}"
+        self._teacher_prefix: list[dict[str, Any]] = []
+        self._student_prefix: list[dict[str, Any]] = []
+        self._rebuild_prefixes()
+
+        # Clients (initialized in init())
+        self.enable_training = enable_training
+        self.lora_rank = lora_rank
+        self.training_client: tinker.TrainingClient | None = None
+        self.sampling_client: tinker.SamplingClient | None = None
+
+    def _load_policy_text(self) -> str:
+        try:
+            return self.policy_path.read_text(encoding="utf-8")
+        except FileNotFoundError:
+            return ""
+        except Exception:
+            # Avoid making the entire session crash on a bad policy file.
+            return ""
+
+    def _rebuild_prefixes(self) -> None:
+        policy_text = self._load_policy_text().strip()
+        if policy_text:
+            teacher_prompt = (
+                self._base_system_prompt
+                + f"\n\n<{self.policy_tag}>\n{policy_text}\n</{self.policy_tag}>\n"
+            )
+        else:
+            teacher_prompt = self._base_system_prompt
+
+        self._teacher_prefix = self.renderer.create_conversation_prefix_with_tools(
+            TOOL_SPECS, teacher_prompt
+        )
+        self._student_prefix = self.renderer.create_conversation_prefix_with_tools(
+            TOOL_SPECS, self._base_system_prompt
+        )
+
+    def reload_policy(self) -> None:
+        """Reload policy from disk for subsequent generations."""
+        self._rebuild_prefixes()
+
+    async def init(self) -> None:
+        """Initialize training/sampling clients."""
+        if self.enable_training:
+            self.training_client = self.service.create_lora_training_client(
+                base_model=self.model, rank=self.lora_rank
+            )
+            self.sampling_client = self.training_client.save_weights_and_get_sampling_client("current")
+        else:
+            self.sampling_client = self.service.create_sampling_client(base_model=self.model)
+
+    async def sample(self) -> tuple[dict[str, Any], bool]:
+        """Sample a completion using the teacher prefix."""
+        if self.sampling_client is None:
+            raise RuntimeError("Session not initialized. Call init() first.")
+
+        # Policy can be updated externally; keep it fresh per turn.
+        self.reload_policy()
+
+        model_input = self.renderer.build_generation_prompt(self._teacher_prefix + self.messages)
+        response = await self.sampling_client.sample_async(
+            prompt=model_input,
+            num_samples=1,
+            sampling_params=types.SamplingParams(
+                stop=self.renderer.get_stop_sequences(),
+                max_tokens=self.max_tokens,
+                temperature=self.temperature,
+            ),
+        )
+        tokens = response.sequences[0].tokens
+        return self.renderer.parse_response(tokens)
+
+    def train_on_approval(
+        self,
+        *,
+        prompt_messages: list[dict[str, Any]],
+        assistant_message: dict[str, Any],
+        scale: float = 1.0,
+    ) -> dict[str, float]:
+        """Train a single update step from an approved action.
+
+        off_policy: cross-entropy imitation of the approved assistant message.
+        on_policy: on-policy distillation (teacher has policy context; student does not).
+        """
+        if self.training_client is None:
+            return {}
+
+        if self.distill_mode == "off_policy":
+            return self._train_off_policy(prompt_messages, assistant_message, scale=scale)
+        return self._train_on_policy(prompt_messages, scale=scale)
+
+    def _train_off_policy(
+        self,
+        prompt_messages: list[dict[str, Any]],
+        assistant_message: dict[str, Any],
+        *,
+        scale: float,
+    ) -> dict[str, float]:
+        """Classic prompt distillation: SFT student(x) -> teacher(p+x) output."""
+        if self.training_client is None:
+            return {}
+
+        train_messages = self._student_prefix + prompt_messages + [assistant_message]
+        model_input, weights = self.renderer.build_supervised_example(
+            train_messages, train_on_what=TrainOnWhat.LAST_ASSISTANT_MESSAGE
+        )
+        if scale != 1.0:
+            weights = [float(w) * float(scale) for w in weights]
+
+        datum = datum_from_model_input_weights(model_input, weights, max_length=None)
+        fwd_bwd = self.training_client.forward_backward([datum], loss_fn="cross_entropy").result()
+        self.training_client.optim_step(
+            types.AdamParams(learning_rate=self.learning_rate, beta1=0.9, beta2=0.95, eps=1e-8)
+        ).result()
+
+        self.sampling_client = self.training_client.save_weights_and_get_sampling_client("current")
+        self.train_steps += 1
+
+        metrics = fwd_bwd.metrics or {}
+        loss = metrics.get("loss:sum", metrics.get("loss", 0.0))
+        num_tokens = int(sum(1 for w in weights if w > 0))
+        return {"step": float(self.train_steps), "loss": float(loss), "tokens": float(num_tokens)}
+
+    def _train_on_policy(self, prompt_messages: list[dict[str, Any]], *, scale: float) -> dict[str, float]:
+        """On-policy distillation: sample student tokens; teacher (with policy) re-scores them."""
+        if self.training_client is None or self.sampling_client is None:
+            return {}
+
+        # Keep policy fresh for teacher scoring (policy may be edited between turns).
+        self.reload_policy()
+
+        student_prompt = self.renderer.build_generation_prompt(self._student_prefix + prompt_messages)
+        student_prompt_len = student_prompt.length
+
+        # 1) Sample a rollout from the student (no policy prompt).
+        response = self.sampling_client.sample(
+            prompt=student_prompt,
+            num_samples=1,
+            sampling_params=types.SamplingParams(
+                stop=self.renderer.get_stop_sequences(),
+                max_tokens=min(self.train_max_tokens, self.max_tokens),
+                temperature=self.train_temperature,
+            ),
+        ).result()
+        seq = response.sequences[0]
+        tokens = seq.tokens
+        student_logps = seq.logprobs or []
+        if not tokens or not student_logps:
+            return {}
+
+        # 2) Compute teacher logprobs on the student's exact tokens (teacher has policy prompt).
+        teacher_prompt = self.renderer.build_generation_prompt(self._teacher_prefix + prompt_messages)
+        teacher_prompt_len = teacher_prompt.length
+        full_teacher_input = teacher_prompt.append(types.EncodedTextChunk(tokens=tokens))
+        teacher_logprobs_full = self.sampling_client.compute_logprobs(full_teacher_input).result()
+
+        # compute_logprobs is shifted by 1 (logprob[i] predicts token[i+1]).
+        teacher_logps = teacher_logprobs_full[
+            teacher_prompt_len - 1 : teacher_prompt_len - 1 + len(tokens)
+        ]
+
+        min_len = min(len(tokens), len(student_logps), len(teacher_logps))
+        if min_len <= 0:
+            return {}
+        tokens = tokens[:min_len]
+        student_logps = student_logps[:min_len]
+        teacher_logps = teacher_logps[:min_len]
+
+        # Reverse KL per token = student - teacher. We optimize negative reverse KL:
+        # advantage[t] = scale * kl_coef * (teacher_lp - student_lp)
+        advantages: list[float] = []
+        total_kl = 0.0
+        for s_lp, t_lp in zip(student_logps, teacher_logps):
+            total_kl += (s_lp - t_lp)
+            advantages.append(float(scale) * float(self.kl_coef) * (t_lp - s_lp))
+
+        # 3) Build datum for importance_sampling.
+        full_input = student_prompt.append(types.EncodedTextChunk(tokens=tokens[:-1]))
+        target_tokens = [0] * (student_prompt_len - 1) + list(tokens)
+        padded_logprobs = [0.0] * (student_prompt_len - 1) + list(student_logps)
+        padded_advantages = [0.0] * (student_prompt_len - 1) + advantages
+
+        datum = types.Datum(
+            model_input=full_input,
+            loss_fn_inputs={
+                "target_tokens": TensorData.from_torch(torch.tensor(target_tokens)),
+                "logprobs": TensorData.from_torch(torch.tensor(padded_logprobs)),
+                "advantages": TensorData.from_torch(torch.tensor(padded_advantages)),
+            },
+        )
+
+        fwd_bwd = self.training_client.forward_backward([datum], loss_fn="importance_sampling").result()
+        self.training_client.optim_step(
+            types.AdamParams(learning_rate=self.learning_rate, beta1=0.9, beta2=0.95, eps=1e-8)
+        ).result()
+
+        # Post-update drift metrics (PPO-style sanity check)
+        approx_kl = 0.0
+        ratio_mean = 1.0
+        try:
+            fwd_after = self.training_client.forward([datum], loss_fn="importance_sampling").result()
+            out = fwd_after.loss_fn_outputs[0] if fwd_after.loss_fn_outputs else {}
+            td = out.get("logprobs") if isinstance(out, dict) else None
+            if td is not None:
+                new_lp = td.to_torch().flatten().to(torch.float32)
+                old_lp = datum.loss_fn_inputs["logprobs"].to_torch().flatten().to(torch.float32)
+                adv = datum.loss_fn_inputs["advantages"].to_torch().flatten().to(torch.float32)
+                mask = adv != 0
+                if mask.any():
+                    new_cat = new_lp[mask]
+                    old_cat = old_lp[mask]
+                    approx_kl = (old_cat - new_cat).mean().item()
+                    ratio_mean = torch.exp(new_cat - old_cat).mean().item()
+        except Exception:
+            pass
+
+        self.sampling_client = self.training_client.save_weights_and_get_sampling_client("current")
+        self.train_steps += 1
+
+        metrics = fwd_bwd.metrics or {}
+        loss = metrics.get("loss:sum", metrics.get("loss", 0.0))
+        kl_st = total_kl / max(1, min_len)
+        adv_mean = sum(advantages) / max(1, len(advantages))
+        return {
+            "step": float(self.train_steps),
+            "loss": float(loss),
+            "tokens": float(min_len),
+            "kl_student_teacher": float(kl_st),
+            "adv_mean": float(adv_mean),
+            "approx_kl": float(approx_kl),
+            "ratio": float(ratio_mean),
+        }
+
+    @property
+    def teacher_prefix_messages(self) -> list[dict[str, Any]]:
+        return list(self._teacher_prefix)
+
+    @property
+    def student_prefix_messages(self) -> list[dict[str, Any]]:
+        return list(self._student_prefix)
+
+    def add_user(self, content: str) -> None:
+        self.messages.append({"role": "user", "content": content})
+
+    def add_assistant(self, msg: dict[str, Any]) -> None:
+        self.messages.append(msg)
+
+    def add_tool_result(self, tool_call_id: str, result: str) -> None:
+        self.messages.append({"role": "tool", "tool_call_id": tool_call_id, "content": result})
+
+    def clear(self) -> None:
+        self.messages.clear()
+
+    def save_checkpoint(self, path: str) -> bool:
+        """Save current LoRA weights to disk. Returns True on success."""
+        if self.training_client is None:
+            return False
+        try:
+            self.training_client.save_state(path).result()
+            return True
+        except Exception:
+            return False
+
+    def load_checkpoint(self, path: str) -> bool:
+        """Load LoRA weights from disk. Returns True on success."""
+        if self.training_client is None:
+            return False
+        try:
+            self.training_client.load_state(path).result()
+            # Update sampling client to use loaded weights
+            self.sampling_client = self.training_client.save_weights_and_get_sampling_client("current")
+            return True
+        except Exception:
+            return False

continualcode/tools.py

@@ -0,0 +1,230 @@
+#!/usr/bin/env python3
+"""
+Core tool implementations and schemas shared by tinkercode harnesses.
+
+Tool names intentionally match the "Claude Code"-style set:
+read, write, edit, glob, grep, bash
+"""
+
+from __future__ import annotations
+
+import glob as globlib
+import os
+import re
+import subprocess
+from pathlib import Path
+from typing import Any
+
+try:
+    # Only required by the Tinker renderer; keep optional for simple imports.
+    from tinker_cookbook.renderers import ToolSpec
+except Exception:  # pragma: no cover
+    ToolSpec = dict  # type: ignore[misc,assignment]
+
+
+def tool_read(args: dict[str, Any]) -> str:
+    path = args["path"]
+    if not os.path.isfile(path):
+        return f"error: file not found: {path}"
+
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            lines = f.readlines()
+    except UnicodeDecodeError:
+        with open(path, "r", encoding="utf-8", errors="replace") as f:
+            lines = f.readlines()
+
+    offset = max(0, int(args.get("offset", 0) or 0))
+    limit = args.get("limit")
+    limit_int = len(lines) if limit is None else max(0, int(limit))
+    selected = lines[offset : offset + limit_int]
+    return "".join(f"{offset + idx + 1:4}| {line}" for idx, line in enumerate(selected))
+
+
+def tool_write(args: dict[str, Any]) -> str:
+    path = args["path"]
+    content = args["content"]
+    Path(path).expanduser().parent.mkdir(parents=True, exist_ok=True)
+    with open(path, "w", encoding="utf-8") as f:
+        f.write(content)
+    return "ok"
+
+
+def tool_edit(args: dict[str, Any]) -> str:
+    path = args["path"]
+    if not os.path.isfile(path):
+        return f"error: file not found: {path}"
+
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            text = f.read()
+    except UnicodeDecodeError:
+        with open(path, "r", encoding="utf-8", errors="replace") as f:
+            text = f.read()
+
+    old = args["old"]
+    new = args["new"]
+    replace_all = bool(args.get("all", False))
+
+    if old not in text:
+        return "error: old_string not found"
+
+    count = text.count(old)
+    if not replace_all and count > 1:
+        return f"error: old_string appears {count} times, must be unique (use all=true)"
+
+    replacement = text.replace(old, new) if replace_all else text.replace(old, new, 1)
+    with open(path, "w", encoding="utf-8") as f:
+        f.write(replacement)
+    return "ok"
+
+
+def tool_glob(args: dict[str, Any]) -> str:
+    base = args.get("path", ".") or "."
+    pattern = (str(base) + "/" + args["pat"]).replace("//", "/")
+    files = globlib.glob(pattern, recursive=True)
+    files = sorted(
+        files,
+        key=lambda p: os.path.getmtime(p) if os.path.isfile(p) else 0,
+        reverse=True,
+    )
+    return "\n".join(files[:50]) or "none"
+
+
+def tool_grep(args: dict[str, Any]) -> str:
+    try:
+        pattern = re.compile(args["pat"])
+    except re.error:
+        return "error: invalid regex pattern"
+
+    base = args.get("path", ".") or "."
+    hits: list[str] = []
+
+    for filepath in globlib.glob(str(base) + "/**", recursive=True):
+        if not os.path.isfile(filepath):
+            continue
+        try:
+            with open(filepath, "r", encoding="utf-8", errors="ignore") as f:
+                for line_num, line in enumerate(f, 1):
+                    if pattern.search(line):
+                        hits.append(f"{filepath}:{line_num}:{line.rstrip()}")
+                        if len(hits) >= 50:
+                            return "\n".join(hits)
+        except Exception:
+            continue
+
+    return "\n".join(hits) or "none"
+
+
+def tool_bash(args: dict[str, Any]) -> str:
+    cmd = args["cmd"]
+    try:
+        result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=30)
+        output = (result.stdout or "") + (result.stderr or "")
+        return output.strip() or "(empty output)"
+    except subprocess.TimeoutExpired:
+        return "error: command timed out after 30s"
+    except Exception as e:
+        return f"error: {e}"
+
+
+TOOL_FUNCTIONS: dict[str, Any] = {
+    "read": tool_read,
+    "write": tool_write,
+    "edit": tool_edit,
+    "glob": tool_glob,
+    "grep": tool_grep,
+    "bash": tool_bash,
+}
+
+
+def execute_tool(name: str, args: dict[str, Any]) -> str:
+    fn = TOOL_FUNCTIONS.get(name)
+    if fn is None:
+        return f"error: unknown tool '{name}'"
+    try:
+        return fn(args)
+    except Exception as e:
+        return f"error: {e}"
+
+
+TOOL_SPECS: list[ToolSpec] = [
+    {
+        "name": "read",
+        "description": "Read a file and return its contents with line numbers",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "path": {"type": "string", "description": "Path to the file to read"},
+                "offset": {"type": "integer", "description": "Line number to start from (0-indexed)"},
+                "limit": {"type": "integer", "description": "Maximum number of lines to read"},
+            },
+            "required": ["path"],
+        },
+    },
+    {
+        "name": "write",
+        "description": "Write content to a file, creating it if it doesn't exist",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "path": {"type": "string", "description": "Path to the file to write"},
+                "content": {"type": "string", "description": "Content to write to the file"},
+            },
+            "required": ["path", "content"],
+        },
+    },
+    {
+        "name": "edit",
+        "description": "Edit a file by replacing old with new. old must appear once unless all=true.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "path": {"type": "string", "description": "Path to the file to edit"},
+                "old": {"type": "string", "description": "The exact string to find and replace"},
+                "new": {"type": "string", "description": "The string to replace it with"},
+                "all": {"type": "boolean", "description": "Replace all occurrences (default false)"},
+            },
+            "required": ["path", "old", "new"],
+        },
+    },
+    {
+        "name": "glob",
+        "description": "Find files matching a glob pattern",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "pat": {"type": "string", "description": "Glob pattern (e.g., '**/*.py')"},
+                "path": {"type": "string", "description": "Base path to search from"},
+            },
+            "required": ["pat"],
+        },
+    },
+    {
+        "name": "grep",
+        "description": "Search for a regex pattern in files",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "pat": {"type": "string", "description": "Regex pattern to search for"},
+                "path": {"type": "string", "description": "Base path to search from"},
+            },
+            "required": ["pat"],
+        },
+    },
+    {
+        "name": "bash",
+        "description": "Run a shell command",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "cmd": {"type": "string", "description": "The shell command to run"},
+            },
+            "required": ["cmd"],
+        },
+    },
+]
+
+
+READONLY_TOOLS = {"read", "glob", "grep"}
+