argot-engine 0.1.4__py3-none-any.whl

argot/__main__.py

@@ -0,0 +1,3 @@
+from argot.extract import main
+
+main()

argot/check.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+import joblib  # type: ignore[import-untyped]
+import pygit2
+import torch
+
+from argot.git_walk import walk_commits
+from argot.jepa.encoder import TokenEncoder
+from argot.jepa.model import JEPAArgot
+from argot.jepa.predictor import ArgotPredictor
+from argot.tokenize import language_for_path, tokenize_lines
+
+
+def _resolve_shas(repo: pygit2.Repository, ref: str) -> set[str]:
+    """Parse a git range (A..B or bare ref) into a set of commit SHAs."""
+    if ".." in ref:
+        start_ref, end_ref = ref.split("..", 1)
+    else:
+        start_ref = ref + "^"
+        end_ref = ref
+
+    end_oid = repo.revparse_single(end_ref).id
+    try:
+        start_oid = repo.revparse_single(start_ref).id
+    except (pygit2.GitError, KeyError):
+        start_oid = None
+
+    shas: set[str] = set()
+    for commit in repo.walk(end_oid, pygit2.enums.SortMode.TOPOLOGICAL):
+        if start_oid is not None and commit.id == start_oid:
+            break
+        shas.add(str(commit.id))
+    return shas
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Check code surprise with argot JEPA model")
+    parser.add_argument("repo_path")
+    parser.add_argument("ref")
+    parser.add_argument("--model", default=".argot/model.pkl")
+    parser.add_argument("--threshold", type=float, default=0.5)
+    args = parser.parse_args()
+
+    model_path = Path(args.model)
+    if not model_path.exists():
+        print(f"error: model not found at {model_path}", file=sys.stderr)
+        sys.exit(2)
+
+    bundle = joblib.load(model_path)
+    vectorizer = bundle["vectorizer"]
+    embed_dim: int = bundle["embed_dim"]
+    input_dim: int = bundle["input_dim"]
+
+    encoder = TokenEncoder(input_dim, embed_dim)
+    encoder.load_state_dict(bundle["encoder_state"])
+    predictor = ArgotPredictor(embed_dim=embed_dim)
+    predictor.load_state_dict(bundle["predictor_state"])
+    model = JEPAArgot(encoder, predictor)
+    model.eval()
+
+    repo = pygit2.Repository(args.repo_path)
+    shas = _resolve_shas(repo, args.ref)
+    if not shas:
+        print("No commits found in range", file=sys.stderr)
+        sys.exit(0)
+
+    results: list[tuple[float, str, int, str]] = []
+
+    context_lines = 50
+    with torch.no_grad():
+        for commit, file_path, post_blob, hunks in walk_commits(args.repo_path, shas):
+            lang = language_for_path(file_path)
+            if lang is None:
+                continue
+            try:
+                source_lines = post_blob.decode("utf-8", errors="replace").splitlines()
+            except Exception:
+                continue
+
+            for hunk in hunks:
+                hunk_start = hunk.new_start - 1
+                hunk_end = hunk_start + hunk.new_lines
+                if hunk_start < 0 or hunk_end > len(source_lines):
+                    continue
+
+                before_start = max(0, hunk_start - context_lines)
+                ctx_tokens = tokenize_lines(source_lines, lang, before_start, hunk_start)
+                hunk_tokens = tokenize_lines(source_lines, lang, hunk_start, hunk_end)
+
+                ctx_text = " ".join(t.text for t in ctx_tokens)
+                hunk_text = " ".join(t.text for t in hunk_tokens)
+
+                ctx_vec = torch.tensor(
+                    vectorizer.transform([ctx_text]).toarray(), dtype=torch.float32
+                )
+                hunk_vec = torch.tensor(
+                    vectorizer.transform([hunk_text]).toarray(), dtype=torch.float32
+                )
+
+                score = model.surprise(ctx_vec, hunk_vec).item()
+                results.append((score, file_path, hunk.new_start, str(commit.id)[:8]))
+
+    if not results:
+        print("No hunks found in range")
+        sys.exit(0)
+
+    results.sort(key=lambda r: r[0], reverse=True)
+
+    print(f"{'SURPRISE':>9}  {'FILE':<48}  {'LINE':>5}  COMMIT")
+    for score, fp, line, sha in results:
+        print(f"{score:>9.4f}  {fp:<48}  {line:>5}  {sha}")
+
+    if any(s > args.threshold for s, *_ in results):
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

argot/dataset.py

@@ -0,0 +1,29 @@
+# Wire format for the dataset JSONL. Every hunk emitted conforms to this schema.
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+Language = Literal["typescript", "javascript", "python"]
+
+
+@dataclass(frozen=True, slots=True)
+class Token:
+    text: str
+    node_type: str  # tree-sitter node kind, e.g. "function_declaration"
+    start_line: int
+    end_line: int
+
+
+@dataclass(frozen=True, slots=True)
+class HunkRecord:
+    commit_sha: str
+    file_path: str
+    language: Language
+    hunk_start_line: int
+    hunk_end_line: int
+    context_before: list[Token]  # up to 50 lines before, tokenized
+    hunk_tokens: list[Token]
+    context_after: list[Token]  # up to 50 lines after
+    parent_sha: str | None
+    author_date_iso: str

argot/explain.py

@@ -0,0 +1,167 @@
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Any
+
+import joblib  # type: ignore[import-untyped]
+import numpy as np
+import pygit2
+import torch
+
+from argot.check import _resolve_shas
+from argot.git_walk import walk_commits
+from argot.jepa.encoder import TokenEncoder
+from argot.jepa.model import JEPAArgot
+from argot.jepa.predictor import ArgotPredictor
+from argot.tokenize import language_for_path, tokenize_lines
+
+
+def percentile_rank(value: float, distribution: list[float]) -> float:
+    arr = np.array(distribution)
+    return float(np.mean(arr < value) * 100)
+
+
+def select_style_examples(records: list[dict[str, Any]], *, n: int = 5) -> list[dict[str, Any]]:
+    """Pick n lowest-surprise records, one per file where possible."""
+    sorted_records = sorted(records, key=lambda r: r["_score"])
+    seen_files: set[str] = set()
+    diverse: list[dict[str, Any]] = []
+    remainder: list[dict[str, Any]] = []
+
+    for r in sorted_records:
+        fp = r.get("file_path", "")
+        if fp not in seen_files:
+            seen_files.add(fp)
+            diverse.append(r)
+        else:
+            remainder.append(r)
+        if len(diverse) >= n:
+            break
+
+    result = diverse[:n]
+    if len(result) < n:
+        result += remainder[: n - len(result)]
+    return result
+
+
+def _score_dataset(
+    model: JEPAArgot,
+    vectorizer: Any,
+    records: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+    ctx_texts = [" ".join(t["text"] for t in r["context_before"]) for r in records]
+    hunk_texts = [" ".join(t["text"] for t in r["hunk_tokens"]) for r in records]
+    ctx_x = torch.tensor(vectorizer.transform(ctx_texts).toarray(), dtype=torch.float32)
+    hunk_x = torch.tensor(vectorizer.transform(hunk_texts).toarray(), dtype=torch.float32)
+    model.eval()
+    scored = []
+    with torch.no_grad():
+        for i, record in enumerate(records):
+            score = model.surprise(ctx_x[i : i + 1], hunk_x[i : i + 1]).item()
+            scored.append({**record, "_score": score})
+    return scored
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Explain style anomalies in a git ref")
+    parser.add_argument("repo_path")
+    parser.add_argument("ref")
+    parser.add_argument("--model", default=".argot/model.pkl")
+    parser.add_argument("--dataset", default=".argot/dataset.jsonl")
+    parser.add_argument("--threshold-percentile", type=float, default=75.0)
+    parser.add_argument("--examples", type=int, default=5)
+    args = parser.parse_args()
+
+    model_path = Path(args.model)
+    if not model_path.exists():
+        print(f"error: model not found at {model_path}", file=sys.stderr)
+        sys.exit(2)
+
+    dataset_path = Path(args.dataset)
+    if not dataset_path.exists():
+        print(f"error: dataset not found at {dataset_path}", file=sys.stderr)
+        sys.exit(2)
+
+    bundle = joblib.load(model_path)
+    vectorizer = bundle["vectorizer"]
+    input_dim: int = bundle["input_dim"]
+    embed_dim: int = bundle["embed_dim"]
+
+    encoder = TokenEncoder(input_dim, embed_dim)
+    encoder.load_state_dict(bundle["encoder_state"])
+    predictor = ArgotPredictor(embed_dim=embed_dim)
+    predictor.load_state_dict(bundle["predictor_state"])
+    model = JEPAArgot(encoder, predictor)
+    model.eval()
+
+    raw_dataset = [
+        json.loads(line) for line in dataset_path.read_text().splitlines() if line.strip()
+    ]
+    scored_dataset = _score_dataset(model, vectorizer, raw_dataset)
+    distribution = [r["_score"] for r in scored_dataset]
+    style_examples = select_style_examples(scored_dataset, n=args.examples)
+    example_texts = [" ".join(t["text"] for t in r["hunk_tokens"]) for r in style_examples]
+
+    repo = pygit2.Repository(args.repo_path)
+    shas = _resolve_shas(repo, args.ref)
+    if not shas:
+        sys.exit(0)
+
+    context_lines = 50
+    with torch.no_grad():
+        for commit, file_path, post_blob, hunks in walk_commits(args.repo_path, shas):
+            lang = language_for_path(file_path)
+            if lang is None:
+                continue
+            try:
+                source_lines = post_blob.decode("utf-8", errors="replace").splitlines()
+            except Exception:
+                continue
+
+            for hunk in hunks:
+                hunk_start = hunk.new_start - 1
+                hunk_end = hunk_start + hunk.new_lines
+                if hunk_start < 0 or hunk_end > len(source_lines):
+                    continue
+
+                before_start = max(0, hunk_start - context_lines)
+                ctx_tokens = tokenize_lines(source_lines, lang, before_start, hunk_start)
+                hunk_tokens = tokenize_lines(source_lines, lang, hunk_start, hunk_end)
+
+                ctx_text = " ".join(t.text for t in ctx_tokens)
+                hunk_text = " ".join(t.text for t in hunk_tokens)
+
+                ctx_vec = torch.tensor(
+                    vectorizer.transform([ctx_text]).toarray(), dtype=torch.float32
+                )
+                hunk_vec = torch.tensor(
+                    vectorizer.transform([hunk_text]).toarray(), dtype=torch.float32
+                )
+
+                score = model.surprise(ctx_vec, hunk_vec).item()
+                pct = percentile_rank(score, distribution)
+
+                if pct < args.threshold_percentile:
+                    continue
+
+                print(
+                    json.dumps(
+                        {
+                            "file_path": file_path,
+                            "line": hunk.new_start,
+                            "commit": str(commit.id)[:8],
+                            "surprise": round(score, 4),
+                            "percentile": round(pct, 1),
+                            "hunk_text": hunk_text,
+                            "context_text": ctx_text,
+                            "style_examples": example_texts,
+                        }
+                    )
+                )
+
+
+if __name__ == "__main__":
+    main()

argot/extract.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from dataclasses import asdict
+from pathlib import Path
+
+import pygit2
+
+from argot.dataset import HunkRecord
+from argot.git_walk import walk_repo
+from argot.tokenize import language_for_path
+
+CONTEXT_LINES = 50
+
+
+def _extract_context(
+    source_lines: list[str],
+    hunk_start: int,
+    hunk_end: int,
+) -> tuple[list[str], list[str], list[str]]:
+    before_start = max(0, hunk_start - CONTEXT_LINES)
+    after_end = min(len(source_lines), hunk_end + CONTEXT_LINES)
+    return (
+        source_lines[before_start:hunk_start],
+        source_lines[hunk_start:hunk_end],
+        source_lines[hunk_end:after_end],
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Extract dataset from git history")
+    parser.add_argument("repo_path", help="Path to git repository")
+    parser.add_argument("--out", default=".argot/dataset.jsonl", help="Output JSONL path")
+    parser.add_argument("--limit", type=int, default=None, help="Max number of records to emit")
+    args = parser.parse_args()
+
+    repo_path = args.repo_path
+    out_path = Path(args.out)
+
+    try:
+        pygit2.Repository(repo_path)
+    except pygit2.GitError:
+        print(f"error: repository not found at {repo_path!r}", file=sys.stderr)
+        sys.exit(2)
+
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+
+    count = 0
+
+    with open(out_path, "w") as fh:
+        for commit, file_path, post_blob, hunks in walk_repo(repo_path):
+            lang = language_for_path(file_path)
+            if lang is None:
+                continue
+
+            try:
+                source_lines = post_blob.decode("utf-8", errors="replace").splitlines()
+            except Exception:
+                continue
+
+            parent_sha = str(commit.parents[0].id) if commit.parents else None
+            author_date_iso = commit.author.time
+
+            for hunk in hunks:
+                hunk_start = hunk.new_start - 1  # convert to 0-indexed
+                hunk_end = hunk_start + hunk.new_lines
+
+                if hunk_start < 0 or hunk_end > len(source_lines):
+                    continue
+
+                ctx_before_lines, hunk_lines, ctx_after_lines = _extract_context(
+                    source_lines, hunk_start, hunk_end
+                )
+
+                before_start_abs = max(0, hunk_start - CONTEXT_LINES)
+                after_start_abs = hunk_end
+
+                from argot.tokenize import tokenize_lines
+
+                context_before = tokenize_lines(source_lines, lang, before_start_abs, hunk_start)
+                hunk_tokens = tokenize_lines(source_lines, lang, hunk_start, hunk_end)
+                context_after = tokenize_lines(
+                    source_lines,
+                    lang,
+                    after_start_abs,
+                    min(len(source_lines), after_start_abs + CONTEXT_LINES),
+                )
+
+                record = HunkRecord(
+                    commit_sha=str(commit.id),
+                    file_path=file_path,
+                    language=lang,
+                    hunk_start_line=hunk_start,
+                    hunk_end_line=hunk_end,
+                    context_before=context_before,
+                    hunk_tokens=hunk_tokens,
+                    context_after=context_after,
+                    parent_sha=parent_sha,
+                    author_date_iso=str(author_date_iso),
+                )
+
+                fh.write(json.dumps(asdict(record)))
+                fh.write("\n")
+                count += 1
+
+                if args.limit is not None and count >= args.limit:
+                    print(f"Reached limit of {args.limit} records", file=sys.stderr)
+                    print(f"Wrote {count} records to {out_path}")
+                    return
+
+    if count == 0:
+        print("error: no hunks found — repository may have no history", file=sys.stderr)
+        sys.exit(2)
+
+    print(f"Wrote {count} records to {out_path}")
+
+
+if __name__ == "__main__":
+    main()

argot/fetch.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import subprocess
+import sys
+import tomllib
+from dataclasses import dataclass
+from pathlib import Path
+
+from argot.extract import main as extract_main
+
+
+@dataclass(frozen=True)
+class RepoConfig:
+    name: str
+    url: str
+    limit: int
+    label: str
+
+
+def load_repos_config(config_path: Path) -> list[RepoConfig]:
+    data = tomllib.loads(config_path.read_text())
+    return [RepoConfig(**r) for r in data["repos"]]
+
+
+def merge_jsonl(sources: list[tuple[Path, str]], dest: Path) -> int:
+    import json
+
+    count = 0
+    with dest.open("w") as fh:
+        for src, repo_name in sources:
+            for line in src.read_text().splitlines():
+                if line.strip():
+                    record = json.loads(line)
+                    record["_repo"] = repo_name
+                    fh.write(json.dumps(record) + "\n")
+                    count += 1
+    return count
+
+
+def _clone_or_update(url: str, dest: Path, depth: int = 2000) -> None:
+    if (dest / ".git").exists():
+        print(f"  already cloned, skipping: {dest}")
+        return
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    print(f"  cloning {url} → {dest}")
+    subprocess.run(
+        ["git", "clone", "--depth", str(depth), url, str(dest)],
+        check=True,
+    )
+
+
+def main() -> None:
+    project_root = Path(__file__).parents[2]  # engine/argot/fetch.py → project root
+    config_path = project_root / "training" / "repos.toml"
+
+    if not config_path.exists():
+        print(f"error: config not found at {config_path}", file=sys.stderr)
+        sys.exit(2)
+
+    repos = load_repos_config(config_path)
+    repos_dir = project_root / ".argot" / "repos"
+    part_files: list[tuple[Path, str]] = []
+
+    for repo in repos:
+        repo_dir = repos_dir / repo.name
+        part_out = repos_dir / f"{repo.name}.jsonl"
+        print(f"\n[{repo.name}] label={repo.label} limit={repo.limit}")
+        _clone_or_update(repo.url, repo_dir)
+
+        sys.argv = [
+            "extract",
+            str(repo_dir),
+            "--out",
+            str(part_out),
+            "--limit",
+            str(repo.limit),
+        ]
+        extract_main()
+        part_files.append((part_out, repo.name))
+
+    out = project_root / ".argot" / "training.jsonl"
+    total = merge_jsonl(part_files, out)
+    print(f"\nMerged {total} records → {out}")
+
+
+if __name__ == "__main__":
+    main()

argot/git_walk.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+from collections.abc import Iterator
+from pathlib import Path
+
+import pygit2
+
+SUPPORTED_EXTENSIONS = frozenset({".ts", ".tsx", ".js", ".jsx", ".py"})
+
+
+def _extension(path: str) -> str:
+    return Path(path).suffix.lower()
+
+
+def walk_repo(repo_path: str) -> Iterator[tuple[pygit2.Commit, str, bytes, list[pygit2.DiffHunk]]]:
+    """Yield (commit, file_path, post_blob_content, hunks) for each changed supported file.
+
+    Skips merge commits. Yields nothing if the repo has no commits.
+    """
+    repo = pygit2.Repository(repo_path)
+    if repo.is_empty:
+        return
+
+    try:
+        start_oid = repo.head.target
+    except pygit2.GitError:
+        branch_refs = [r for r in repo.references if r.startswith("refs/heads/")]
+        if not branch_refs:
+            return
+        start_oid = repo.references[branch_refs[0]].target
+
+    for commit in repo.walk(start_oid, pygit2.enums.SortMode.TOPOLOGICAL):
+        if len(commit.parents) != 1:
+            # Skip merge commits and root commits
+            continue
+
+        parent = commit.parents[0]
+        diff = parent.tree.diff_to_tree(commit.tree)
+        diff.find_similar()
+
+        for patch in diff:
+            if patch is None:
+                continue
+            file_path = patch.delta.new_file.path
+            if _extension(file_path) not in SUPPORTED_EXTENSIONS:
+                continue
+
+            hunks = list(patch.hunks)
+            if not hunks:
+                continue
+
+            try:
+                obj = commit.tree / file_path
+                if not isinstance(obj, pygit2.Blob):
+                    continue
+                post_blob_content = obj.data
+            except KeyError:
+                # File deleted in this commit
+                continue
+
+            yield commit, file_path, post_blob_content, hunks
+
+
+def walk_commits(
+    repo_path: str, shas: set[str]
+) -> Iterator[tuple[pygit2.Commit, str, bytes, list[pygit2.DiffHunk]]]:
+    """Yield (commit, file_path, post_blob_content, hunks) for commits in shas."""
+    repo = pygit2.Repository(repo_path)
+    if repo.is_empty:
+        return
+
+    try:
+        start_oid = repo.head.target
+    except pygit2.GitError:
+        branch_refs = [r for r in repo.references if r.startswith("refs/heads/")]
+        if not branch_refs:
+            return
+        start_oid = repo.references[branch_refs[0]].target
+
+    for commit in repo.walk(start_oid, pygit2.enums.SortMode.TOPOLOGICAL):
+        if str(commit.id) not in shas:
+            continue
+        if len(commit.parents) != 1:
+            continue
+
+        parent = commit.parents[0]
+        diff = parent.tree.diff_to_tree(commit.tree)
+        diff.find_similar()
+
+        for patch in diff:
+            if patch is None:
+                continue
+            file_path = patch.delta.new_file.path
+            if _extension(file_path) not in SUPPORTED_EXTENSIONS:
+                continue
+
+            hunks = list(patch.hunks)
+            if not hunks:
+                continue
+
+            try:
+                obj = commit.tree / file_path
+                if not isinstance(obj, pygit2.Blob):
+                    continue
+                post_blob_content = obj.data
+            except KeyError:
+                continue
+
+            yield commit, file_path, post_blob_content, hunks

argot/jepa/encoder.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from typing import cast
+
+import torch
+from torch import nn
+
+
+class TokenEncoder(nn.Module):
+    """MLP encoder from TF-IDF vectors to latent embeddings."""
+
+    def __init__(self, input_dim: int = 5000, embed_dim: int = 192) -> None:
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, 512),
+            nn.GELU(),
+            nn.Linear(512, embed_dim),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """x: (B, input_dim) → (B, embed_dim)"""
+        return cast(torch.Tensor, self.net(x))