whetkit 0.3.0__tar.gz → 0.4.0__tar.gz

{whetkit-0.3.0 → whetkit-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: whetkit
-Version: 0.3.0
+Version: 0.4.0
 Summary: Measure and improve how well LLM agents select and use the tools exposed by an MCP server
 Keywords: mcp,model-context-protocol,llm,agents,evals,tool-use,tool-selection,benchmark
 Author: Mohammed Benlamlih

{whetkit-0.3.0 → whetkit-0.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "whetkit"
-version = "0.3.0"
+version = "0.4.0"
 description = "Measure and improve how well LLM agents select and use the tools exposed by an MCP server"
 readme = "README.md"
 license = "Apache-2.0"

{whetkit-0.3.0 → whetkit-0.4.0}/src/whetkit/cli.py

@@ -255,6 +255,62 @@ def generate(
     typer.echo(f"  whetkit run --server {server} --tasks {out}")
 
 
+@app.command(name="plan-init")
+def plan_init(
+    server: Annotated[
+        str,
+        typer.Option("--server", help="MCP server: URL, directory, server.json, or server.py"),
+    ],
+    keep: Annotated[
+        str, typer.Option("--keep", help="Comma-separated tool names to keep visible")
+    ] = "",
+    from_tasks: Annotated[
+        str | None,
+        typer.Option(
+            "--from-tasks",
+            help="Also keep every tool referenced by these tasks' expected_tools",
+        ),
+    ] = None,
+    out: Annotated[
+        str, typer.Option("--out", help="Where to write the plan YAML")
+    ] = ".whetkit/curation-plan.yaml",
+    http_mode: Annotated[HttpMode, typer.Option("--http-mode")] = HttpMode.STATEFUL,
+) -> None:
+    """Scaffold a view plan: keep the named tools, hide everything else.
+    The fastest way to serve a lean read-only slice of a big server."""
+    from whetkit.curation import CurationPlan, ToolOverride, save_plan
+    from whetkit.datasets import load_tasks
+
+    keep_set = {name.strip() for name in keep.split(",") if name.strip()}
+    if from_tasks:
+        for task in load_tasks(from_tasks):
+            for slot in task.expected_tool_slots:
+                keep_set.update(slot)
+    if not keep_set:
+        raise typer.BadParameter("nothing to keep — pass --keep and/or --from-tasks")
+
+    spec = _resolve_server(server, http_mode)
+    inventory = asyncio.run(inspect_server(spec))
+    names = {t.name for t in inventory.tools}
+    if unknown := sorted(keep_set - names):
+        typer.echo(f"warning: not on the server, ignoring: {', '.join(unknown)}", err=True)
+
+    hidden = [
+        ToolOverride(original_name=t.name, hidden=True, reason="Not part of this view's workflows.")
+        for t in inventory.tools
+        if t.name not in keep_set
+    ]
+    plan = CurationPlan(
+        server=spec.label(),
+        notes=f"View plan: keep {len(names & keep_set)} tool(s), hide the rest.",
+        overrides=hidden,
+    )
+    Path(out).parent.mkdir(parents=True, exist_ok=True)
+    save_plan(plan, out)
+    typer.echo(f"kept {len(names & keep_set)}, hidden {len(hidden)} — wrote {out}")
+    typer.echo(f"score it:  whetkit run --server {server} --tasks <tasks> --plan {out}")
+
+
 @app.command()
 def export(
     plan: Annotated[
@@ -617,6 +673,17 @@ def curate(
             help="Completion-token budget per model turn (raise for reasoning models)",
         ),
     ] = 1024,
+    prune_unused: Annotated[
+        bool,
+        typer.Option(
+            "--prune-unused",
+            help=(
+                "Additionally hide every tool the eval never touched — the cost "
+                "play for big servers. Only sound when the tasks cover all the "
+                "workflows the curated view will serve."
+            ),
+        ),
+    ] = False,
     store: Annotated[str | None, typer.Option("--store")] = None,
     http_mode: Annotated[HttpMode, typer.Option("--http-mode")] = HttpMode.STATEFUL,
 ) -> None:
@@ -624,6 +691,7 @@ def curate(
     it, and show the before/after hit-rate."""
     from whetkit.curation import CuratedMCPClient, propose_plan, save_plan
     from whetkit.curation.optimizer import OptimizerConfig
+    from whetkit.curation.optimizer import prune_unused as apply_prune_unused
     from whetkit.datasets import load_tasks
     from whetkit.mcp import inspect_server
     from whetkit.runner import RunConfig, run_task
@@ -666,6 +734,9 @@ def curate(
             )
             for warning in warnings:
                 typer.echo(f"warning: {warning}", err=True)
+            if prune_unused:
+                pruned = apply_prune_unused(plan, inventory, task_list, baseline_runs)
+                typer.echo(f"--prune-unused: hid {pruned} untouched tool(s)", err=True)
             save_plan(plan, plan_path)
             typer.echo(f"curation plan written to {plan_path}", err=True)
 
@@ -742,6 +813,159 @@ def curate(
     asyncio.run(_curate())
 
 
+@app.command()
+def fix(
+    tasks: Annotated[
+        str, typer.Option("--tasks", help="Task YAML file or directory of task files")
+    ],
+    server: Annotated[
+        str | None, typer.Option("--server", help="Override the server every task runs against")
+    ] = None,
+    model: Annotated[
+        str, typer.Option("--model", help="Agent model as provider:model_id")
+    ] = "anthropic:claude-sonnet-5",
+    optimizer_model: Annotated[
+        str, typer.Option("--optimizer-model")
+    ] = "anthropic:claude-sonnet-5",
+    judge: Annotated[str, typer.Option("--judge", help="'auto', 'on', or 'off'")] = "auto",
+    judge_model: Annotated[str, typer.Option("--judge-model")] = "anthropic:claude-sonnet-5",
+    max_iterations: Annotated[
+        int, typer.Option("--max-iterations", help="Propose→eval→revise rounds (≥1)")
+    ] = 3,
+    plan_path: Annotated[
+        str, typer.Option("--plan", help="Where to write the best plan YAML")
+    ] = ".whetkit/curation-plan.yaml",
+    match_mode: Annotated[str, typer.Option("--match-mode")] = "order_tolerant",
+    max_turns: Annotated[int, typer.Option("--max-turns")] = 10,
+    max_tokens: Annotated[int, typer.Option("--max-tokens")] = 1024,
+    store: Annotated[str | None, typer.Option("--store")] = None,
+    http_mode: Annotated[HttpMode, typer.Option("--http-mode")] = HttpMode.STATEFUL,
+) -> None:
+    """Self-correcting curation: propose a plan, eval through it, feed the
+    regressions and remaining waste back to the optimizer, revise — up to
+    --max-iterations — and keep the best plan by measured results."""
+    from functools import partial
+
+    from whetkit.curation import CuratedMCPClient, propose_plan, save_plan
+    from whetkit.curation.optimizer import OptimizerConfig, propose_revision
+    from whetkit.datasets import load_tasks
+    from whetkit.runner import RunConfig, run_task
+    from whetkit.scoring import JudgeCache, JudgeConfig, MatchMode, score_runs
+    from whetkit.tracing import TraceStore, default_store_path
+
+    if max_iterations < 1:
+        raise typer.BadParameter("--max-iterations must be at least 1")
+    task_list = load_tasks(tasks)
+    servers = _resolve_task_servers(task_list, server, http_mode)
+    config = RunConfig(model=model, max_turns=max_turns, max_tokens=max_tokens)
+    use_judge = _judge_enabled(judge, judge_model)
+    judge_config = JudgeConfig(model=judge_model)
+    mode = MatchMode(match_mode)
+    store_path = store or str(default_store_path())
+    optimizer_config = OptimizerConfig(model=optimizer_model)
+
+    from whetkit.mcp import MCPClient
+
+    async def _eval(group: str, cache, client_factory=MCPClient):
+        runs = []
+        for task in task_list:
+            typer.echo(f"running {task.id} ...", err=True)
+            runs.append(
+                await run_task(task, servers[task.server], config, client_factory=client_factory)
+            )
+        with TraceStore(store_path) as trace_store:
+            trace_store.save_runs(runs, run_group=group)
+        return runs
+
+    async def _fix() -> None:
+        from whetkit.mcp import inspect_server as _inspect
+
+        cache = JudgeCache(default_store_path().parent / "judge_cache.sqlite3")
+        try:
+            typer.echo("== baseline ==", err=True)
+            baseline_runs = await _eval("baseline", cache)
+            baseline = await score_runs(
+                task_list,
+                baseline_runs,
+                mode=mode,
+                judge_config=judge_config,
+                judge_cache=cache,
+                use_judge=use_judge,
+            )
+            inventory = await _inspect(next(iter(servers.values())))
+
+            typer.echo("== proposing plan (iteration 1) ==", err=True)
+            plan, warnings = await propose_plan(
+                inventory, task_list, baseline_runs, baseline.scores, optimizer_config
+            )
+            for w in warnings:
+                typer.echo(f"warning: {w}", err=True)
+
+            def metric(summary):
+                return (summary.hit_rate, -summary.avg_extra_calls, summary.avg_precision)
+
+            best_plan, best_summary = None, None
+            for iteration in range(1, max_iterations + 1):
+                typer.echo(f"== eval through plan (iteration {iteration}) ==", err=True)
+                curated_runs = await _eval(
+                    f"fix-{iteration}",
+                    cache,
+                    client_factory=partial(CuratedMCPClient, plan=plan),
+                )
+                curated = await score_runs(
+                    task_list,
+                    curated_runs,
+                    mode=mode,
+                    judge_config=judge_config,
+                    judge_cache=cache,
+                    use_judge=use_judge,
+                    name_map=plan.rename_map(),
+                )
+                typer.echo(
+                    f"iteration {iteration}: hit {curated.hit_rate:.0%} "
+                    f"(baseline {baseline.hit_rate:.0%}), "
+                    f"extra calls {curated.avg_extra_calls:.1f}/task",
+                    err=True,
+                )
+                if best_summary is None or metric(curated) > metric(best_summary):
+                    best_plan, best_summary = plan, curated
+
+                converged = (
+                    curated.hit_rate >= baseline.hit_rate and curated.avg_extra_calls <= 0.25
+                )
+                if converged or iteration == max_iterations:
+                    if converged:
+                        typer.echo("converged — no regressions, negligible waste", err=True)
+                    break
+
+                typer.echo("== revising plan ==", err=True)
+                plan, warnings = await propose_revision(
+                    plan,
+                    inventory,
+                    task_list,
+                    baseline_runs,
+                    baseline.scores,
+                    curated_runs,
+                    curated.scores,
+                    optimizer_config,
+                )
+                for w in warnings:
+                    typer.echo(f"warning: {w}", err=True)
+        finally:
+            cache.close()
+
+        save_plan(best_plan, plan_path)
+        typer.echo(f"\nbest plan (of {max_iterations} max iterations) -> {plan_path}")
+        typer.echo(
+            f"Hit-rate: {baseline.hit_rate:.0%} -> {best_summary.hit_rate:.0%}   "
+            f"Extra calls: {baseline.avg_extra_calls:.1f} -> "
+            f"{best_summary.avg_extra_calls:.1f}/task"
+        )
+        typer.echo(f"serve it:  whetkit overlay --server <origin> --plan {plan_path}")
+
+    asyncio.run(_fix())
+
+
 @app.command()
 def report(
     tasks: Annotated[

{whetkit-0.3.0 → whetkit-0.4.0}/src/whetkit/curation/optimizer.py

@@ -197,3 +197,132 @@ async def propose_plan(
             plan.overrides = []
             break
     return plan, warnings
+
+
+def prune_unused(
+    plan: CurationPlan,
+    inventory: ServerInventory,
+    tasks: list[TaskSpec],
+    runs: list[TaskRun],
+) -> int:
+    """Hide every tool the eval never touched: not called in any baseline
+    run, not accepted by any task's expected_tools, and not already
+    overridden. The cost play for big servers — only sound when the task
+    set really covers the workflows the curated view will serve, which is
+    why it's opt-in. Returns how many tools were hidden."""
+    used = {name for run in runs for name in run.called_tool_names}
+    for task in tasks:
+        for slot in task.expected_tool_slots:
+            used.update(slot)
+    overridden = {o.original_name for o in plan.overrides}
+    added = 0
+    for tool in inventory.tools:
+        if tool.name in used or tool.name in overridden:
+            continue
+        plan.overrides.append(
+            ToolOverride(
+                original_name=tool.name,
+                hidden=True,
+                reason="Unused by the eval workflows (pruned by --prune-unused).",
+            )
+        )
+        added += 1
+    return added
+
+
+REVISION_INSTRUCTIONS = """\
+You previously proposed the curation plan below, and it was re-evaluated.
+Revise it: KEEP the changes that worked, FIX or DROP the ones implicated in
+regressions or remaining waste. A rename the agent still misused needs a
+better name or a sharper description; a prune that broke a task must be
+un-hidden. Respond with the SAME JSON shape as before — the FULL revised
+plan, not a delta.
+"""
+
+
+def build_revision_prompt(
+    plan: CurationPlan,
+    inventory: ServerInventory,
+    tasks: list[TaskSpec],
+    baseline_runs: list[TaskRun],
+    baseline_scores: list[TaskScore],
+    curated_runs: list[TaskRun],
+    curated_scores: list[TaskScore],
+) -> str:
+    import yaml
+
+    baseline_by_id = {s.task_id: s for s in baseline_scores}
+    regressed = [
+        s.task_id
+        for s in curated_scores
+        if not s.hit and (b := baseline_by_id.get(s.task_id)) is not None and b.hit
+    ]
+    regression_line = (
+        f"REGRESSIONS your plan caused (fix these first): {', '.join(regressed)}"
+        if regressed
+        else "No regressions — focus on remaining waste (unnecessary calls)."
+    )
+    plan_yaml = yaml.safe_dump(plan.model_dump(exclude_defaults=True), sort_keys=False)
+    return (
+        f"## Tools exposed by the server\n{_inventory_block(inventory)}\n\n"
+        f"## Baseline eval (no plan)\n{_trace_block(tasks, baseline_runs, baseline_scores)}\n\n"
+        f"## Your previous plan\n```yaml\n{plan_yaml}```\n\n"
+        f"## Eval THROUGH your plan\n{_trace_block(tasks, curated_runs, curated_scores)}\n\n"
+        f"{regression_line}\n\nRevise the plan now."
+    )
+
+
+async def propose_revision(
+    plan: CurationPlan,
+    inventory: ServerInventory,
+    tasks: list[TaskSpec],
+    baseline_runs: list[TaskRun],
+    baseline_scores: list[TaskScore],
+    curated_runs: list[TaskRun],
+    curated_scores: list[TaskScore],
+    config: OptimizerConfig | None = None,
+    provider: LLMProvider | None = None,
+) -> tuple[CurationPlan, list[str]]:
+    """One revision round: same parsing/validation as propose_plan, but the
+    model sees its previous plan and what happened through it."""
+    config = config or OptimizerConfig()
+    provider_name, model_id = parse_model(config.model)
+    provider = provider or get_provider(provider_name)
+
+    prompt = build_revision_prompt(
+        plan, inventory, tasks, baseline_runs, baseline_scores, curated_runs, curated_scores
+    )
+    parsed = None
+    for _attempt in range(2):
+        turn = await provider.complete(
+            model=model_id,
+            system=OPTIMIZER_SYSTEM_PROMPT + "\n" + REVISION_INSTRUCTIONS,
+            messages=[ChatMessage(role="user", content=prompt)],
+            tools=[],
+            max_tokens=config.max_tokens,
+        )
+        parsed = _parse_overrides(turn.text or "")
+        if parsed is not None:
+            break
+    if parsed is None:
+        return plan, ["revision output was not valid JSON after 2 attempts; keeping previous plan"]
+
+    notes, overrides = parsed
+    revised = CurationPlan(server=inventory.server, notes=notes, overrides=overrides)
+    origin_names = {t.name for t in inventory.tools}
+    warnings: list[str] = []
+    while problems := revised.validate_against(origin_names):
+        problem = problems[0]
+        name_match = re.search(r"'([^']+)'", problem)
+        offender = name_match.group(1) if name_match else None
+        before = len(revised.overrides)
+        revised.overrides = [
+            o
+            for o in revised.overrides
+            if offender not in (o.original_name, o.new_name, o.presented_name)
+        ]
+        warnings.append(f"dropped unsafe override(s): {problem}")
+        if len(revised.overrides) == before:
+            warnings.append("could not repair revised plan; keeping previous plan")
+            return plan, warnings
+    return revised, warnings