Source code for best_engine_ai_helper.recommend

"""
recommend — the end-to-end "which engine" algorithm.

Takes three inputs, at least one of them allowed to be vague:

1. a **hardware description** (memory pool + compute profile),
2. the **benchmark catalog** (models with public scores), and
3. a **task**, which may be a free-text phrase ("retail product descriptions
   and image-quality checks") rather than a fixed keyword,

and returns the best local engine to pull for each kind the task needs
(LLM for text, VLM for images), justified by three factors it makes explicit:

- **task fit** — the model's benchmark on the axis the task maps to,
- **memory fit** — whether it fits the accelerator's usable budget (so it runs
  on the GPU, not spilled to CPU), and
- **throughput** — an estimate of decode tokens/s from memory bandwidth.

The result is a plain dict (ready for ``json.dumps``) and, via
:func:`to_markdown`, a human-readable report. Nothing here calls a model or the
network; it is pure ranking over the catalog and the detected hardware.

Author
------
Warith Harchaoui <warith.harchaoui@deraison.ai>
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import os_helper as osh

from .score import effective_budget, estimated_tokens_per_second, rank

# Keyword -> benchmark axis. The task text is scanned for these; the first axis
# whose keywords appear wins for the text model. Vision keywords additionally
# turn on the VLM recommendation. Deliberately small and legible: this is a
# heuristic to turn a vague phrase into a scoring axis, not an NLP model.
_AXIS_KEYWORDS: dict[str, tuple[str, ...]] = {
    "ocr": ("ocr", "scan", "receipt", "invoice", "document", "handwrit"),
    "code": ("code", "coding", "program", "python", "sql", "javascript", "bug"),
    "math": ("math", "calculation", "arithmetic", "equation", "proof"),
    "vision": ("image", "vision", "visual", "photo", "picture", "diagram",
               "screenshot", "chart", "quality assessment", "aesthetic"),
}
_VISION_AXES = {"vision", "ocr"}


[docs] def parse_task(task: str | None) -> dict[str, Any]: """ Turn a vague task phrase into the model kinds and benchmark axis it implies. Returns a dict with ``kinds`` (subset of ``["llm", "vlm"]`` in pull order), ``application`` (the benchmark axis for the text model), and ``matched`` (the keywords that fired, for the report's justification). A task that mentions nothing visual still gets an LLM on the ``generalist`` axis; any vision keyword adds a VLM. """ text = (task or "").lower() matched: list[str] = [] axis = "generalist" for candidate_axis, words in _AXIS_KEYWORDS.items(): hit = [w for w in words if w in text] if hit: matched.extend(hit) # A non-vision axis (code/math) sets the text model's axis; vision # keywords are recorded but the text axis stays generalist unless a # text-specific axis also matched. if candidate_axis not in _VISION_AXES: axis = candidate_axis needs_vlm = any( w in text for a in _VISION_AXES for w in _AXIS_KEYWORDS[a] ) kinds = ["llm"] if needs_vlm: kinds.append("vlm") # If the task is purely an OCR axis, the VLM should be scored on ocr. vlm_axis = "ocr" if any(w in text for w in _AXIS_KEYWORDS["ocr"]) else "vision" return {"kinds": kinds, "application": axis, "vlm_application": vlm_axis, "matched": sorted(set(matched))}
def _candidate_row( entry: dict[str, Any], budget: float, bandwidth_gbs: float | None, axis: str ) -> dict[str, Any]: """One catalog entry annotated with the three decision factors.""" from .score import _benchmark_score # local: internal scorer ram = float(entry.get("ram_gb", 0) or 0) kind = entry.get("kind", "llm") return { "id": entry.get("id"), "kind": kind, "ram_gb": ram, "score": _benchmark_score(entry, kind, axis), "fits": ram <= budget, # False = can't honour Ollama structured JSON output, so it is never # auto-chosen for the helper's schema-driven tasks (see score.rank). "structured_output": entry.get("structured_output", True) is not False, "est_tokens_per_s": estimated_tokens_per_second(entry, bandwidth_gbs), "notes": entry.get("notes", ""), }
[docs] def recommend( hw: dict[str, float | None], catalog: list[dict[str, Any]], task: str | None = None, *, headroom: float = 0.85, compute: dict[str, Any] | None = None, ) -> dict[str, Any]: """ Recommend the best engine per needed kind for this hardware and task. Parameters ---------- hw : dict Memory description from :func:`detect.available_memory`. catalog : list of dict Benchmark catalog from :func:`catalog.load_catalog`. task : str or None Free-text or keyword task. None means a generalist text assistant. headroom : float Memory safety margin passed to :func:`score.effective_budget`. compute : dict or None Compute profile from :func:`detect.compute_profile` (accelerator + bandwidth). None disables the throughput estimate. Returns ------- dict A JSON-ready report: the parsed task, hardware, memory budget, and for each needed kind the chosen model plus the full ranked candidate table, with per-model fit and estimated throughput. """ compute = compute or {} bandwidth = compute.get("bandwidth_gbs") budget = effective_budget(hw, headroom=headroom) parsed = parse_task(task) osh.info( f"Recommending for task axis '{parsed['application']}' " f"(kinds: {', '.join(parsed['kinds'])}); memory budget {budget:.1f} GB" ) picks: dict[str, Any] = {} for kind in parsed["kinds"]: axis = parsed["vlm_application"] if kind == "vlm" else parsed["application"] ranked = rank(hw, catalog, kind, headroom=headroom, application=axis) rows = [_candidate_row(e, budget, bandwidth, axis) for e in ranked] fitting = [r for r in rows if r["fits"]] chosen = fitting[0] if fitting else (rows[0] if rows else None) if chosen is None: osh.warning(f"No candidate found for kind '{kind}' on axis '{axis}'") else: osh.info( f"Chosen {kind} on axis '{axis}': {chosen['id']} " f"({chosen['ram_gb']:.1f} GB, {'fits' if chosen['fits'] else 'over budget'})" ) if not chosen.get("structured_output", True): # Only happens when no structured-capable model fits the budget # (e.g. a small machine, where the only VLMs that fit are the # Qwen3-VL family). Say so plainly so the caller can warn the user # rather than silently getting empty structured responses. osh.warning( f"Chosen {kind} '{chosen['id']}' does NOT support Ollama structured " "JSON output; schema-driven tasks (intent, critique) will fail on it. " "Pull a structured-capable model or free up memory for one." ) # Lightest model within 3 benchmark points of the chosen one: the # "good enough but leaner / faster" alternative worth surfacing. It must # not be less structured-output-capable than the chosen model, otherwise # the "leaner" pick would silently fail the suite's schema-driven tasks # (e.g. suggesting a Qwen3-VL under a structured-capable chosen). lighter = None if chosen: near = [ r for r in fitting if r["score"] >= chosen["score"] - 3 and (r["structured_output"] or not chosen["structured_output"]) ] if near: lighter = min(near, key=lambda r: r["ram_gb"]) if lighter["id"] == chosen["id"]: lighter = None picks[kind] = { "axis": axis, "chosen": chosen, "lighter_alternative": lighter, "candidates": rows, } return { "task": {"input": task, **parsed}, "hardware": {**hw, **({"compute": compute} if compute else {})}, "memory_budget_gb": budget, "headroom": headroom, "recommendations": picks, "method": _METHOD_NOTE, "sources": _SOURCES, }
_METHOD_NOTE = ( "Each model is judged on four explicit factors: structured-output capability " "(whether it honours Ollama JSON-schema output -- the helper routes every " "task through a schema, so a model that can't is never auto-chosen however " "high it scores), task fit (its benchmark on the task's axis), memory fit " "(whether its peak-inference RAM fits the accelerator's usable budget, so it " "runs on the GPU rather than spilling to CPU), and throughput (estimated " "decode tokens/s = memory bandwidth / model size x 0.65 efficiency, because " "generation is memory-bandwidth bound). The chosen model is the highest " "task-fit, structured-capable model that fits the memory budget; a lighter " "alternative is offered when one is nearly as strong but smaller and faster." ) _SOURCES = [ "Apple Metal recommendedMaxWorkingSetSize ~66% (<=36 GB) / ~75% (>36 GB) of " "unified memory (apple-specs, Metal docs).", "Inference RAM = weights + KV cache (15-20%) + overhead (5-10%), up to ~2x " "weights at long context (local-LLM sizing guides).", "Decode is memory-bandwidth bound: tokens/s ~= bandwidth / active-model-bytes, " "derated to ~50-80% real-world (llama.cpp / MLX benchmarks).", ] def _fmt_tps(v: float | None) -> str: return f"{v:.0f}" if v else "-"
[docs] def to_markdown(report: dict[str, Any]) -> str: """Render a :func:`recommend` report as a Markdown document.""" hw = report["hardware"] comp = hw.get("compute", {}) or {} task = report["task"] lines: list[str] = [] lines.append("# Best local engine — recommendation") lines.append("") lines.append(f"**Task:** {task.get('input') or 'general text assistant'} ") if task.get("matched"): lines.append(f"**Matched keywords:** {', '.join(task['matched'])} ") lines.append(f"**Needs:** {', '.join(k.upper() for k in task['kinds'])}") lines.append("") lines.append("## Hardware") lines.append("") lines.append(f"- Chip / accelerator: {comp.get('chip') or '?'} " f"({comp.get('accelerator', 'unknown')})") pool = hw.get("unified_gb") or hw.get("vram_gb") or hw.get("ram_gb") lines.append(f"- Memory pool: {pool} GB") if comp.get("bandwidth_gbs"): lines.append(f"- Memory bandwidth: {comp['bandwidth_gbs']:.0f} GB/s " "(sets the decode-speed ceiling)") lines.append(f"- Usable model budget: **{report['memory_budget_gb']:.1f} GB** " f"(headroom {report['headroom']})") lines.append("") for kind, block in report["recommendations"].items(): chosen = block["chosen"] lines.append(f"## Best {kind.upper()} (axis: {block['axis']})") lines.append("") if chosen: tps = _fmt_tps(chosen["est_tokens_per_s"]) lines.append(f"**→ `{chosen['id']}`** — {chosen['ram_gb']:.1f} GB, " f"score {chosen['score']:.0f}, ~{tps} tok/s" + ("" if chosen["fits"] else " ⚠️ exceeds budget (will be slow)")) alt = block["lighter_alternative"] if alt: lines.append("") lines.append(f"Lighter alternative: `{alt['id']}` — " f"{alt['ram_gb']:.1f} GB, score {alt['score']:.0f}, " f"~{_fmt_tps(alt['est_tokens_per_s'])} tok/s.") else: lines.append("_No candidate found._") lines.append("") lines.append("| model | RAM GB | score | fits | ~tok/s |") lines.append("|---|---|---|---|---|") for r in block["candidates"]: lines.append(f"| `{r['id']}` | {r['ram_gb']:.1f} | {r['score']:.0f} | " f"{'yes' if r['fits'] else 'no'} | {_fmt_tps(r['est_tokens_per_s'])} |") lines.append("") lines.append("## How this was decided") lines.append("") lines.append(report["method"]) lines.append("") lines.append("Sources:") for s in report["sources"]: lines.append(f"- {s}") lines.append("") return "\n".join(lines)
[docs] def write_report(report: dict[str, Any], stem: str | Path) -> tuple[Path, Path]: """Write both a JSON and a Markdown rendering; return (md_path, json_path).""" stem = Path(stem) md_path = stem.with_suffix(".md") json_path = stem.with_suffix(".json") osh.make_directory(str(md_path.parent)) md_path.write_text(to_markdown(report), encoding="utf-8") json_path.write_text(json.dumps(report, indent=2), encoding="utf-8") osh.info(f"Wrote recommendation report:\n\t{md_path}\n\t{json_path}") return md_path, json_path