jfl 0.4.4 → 0.6.0

package/scripts/train/v2/eval.py

@@ -0,0 +1,196 @@
+"""
+v2 Policy Head Evaluation — per-tool accuracy, confusion matrix, confidence analysis.
+"""
+
+import json
+import os
+import sys
+import argparse
+from collections import defaultdict
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+
+from model import PolicyHead
+from dataset import PolicyHeadDataset, load_embedding_cache
+
+
+def load_model(checkpoint_path: str, device: str = "cpu") -> tuple:
+    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    config = ckpt["config"]
+    tool_to_index = ckpt["tool_to_index"]
+    index_to_tool = {v: k for k, v in tool_to_index.items()}
+
+    model = PolicyHead(
+        embedding_dim=config["embedding_dim"],
+        hidden_dim=config["hidden_dim"],
+        num_tools=ckpt["num_tools"],
+        num_layers=config["num_layers"],
+        num_heads=config["num_heads"],
+        dropout=config.get("dropout", 0.1),
+    ).to(device)
+
+    model.load_state_dict(ckpt["model_state_dict"])
+    model.eval()
+
+    return model, tool_to_index, index_to_tool
+
+
+@torch.no_grad()
+def evaluate_detailed(model, dataloader, index_to_tool, device):
+    model.eval()
+
+    per_tool_correct = defaultdict(int)
+    per_tool_total = defaultdict(int)
+    confusion_pairs = defaultdict(int)
+    all_predictions = []
+
+    for batch in dataloader:
+        state_emb = batch["state_emb"].to(device)
+        goal_emb = batch["goal_emb"].to(device)
+        labels = batch["label"].to(device)
+        tool_names = batch["tool_name"]
+
+        logits = model(state_emb, goal_emb)
+        probs = torch.softmax(logits, dim=-1)
+        preds = logits.argmax(dim=-1)
+
+        for i in range(labels.size(0)):
+            true_tool = tool_names[i]
+            pred_tool = index_to_tool[preds[i].item()]
+            confidence = probs[i, preds[i]].item()
+
+            per_tool_total[true_tool] += 1
+            if preds[i] == labels[i]:
+                per_tool_correct[true_tool] += 1
+            else:
+                confusion_pairs[(true_tool, pred_tool)] += 1
+
+            all_predictions.append({
+                "true": true_tool,
+                "predicted": pred_tool,
+                "correct": preds[i].item() == labels[i].item(),
+                "confidence": confidence,
+            })
+
+    return per_tool_correct, per_tool_total, confusion_pairs, all_predictions
+
+
+def print_report(per_tool_correct, per_tool_total, confusion_pairs, all_predictions):
+    total_correct = sum(per_tool_correct.values())
+    total = sum(per_tool_total.values())
+    overall_acc = total_correct / max(total, 1)
+
+    print("=" * 70)
+    print("  v2 POLICY HEAD — EVALUATION REPORT")
+    print("=" * 70)
+    print(f"\n  Overall Accuracy: {overall_acc:.1%} ({total_correct}/{total})")
+    print(f"\n{'Tool':<30} {'Accuracy':>10} {'Correct':>9} {'Total':>7}")
+    print("-" * 60)
+
+    for tool in sorted(per_tool_total.keys()):
+        correct = per_tool_correct.get(tool, 0)
+        total_t = per_tool_total[tool]
+        acc = correct / total_t if total_t > 0 else 0
+        bar = "█" * int(acc * 20) + "░" * (20 - int(acc * 20))
+        print(f"  {tool:<28} {acc:>8.1%}  {correct:>5}/{total_t:<5} {bar}")
+
+    if confusion_pairs:
+        print(f"\n{'Top Confusion Pairs':<45} {'Count':>7}")
+        print("-" * 55)
+        for (true_t, pred_t), count in sorted(
+            confusion_pairs.items(), key=lambda x: -x[1]
+        )[:10]:
+            print(f"  {true_t} -> {pred_t:<25} {count:>5}")
+
+    correct_confs = [p["confidence"] for p in all_predictions if p["correct"]]
+    wrong_confs = [p["confidence"] for p in all_predictions if not p["correct"]]
+
+    print(f"\n  Confidence Analysis:")
+    if correct_confs:
+        print(f"    Correct predictions:  avg={sum(correct_confs) / len(correct_confs):.3f}")
+    if wrong_confs:
+        print(f"    Wrong predictions:    avg={sum(wrong_confs) / len(wrong_confs):.3f}")
+
+    if overall_acc < 0.80:
+        print(f"\n  ⚠️  Accuracy below 80% — consider more data or warm-start")
+    elif overall_acc < 0.90:
+        print(f"\n  ⚠️  Accuracy below 90% — good but can improve")
+    else:
+        print(f"\n  ✅ Accuracy ≥90% — ready for deployment")
+
+    print("=" * 70)
+
+    return {
+        "overall_accuracy": overall_acc,
+        "total": total,
+        "per_tool": {
+            tool: {
+                "accuracy": per_tool_correct.get(tool, 0) / per_tool_total[tool] if per_tool_total[tool] > 0 else 0,
+                "correct": per_tool_correct.get(tool, 0),
+                "total": per_tool_total[tool],
+            }
+            for tool in per_tool_total
+        },
+        "top_confusion": [
+            {"true": t, "predicted": p, "count": c}
+            for (t, p), c in sorted(confusion_pairs.items(), key=lambda x: -x[1])[:10]
+        ],
+        "avg_confidence_correct": sum(correct_confs) / len(correct_confs) if correct_confs else 0,
+        "avg_confidence_wrong": sum(wrong_confs) / len(wrong_confs) if wrong_confs else 0,
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Evaluate v2 policy head")
+    parser.add_argument("--checkpoint", default=".jfl/checkpoints/best_policy_head.pt", help="Path to checkpoint")
+    parser.add_argument("--data-dir", default=".jfl/v2-data", help="Directory with test.jsonl + embeddings")
+    parser.add_argument("--split", default="test", help="Which split to evaluate (train/val/test)")
+    parser.add_argument("--json", action="store_true", help="Output results as JSON")
+    args = parser.parse_args()
+
+    if not os.path.exists(args.checkpoint):
+        print(f"Checkpoint not found: {args.checkpoint}")
+        sys.exit(1)
+
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        device = "mps"
+    else:
+        device = "cpu"
+
+    model, tool_to_index, index_to_tool = load_model(args.checkpoint, device)
+
+    data_path = os.path.join(args.data_dir, f"{args.split}.jsonl")
+    if not os.path.exists(data_path):
+        print(f"Data not found: {data_path}")
+        sys.exit(1)
+
+    embeddings_matrix, text_to_idx = load_embedding_cache(args.data_dir)
+
+    test_ds = PolicyHeadDataset(data_path, tool_to_index, embeddings_matrix, text_to_idx)
+    test_loader = DataLoader(test_ds, batch_size=64, shuffle=False)
+
+    print(f"Evaluating on {args.split} split: {len(test_ds)} examples")
+    print(f"Device: {device}\n")
+
+    per_tool_correct, per_tool_total, confusion_pairs, all_predictions = evaluate_detailed(
+        model, test_loader, index_to_tool, device
+    )
+
+    results = print_report(per_tool_correct, per_tool_total, confusion_pairs, all_predictions)
+
+    if args.json:
+        print(json.dumps(results, indent=2))
+
+    # Write eval results alongside checkpoint
+    results_path = os.path.join(os.path.dirname(args.checkpoint), "eval-results.json")
+    with open(results_path, "w") as f:
+        json.dump(results, f, indent=2)
+    print(f"\nResults saved to: {results_path}")
+
+
+if __name__ == "__main__":
+    main()

package/scripts/train/v2/generate_data.py

@@ -0,0 +1,219 @@
+"""
+Synthetic training data generator for v2 policy head.
+
+Generates (current_state, goal, correct_tool) tuples using LLM API.
+Adapted from Drew's generate_data.py for JFL's action taxonomy.
+"""
+
+import json
+import os
+import sys
+import random
+import argparse
+import time
+from pathlib import Path
+
+def load_domain(path: str) -> dict:
+    with open(path) as f:
+        return json.load(f)
+
+def get_llm_client():
+    api_key = os.environ.get("OPENAI_API_KEY") or os.environ.get("ANTHROPIC_API_KEY")
+    if not api_key:
+        print("Set OPENAI_API_KEY or ANTHROPIC_API_KEY for synthetic data generation")
+        sys.exit(1)
+
+    if os.environ.get("OPENAI_API_KEY"):
+        from openai import OpenAI
+        return OpenAI(), "openai"
+    else:
+        from anthropic import Anthropic
+        return Anthropic(), "anthropic"
+
+def generate_goals_for_tool(client, client_type: str, tool: dict, all_tools: list, n: int = 50) -> list:
+    tool_list_str = "\n".join(
+        f"- {t['name']}: {t['description']}" for t in all_tools
+    )
+
+    prompt = f"""You are generating training data for an AI action-routing system in software development.
+
+Given these {len(all_tools)} available action types:
+{tool_list_str}
+
+Generate exactly {n} realistic goals/situations that should route to:
+  Action: {tool["name"]}
+  Description: {tool["description"]}
+
+Requirements:
+1. Each goal should describe a real software development situation or user request
+2. Include variety: specific bug reports, vague feature requests, urgent fixes, routine maintenance
+3. Some should be indirect (describe the symptom, not the action needed)
+4. Some should use technical jargon (CI/CD, linting, test coverage, etc.)
+5. Some should include context about current project state
+6. Include goals that could be confused with OTHER action types
+7. Do NOT include the action type name in the goal text
+8. Goals should reflect what an AI coding agent would encounter
+
+Output as a JSON object with a "goals" key containing an array of strings. No explanation."""
+
+    try:
+        if client_type == "openai":
+            response = client.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=[{"role": "user", "content": prompt}],
+                temperature=0.9,
+                response_format={"type": "json_object"},
+            )
+            result = json.loads(response.choices[0].message.content)
+        else:
+            response = client.messages.create(
+                model="claude-sonnet-4-5-20250514",
+                max_tokens=4096,
+                messages=[{"role": "user", "content": prompt}],
+            )
+            text = response.content[0].text
+            json_match = text[text.index("{"):text.rindex("}") + 1]
+            result = json.loads(json_match)
+
+        goals = result.get("goals", result.get("queries", []))
+        return goals[:n] if isinstance(goals, list) else []
+    except Exception as e:
+        print(f"  Error generating goals for {tool['name']}: {e}")
+        return []
+
+def generate_states_for_goal(client, client_type: str, goal: str, domain_desc: str, n: int = 2) -> list:
+    prompt = f"""You are generating training data for an AI planning system in:
+"{domain_desc}"
+
+Given this development goal: "{goal}"
+
+Generate {n} different realistic "current state" descriptions that an AI coding agent might be in
+when this goal arises. Each state should describe:
+- Current codebase state (test results, scores, recent changes)
+- What the agent currently knows about the project
+- Any relevant trajectory information (recent actions taken)
+
+States should vary in how much context is available and what prior work has been done.
+
+Output as a JSON object with a "states" key containing an array of strings (1-3 sentences each)."""
+
+    try:
+        if client_type == "openai":
+            response = client.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=[{"role": "user", "content": prompt}],
+                temperature=0.8,
+                response_format={"type": "json_object"},
+            )
+            result = json.loads(response.choices[0].message.content)
+        else:
+            response = client.messages.create(
+                model="claude-sonnet-4-5-20250514",
+                max_tokens=2048,
+                messages=[{"role": "user", "content": prompt}],
+            )
+            text = response.content[0].text
+            json_match = text[text.index("{"):text.rindex("}") + 1]
+            result = json.loads(json_match)
+
+        states = result.get("states", [])
+        return states[:n]
+    except Exception as e:
+        print(f"  Error generating states: {e}")
+        return [f"Agent working on codebase. Goal context: {goal[:100]}"]
+
+def generate_dataset(
+    domain_path: str,
+    goals_per_tool: int = 50,
+    states_per_goal: int = 2,
+    output_dir: str = "data",
+    seed: int = 42,
+):
+    domain = load_domain(domain_path)
+    tools = domain["tools"]
+    domain_desc = domain["description"]
+
+    client, client_type = get_llm_client()
+    print(f"Using {client_type} API for generation")
+
+    all_examples = []
+
+    for tool_idx, tool in enumerate(tools):
+        print(f"[{tool_idx + 1}/{len(tools)}] Generating goals for: {tool['name']}")
+
+        goals = generate_goals_for_tool(client, client_type, tool, tools, n=goals_per_tool)
+        print(f"  Generated {len(goals)} goals")
+
+        for goal_idx, goal in enumerate(goals):
+            if goal_idx % 10 == 0:
+                print(f"  Generating states: {goal_idx}/{len(goals)}")
+
+            states = generate_states_for_goal(client, client_type, goal, domain_desc, n=states_per_goal)
+
+            for state in states:
+                example = {
+                    "current_state": state,
+                    "goal": goal,
+                    "correct_tool": tool["name"],
+                    "tool_category": tool["category"],
+                    "source": "synthetic",
+                }
+                all_examples.append(example)
+
+            time.sleep(0.1)
+
+    random.seed(seed)
+    random.shuffle(all_examples)
+
+    n = len(all_examples)
+    train_end = int(n * 0.7)
+    val_end = int(n * 0.85)
+
+    splits = {
+        "train": all_examples[:train_end],
+        "val": all_examples[train_end:val_end],
+        "test": all_examples[val_end:],
+    }
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    for split_name, split_data in splits.items():
+        path = os.path.join(output_dir, f"{split_name}.jsonl")
+        with open(path, "w") as f:
+            for ex in split_data:
+                f.write(json.dumps(ex) + "\n")
+        print(f"  {split_name}: {len(split_data)} examples -> {path}")
+
+    print(f"\nTotal: {n} examples across {len(tools)} tools")
+    print(f"  Per tool: ~{n // len(tools)} examples")
+    return splits
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate synthetic training data for v2 policy head")
+    parser.add_argument("--domain", default=None, help="Path to domain.json")
+    parser.add_argument("--output", default=".jfl/v2-data", help="Output directory")
+    parser.add_argument("--goals-per-tool", type=int, default=50, help="Goals to generate per tool")
+    parser.add_argument("--states-per-goal", type=int, default=2, help="States to generate per goal")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    args = parser.parse_args()
+
+    domain_path = args.domain
+    if domain_path is None:
+        domain_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "domain.json")
+
+    if not os.path.exists(domain_path):
+        print(f"Domain file not found: {domain_path}")
+        sys.exit(1)
+
+    generate_dataset(
+        domain_path=domain_path,
+        goals_per_tool=args.goals_per_tool,
+        states_per_goal=args.states_per_goal,
+        output_dir=args.output,
+        seed=args.seed,
+    )
+
+
+if __name__ == "__main__":
+    main()

package/scripts/train/v2/infer.py

@@ -0,0 +1,188 @@
+"""
+v2 Policy Head Inference — CLI script for action selection.
+
+Usage:
+  python infer.py --checkpoint .jfl/checkpoints/best_policy_head.pt --state "..." --goal "..." --top-k 3
+
+Also supports JSON mode for TypeScript bridge:
+  python infer.py --checkpoint ... --state "..." --goal "..." --json
+"""
+
+import json
+import os
+import sys
+import argparse
+
+import torch
+import numpy as np
+
+from model import PolicyHead
+
+
+def load_model(checkpoint_path: str, device: str = "cpu"):
+    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    config = ckpt["config"]
+    tool_to_index = ckpt["tool_to_index"]
+    index_to_tool = {v: k for k, v in tool_to_index.items()}
+
+    model = PolicyHead(
+        embedding_dim=config["embedding_dim"],
+        hidden_dim=config["hidden_dim"],
+        num_tools=ckpt["num_tools"],
+        num_layers=config["num_layers"],
+        num_heads=config["num_heads"],
+        dropout=config.get("dropout", 0.1),
+    ).to(device)
+
+    model.load_state_dict(ckpt["model_state_dict"])
+    model.eval()
+
+    return model, tool_to_index, index_to_tool, config
+
+
+def get_embedding(text: str, api_url: str, api_key: str) -> list[float]:
+    import requests
+
+    response = requests.post(
+        f"{api_url}/v1/embeddings",
+        headers={
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+        },
+        json={
+            "model": "stratus-x1ac-base",
+            "input": text,
+        },
+        timeout=15,
+    )
+    response.raise_for_status()
+    data = response.json()
+    return data["data"][0]["embedding"]
+
+
+def infer(args):
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        device = "mps"
+    else:
+        device = "cpu"
+
+    model, tool_to_index, index_to_tool, config = load_model(args.checkpoint, device)
+
+    api_url = os.environ.get("STRATUS_API_URL", "https://api.stratus.run")
+    api_key = os.environ.get("STRATUS_API_KEY", "")
+
+    if not api_key:
+        print("STRATUS_API_KEY not set", file=sys.stderr)
+        sys.exit(1)
+
+    state_emb = get_embedding(args.state, api_url, api_key)
+    goal_emb = get_embedding(args.goal, api_url, api_key)
+
+    state_tensor = torch.tensor([state_emb], dtype=torch.float32).to(device)
+    goal_tensor = torch.tensor([goal_emb], dtype=torch.float32).to(device)
+
+    result = model.predict(state_tensor, goal_tensor, top_k=args.top_k)
+
+    top_indices = result["top_k_indices"][0].cpu().tolist()
+    top_probs = result["top_k_probs"][0].cpu().tolist()
+
+    predictions = []
+    for idx, prob in zip(top_indices, top_probs):
+        predictions.append({
+            "action": index_to_tool[idx],
+            "confidence": round(prob, 4),
+        })
+
+    if args.json:
+        output = {
+            "action": predictions[0]["action"],
+            "confidence": predictions[0]["confidence"],
+            "alternatives": predictions[1:],
+        }
+        print(json.dumps(output))
+    else:
+        print(f"\nv2 Policy Head Prediction")
+        print(f"{'─' * 40}")
+        print(f"State: {args.state[:80]}...")
+        print(f"Goal:  {args.goal[:80]}...")
+        print(f"\nTop {args.top_k} actions:")
+        for i, pred in enumerate(predictions):
+            marker = "→" if i == 0 else " "
+            bar = "█" * int(pred["confidence"] * 20)
+            print(f"  {marker} {pred['action']:25s} {pred['confidence']:6.1%} {bar}")
+
+
+def batch_infer(args):
+    """Batch inference from JSONL input on stdin. Each line: {"state": "...", "goal": "..."}"""
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        device = "mps"
+    else:
+        device = "cpu"
+
+    model, tool_to_index, index_to_tool, config = load_model(args.checkpoint, device)
+
+    api_url = os.environ.get("STRATUS_API_URL", "https://api.stratus.run")
+    api_key = os.environ.get("STRATUS_API_KEY", "")
+
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+
+        try:
+            req = json.loads(line)
+            state_emb = get_embedding(req["state"], api_url, api_key)
+            goal_emb = get_embedding(req["goal"], api_url, api_key)
+
+            state_tensor = torch.tensor([state_emb], dtype=torch.float32).to(device)
+            goal_tensor = torch.tensor([goal_emb], dtype=torch.float32).to(device)
+
+            result = model.predict(state_tensor, goal_tensor, top_k=args.top_k)
+
+            top_indices = result["top_k_indices"][0].cpu().tolist()
+            top_probs = result["top_k_probs"][0].cpu().tolist()
+
+            predictions = []
+            for idx, prob in zip(top_indices, top_probs):
+                predictions.append({
+                    "action": index_to_tool[idx],
+                    "confidence": round(prob, 4),
+                })
+
+            output = {
+                "action": predictions[0]["action"],
+                "confidence": predictions[0]["confidence"],
+                "alternatives": predictions[1:],
+            }
+            print(json.dumps(output))
+            sys.stdout.flush()
+        except Exception as e:
+            print(json.dumps({"error": str(e)}))
+            sys.stdout.flush()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="v2 policy head inference")
+    parser.add_argument("--checkpoint", required=True, help="Path to .pt checkpoint")
+    parser.add_argument("--state", default=None, help="Current state text")
+    parser.add_argument("--goal", default=None, help="Goal text")
+    parser.add_argument("--top-k", type=int, default=3, help="Number of top actions")
+    parser.add_argument("--json", action="store_true", help="JSON output for TypeScript bridge")
+    parser.add_argument("--batch", action="store_true", help="Batch mode: read JSONL from stdin")
+    args = parser.parse_args()
+
+    if args.batch:
+        batch_infer(args)
+    elif args.state and args.goal:
+        infer(args)
+    else:
+        print("Provide --state and --goal, or use --batch for stdin JSONL")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()