Add xagent labels to router training

Author: qinxuye

AGENTS.md

@@ -34,8 +34,8 @@ model will usually win raw score while losing the cost objective.
 
 ## Current Data Strategy
 
-The production difficulty model is trained on **three datasets combined**
-(377,997 rows / ~14,364 prompts / 283 subjects):
+The production difficulty model is trained on **four datasets combined**
+(378,397 rows / ~14,463 prompts / 287 subjects):
 
 ```text
 NPULH/LLMRouterBench (350k stream sample)   37 models x ~13,775 prompts, single-turn
@@ -45,20 +45,23 @@ agent-psychometrics Terminal-Bench 2.0      89 tasks x 112 subjects, terminal ag
   data/agentic/terminalbench  (--dataset agentic:agentic/terminalbench)
 agent-psychometrics SWE-bench Verified      500 tasks x 134 subjects, coding agent
   data/agentic/swebench_verified  (task text joined from princeton-nlp/SWE-bench_Verified)
+Xorbits/xagent-xrouter-labels                100 prompts x 4 OpenRouter models, real xagent
+  --dataset xagent-labels:Xorbits/xagent-xrouter-labels:full
 ```
 
-Only the 37 llmrouterbench models have benchmark profiles, so they alone feed
-the capability/combine logistic; the agentic subjects (agent+scaffold combos)
-feed the difficulty axis only. agent-psychometrics swebench_pro (730x14) and
-gso (102x15) are loadable but NOT wired in (they ship no local task text and
-need an external join).
+Only models with benchmark profiles feed the capability/combine logistic: the
+37 llmrouterbench models plus profiled xagent OpenRouter candidates. The
+agent-psychometrics subjects (agent+scaffold combos) feed the difficulty axis
+only. agent-psychometrics swebench_pro (730x14) and gso (102x15) are loadable
+but NOT wired in (they ship no local task text and need an external join).
 
 - agent-psychometrics matrices load via `agentic.py` and the CLI `agentic:`
   dataset kind. **terminalbench** (89x112, local `tasks.jsonl`) and
   **swebench_verified** (500x134, task text joined from
   `princeton-nlp/SWE-bench_Verified`) are wired in; see "Agentic training data".
 - All sources feed the difficulty axis; only the profiled llmrouterbench models
-  feed the capability/combine logistic (agentic subjects have no profile).
+  plus profiled xagent OpenRouter candidates feed the capability/combine
+  logistic (agentic subjects have no profile).
 - RouterBench (`withmartian/routerbench`) is kept as a smaller legacy baseline.
 - `swebench_pro` (730x14), `gso` (102x15) ship no local task text and need an
   external join (not wired yet).
@@ -242,7 +245,8 @@ PYTHONPATH=src python3 -m xrouter_llm.cli train-irt \
   --dataset llmrouterbench:data/raw/llmrouterbench_stream_sample_350k \
   --dataset agentic:agentic/terminalbench \
   --dataset agentic:agentic/swebench_verified \
-  --benchmark-profiles artifacts/profiles/llmrouterbench_350k_profiles_priority_collected.json \
+  --dataset xagent-labels:Xorbits/xagent-xrouter-labels:full \
+  --benchmark-profiles artifacts/profiles/llmrouterbench_350k_profiles_priority_collected.json,src/xrouter_llm/resources/config/models \
   --output artifacts/models/irt_router_350k.joblib
 ```
 
@@ -270,7 +274,8 @@ PYTHONPATH=src python3 -m xrouter_llm.cli train-irt \
   --dataset llmrouterbench:data/raw/llmrouterbench_stream_sample_350k \
   --dataset agentic:agentic/terminalbench \
   --dataset agentic:agentic/swebench_verified \
-  --benchmark-profiles artifacts/profiles/llmrouterbench_350k_profiles_priority_collected.json \
+  --dataset xagent-labels:Xorbits/xagent-xrouter-labels:full \
+  --benchmark-profiles artifacts/profiles/llmrouterbench_350k_profiles_priority_collected.json,src/xrouter_llm/resources/config/models \
   --output artifacts/models/irt_router_350k.joblib
 ```
 
@@ -291,7 +296,8 @@ PYTHONPATH=src python3 -m xrouter_llm.cli train-irt \
   --dataset llmrouterbench:data/raw/llmrouterbench_stream_sample_350k \
   --dataset agentic:agentic/terminalbench \
   --dataset agentic:agentic/swebench_verified \
-  --benchmark-profiles artifacts/profiles/llmrouterbench_350k_profiles_priority_collected.json \
+  --dataset xagent-labels:Xorbits/xagent-xrouter-labels:full \
+  --benchmark-profiles artifacts/profiles/llmrouterbench_350k_profiles_priority_collected.json,src/xrouter_llm/resources/config/models \
   --output artifacts/models/irt_router_350k.joblib
 ```
 
@@ -301,6 +307,8 @@ PYTHONPATH=src python3 -m xrouter_llm.cli train-irt \
   joined from `princeton-nlp/SWE-bench_Verified` `problem_statement` into a local
   `data/agentic/swebench_verified/tasks.jsonl` (regenerate on a fresh checkout
   via `datasets.load_dataset("princeton-nlp/SWE-bench_Verified", split="test")`).
+- **xagent labels** are loaded from `Xorbits/xagent-xrouter-labels` via
+  `xagent_labels.py` and the CLI `xagent-labels:` dataset kind.
 - **swebench_pro** (730x14), **gso** (102x15) ship no local task text and need
   an external join -- not wired yet.
 
@@ -312,12 +320,10 @@ agentic subjects have no
 benchmark profiles, so they feed ONLY the difficulty axis; the capability/combine
 logistic still fits on profiled models.
 
-Limitation (verified): real xagent prompts (e.g. Chinese business + image-gen
-agentic tasks) are NOT covered by SWE-bench/Terminal-Bench either, so they stay
-out-of-distribution and get a near-max (clamped) difficulty. Difficulty is
-clamped to the training range so P never collapses, but the only way to make it
-accurate for a specific task mix is that deployment's own logged
-prompts + outcomes.
+Limitation (verified): the public xagent labels are only a 100-prompt seed. They
+slightly improve the public benchmark headline under a controlled base test set,
+but real deployment accuracy for a specific task mix still needs that
+deployment's own logged prompts + outcomes.
 
 ## Evaluation Rules
 

README.md

@@ -3,16 +3,16 @@
 
 # xrouter-llm
 
-<img src="./assets/xrouter-llm-hero-clean.png" alt="xrouter-llm: 52.4% lower cost and +1.7 pts completion on our tested dataset" />
+<img src="./assets/xrouter-llm-hero-clean.png" alt="xrouter-llm: 53.2% lower cost and +1.9 pts completion on our tested dataset" />
 
 </div>
 
 Stop sending every prompt to your most expensive LLM.
 
 `xrouter-llm` is a prompt-aware LLM **routing-decision** service: it predicts
 which models can complete a prompt, then chooses the cheapest model that clears
-the bar. On our tested dataset, it cuts realized cost by **52.4%** while
-improving completion by **+1.7 pts**.
+the bar. On our tested dataset, it cuts realized cost by **53.2%** while
+improving completion by **+1.9 pts**.
 
 It answers "which model should serve this prompt?" and records the choice — it
 does NOT call the underlying LLMs.
@@ -116,10 +116,12 @@ The production difficulty model is trained on **multiple datasets combined**
 | `NPULH/LLMRouterBench` (350k stream sample) | single-turn QA / code / math (22 tasks) | 37 models x ~13.8k prompts | ✅ |
 | agent-psychometrics — Terminal-Bench 2.0 | terminal agent | 89 tasks x 112 subjects | ✅ `--dataset agentic:agentic/terminalbench` |
 | agent-psychometrics — SWE-bench Verified | coding agent | 500 tasks x 134 subjects | ✅ task text joined from `princeton-nlp/SWE-bench_Verified` |
+| `Xorbits/xagent-xrouter-labels` | real xagent internal prompts | 100 prompts x 4 OpenRouter models | ✅ `--dataset xagent-labels:Xorbits/xagent-xrouter-labels:full` |
 | agent-psychometrics — SWE-bench Pro / GSO | coding agent | 730x14 / 102x15 | ⛔ ship no local task text, external join needed |
 
 The current artifact trains on LLMRouterBench 350k **+ Terminal-Bench +
-SWE-bench Verified** (377,997 rows / ~14,364 prompts / 283 subjects). The
+SWE-bench Verified + xagent labels** (378,397 rows / ~14,463 prompts /
+287 subjects). The
 agentic matrices come from
 [agent-psychometrics](https://github.com/dariakryvosheieva/agent-psychometrics)
 (MIT) via `agentic.py`. In `IRTRouter`, only the 37 profiled llmrouterbench
@@ -137,7 +139,8 @@ xrouter-llm train-irt \
   --dataset llmrouterbench:data/raw/llmrouterbench_stream_sample_350k \
   --dataset agentic:agentic/terminalbench \
   --dataset agentic:agentic/swebench_verified \
-  --benchmark-profiles artifacts/profiles/llmrouterbench_350k_profiles_priority_collected.json \
+  --dataset xagent-labels:Xorbits/xagent-xrouter-labels:full \
+  --benchmark-profiles artifacts/profiles/llmrouterbench_350k_profiles_priority_collected.json,src/xrouter_llm/resources/config/models \
   --output artifacts/models/irt_router_350k.joblib
 ```
 

assets/xrouter-llm-hero-clean.png

@@ -58,6 +58,7 @@
     RouteDecision,
     UtilityBreakdown,
 )
+from xrouter_llm.xagent_labels import load_xagent_labels, load_xagent_openrouter_labels
 
 __all__ = [
     "BenchmarkRow",
@@ -105,6 +106,8 @@
     "load_llmrouterbench",
     "load_routerbench_pickle",
     "load_jsonl",
+    "load_xagent_labels",
+    "load_xagent_openrouter_labels",
     "sample_llmrouterbench",
     "split_by_prompt",
 ]

src/xrouter_llm/__init__.py

@@ -65,7 +65,7 @@ def main(argv: list[str] | None = None) -> int:
     sweep_parser.add_argument("--input", default=None)
     sweep_parser.add_argument(
         "--format",
-        choices=["jsonl", "csv", "routerbench-pkl", "llmrouterbench"],
+        choices=["jsonl", "csv", "routerbench-pkl", "llmrouterbench", "xagent-labels"],
         default="jsonl",
     )
     _add_sweep_args(sweep_parser)
@@ -83,7 +83,7 @@ def main(argv: list[str] | None = None) -> int:
     holdout_parser.add_argument("--input", default=None)
     holdout_parser.add_argument(
         "--format",
-        choices=["jsonl", "csv", "routerbench-pkl", "llmrouterbench"],
+        choices=["jsonl", "csv", "routerbench-pkl", "llmrouterbench", "xagent-labels"],
         default="jsonl",
     )
     holdout_parser.add_argument(
@@ -104,7 +104,7 @@ def main(argv: list[str] | None = None) -> int:
     train_irt_parser.add_argument("--input", default=None)
     train_irt_parser.add_argument(
         "--format",
-        choices=["jsonl", "csv", "routerbench-pkl", "llmrouterbench"],
+        choices=["jsonl", "csv", "routerbench-pkl", "llmrouterbench", "xagent-labels"],
         default="llmrouterbench",
     )
     train_irt_parser.add_argument("--benchmark-profiles", default=default_models_dir())
@@ -416,6 +416,14 @@ def _load_dataset(kind: str, path: str, *, args: argparse.Namespace) -> list[obj
             max_prompts=args.max_prompts,
             random_state=args.random_state,
         )
+    if kind == "xagent-labels":
+        from xrouter_llm.xagent_labels import load_xagent_openrouter_labels
+
+        return limit_rows_by_prompt(
+            load_xagent_openrouter_labels(path),
+            max_prompts=args.max_prompts,
+            random_state=args.random_state,
+        )
     raise ValueError(f"Unsupported dataset kind {kind!r}")
 
 

src/xrouter_llm/cli.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Iterable
+
+from huggingface_hub import hf_hub_download
+
+from xrouter_llm.types import BenchmarkRow
+
+
+DEFAULT_XAGENT_LABEL_REPO = "Xorbits/xagent-xrouter-labels"
+
+XAGENT_LABEL_FILES = {
+    "full": "xagent_openrouter_official_candidates_100_4models.public.jsonl",
+    "train80": "xagent_openrouter_official_candidates_100_4models_train80.public.jsonl",
+    "holdout20": "xagent_openrouter_official_candidates_100_4models_holdout20.public.jsonl",
+    "sample": "xagent_openrouter_official_sample.public.jsonl",
+}
+
+
+def load_xagent_openrouter_labels(
+    source: str | Path = DEFAULT_XAGENT_LABEL_REPO,
+    *,
+    split: str = "full",
+) -> list[BenchmarkRow]:
+    """Load scrubbed xagent routing labels from a local JSONL file or HF dataset.
+
+    The public HF dataset currently contains plain JSONL files. Loading through
+    ``huggingface_hub`` keeps the dependency surface small and avoids relying on
+    a generated datasets schema.
+    """
+
+    path = Path(source).expanduser()
+    if path.exists():
+        return _load_xagent_jsonl(path)
+
+    repo_id, resolved_split = _parse_hf_source(str(source), split=split)
+    filename = XAGENT_LABEL_FILES.get(resolved_split, resolved_split)
+    downloaded = hf_hub_download(repo_id, filename, repo_type="dataset")
+    return _load_xagent_jsonl(Path(downloaded))
+
+
+def _parse_hf_source(source: str, *, split: str) -> tuple[str, str]:
+    if ":" not in source:
+        return source, split
+
+    repo_id, maybe_split = source.rsplit(":", 1)
+    if "/" not in maybe_split:
+        return repo_id, maybe_split
+    return source, split
+
+
+def _load_xagent_jsonl(path: Path) -> list[BenchmarkRow]:
+    output: list[BenchmarkRow] = []
+    for line_number, record in _iter_jsonl(path):
+        score = _score(record, line_number=line_number)
+        prompt = str(record.get("prompt") or "")
+        model_id = str(record.get("candidate_model") or "")
+        if not prompt:
+            raise ValueError(f"Missing prompt in xagent label line {line_number}")
+        if not model_id:
+            raise ValueError(f"Missing candidate_model in xagent label line {line_number}")
+
+        prompt_sha = record.get("prompt_sha256")
+        prompt_id = f"xagent:{prompt_sha}" if prompt_sha else _fallback_prompt_id(record)
+        output.append(
+            BenchmarkRow(
+                prompt_id=prompt_id,
+                prompt=prompt,
+                model_id=model_id,
+                score=score,
+                cost_usd=_usage_cost(record.get("candidate_usage")),
+                task=f"xagent:{record.get('category') or 'unknown'}",
+            )
+        )
+    if not output:
+        raise ValueError(f"No xagent labels loaded from {path}")
+    return output
+
+
+load_xagent_labels = load_xagent_openrouter_labels
+
+
+def _iter_jsonl(path: Path) -> Iterable[tuple[int, dict[str, Any]]]:
+    with path.open("r", encoding="utf-8") as file:
+        for line_number, line in enumerate(file, start=1):
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                value = json.loads(line)
+            except json.JSONDecodeError as exc:
+                raise ValueError(f"Invalid xagent JSONL at line {line_number}: {exc}") from exc
+            if not isinstance(value, dict):
+                raise ValueError(f"Invalid xagent JSONL at line {line_number}: expected object")
+            yield line_number, value
+
+
+def _score(record: dict[str, Any], *, line_number: int) -> float:
+    judge = record.get("judge")
+    if not isinstance(judge, dict) or judge.get("score") is None:
+        raise ValueError(f"Missing judge.score in xagent label line {line_number}")
+    score = float(judge["score"])
+    if not 0.0 <= score <= 1.0:
+        raise ValueError(f"judge.score must be in [0, 1] at line {line_number}: {score}")
+    return score
+
+
+def _fallback_prompt_id(record: dict[str, Any]) -> str:
+    event_id = record.get("event_id", "unknown")
+    end_event_id = record.get("end_event_id", "unknown")
+    category = record.get("category", "unknown")
+    return f"xagent:{event_id}:{end_event_id}:{category}"
+
+
+def _usage_cost(usage: object) -> float | None:
+    if not isinstance(usage, dict):
+        return None
+    value = usage.get("cost")
+    if value is None:
+        return None
+    return float(value)

src/xrouter_llm/resources/models/irt_router_350k.joblib

@@ -0,0 +1,64 @@
+import json
+
+from xrouter_llm.xagent_labels import load_xagent_openrouter_labels
+
+
+def test_load_xagent_labels_from_local_jsonl(tmp_path) -> None:
+    path = tmp_path / "xagent.public.jsonl"
+    path.write_text(
+        "\n".join(
+            json.dumps(record)
+            for record in [
+                {
+                    "prompt_sha256": "abc123",
+                    "prompt": "Current user request:\nhello",
+                    "candidate_model": "google/gemini-2.5-flash-lite",
+                    "candidate_usage": {"cost": 0.001},
+                    "category": "skill",
+                    "judge": {"score": 0.75, "can_complete": True},
+                },
+                {
+                    "prompt_sha256": "abc123",
+                    "prompt": "Current user request:\nhello",
+                    "candidate_model": "deepseek/deepseek-v4-flash",
+                    "candidate_usage": {"cost": 0.0005},
+                    "category": "skill",
+                    "judge": {"score": 0.25, "can_complete": False},
+                },
+            ]
+        ),
+        encoding="utf-8",
+    )
+
+    rows = load_xagent_openrouter_labels(path)
+
+    assert len(rows) == 2
+    assert {row.prompt_id for row in rows} == {"xagent:abc123"}
+    assert {row.model_id for row in rows} == {
+        "deepseek/deepseek-v4-flash",
+        "google/gemini-2.5-flash-lite",
+    }
+    assert [row.score for row in rows] == [0.75, 0.25]
+    assert rows[0].cost_usd == 0.001
+    assert rows[0].task == "xagent:skill"
+
+
+def test_load_xagent_labels_rejects_unlabeled_candidates(tmp_path) -> None:
+    path = tmp_path / "xagent.public.jsonl"
+    path.write_text(
+        json.dumps(
+            {
+                "prompt_sha256": "abc123",
+                "prompt": "prompt",
+                "category": "skill",
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    try:
+        load_xagent_openrouter_labels(path)
+    except ValueError as exc:
+        assert "Missing judge.score" in str(exc)
+    else:
+        raise AssertionError("expected missing judge.score to fail")