SPB Git forge

spb/llm-api

Public
0commits 0branches 0releases
0 Bsize
maindefault branch
—last push
3.0 KB · 71 lines python
Raw Blame History
1"""Optional `model: auto` routing. Never used when the client names a model explicitly."""23from __future__ import annotations45import re67CODE_RE = re.compile(r"```|\bdef \b|\bclass \b|\bimport \b|\bfunction\b|=>|#include|SELECT .* FROM", re.I)8REASON_RE = re.compile(r"\b(prove|proof|step by step|reason|derive|why|analy[sz]e|plan|strategy|compare|evaluate|explain in depth)\b", re.I)91011def _text_of(body: dict) -> tuple[str, bool]:12    text = ""13    has_image = False14    for m in body.get("messages") or []:15        c = m.get("content")16        if isinstance(c, str):17            text += c + "\n"18        elif isinstance(c, list):19            for p in c:20                if isinstance(p, dict):21                    if p.get("type") == "text":22                        text += p.get("text", "") + "\n"23                    elif p.get("type") in ("image_url", "input_image", "image"):24                        has_image = True25    if body.get("prompt"):26        text += str(body["prompt"])27    return text, has_image282930async def choose_auto(state, body: dict, endpoint: str) -> str:31    registry = state.manager.registry32    manager = state.manager33    aliases = await registry.aliases()34    text, has_image = _text_of(body)35    if endpoint.endswith("/embeddings"):36        return aliases.get("embedding") or await _first(registry, lambda m: m["embedding"])37    if endpoint.endswith("/rerank"):38        return aliases.get("reranker") or await _first(registry, lambda m: m["reranker"])39    if has_image:40        return aliases.get("vision") or await _first(registry, lambda m: m["vision"] and not m["embedding"])41    long_prompt = len(text) > 1200042    if CODE_RE.search(text) and "coder" in aliases:43        return aliases["coder"]44    if (REASON_RE.search(text) or long_prompt) and "reasoning" in aliases:45        return aliases["reasoning"]46    # prefer the currently loaded text model to avoid a switch47    cur = manager.current_model()48    if cur and not (cur.model.get("embedding") or cur.model.get("reranker") or cur.model.get("vision") and not text):49        return cur.model["id"]50    if "default" in aliases:51        return aliases["default"]52    if "fast" in aliases:53        return aliases["fast"]54    default = await state.db.get_setting("default_model")55    if default:56        return default57    return await _first(registry, lambda m: not (m["embedding"] or m["reranker"]) and m["compatible"])585960async def _first(registry, pred) -> str:61    models = await registry.list_models()62    fav = [m for m in models if pred(m) and m["favorite"] and m["enabled"] and m["compatible"]]63    ok = fav or [m for m in models if pred(m) and m["enabled"] and m["compatible"]]64    if not ok:65        from .errors import ModelNotFound66        raise ModelNotFound("No suitable model installed for automatic routing.")67    # smallest that is reasonably capable: prefer the largest under 20 GB, else smallest68    mid = sorted(ok, key=lambda m: (m["estimated_ram_gb"] or 0))69    under = [m for m in mid if (m["estimated_ram_gb"] or 0) <= 20]70    return (under[-1] if under else mid[0])["id"]71