stabbur 0.6.0__py3-none-any.whl

stabbur/__init__.py

@@ -0,0 +1,37 @@
+"""Backup and browse local LLM models from Hugging Face, Ollama, and LM Studio.
+
+This package's ``__init__`` deliberately runs two Hugging Face environment tweaks **at import
+time**, before ``huggingface_hub`` is first imported. That timing is load-bearing, not
+incidental: ``huggingface_hub`` freezes its cache path (from ``HF_HOME`` / ``HF_HUB_CACHE``) at
+*its* import, and importing almost any stabbur module (e.g. ``stabbur.cli``) transitively imports
+``huggingface_hub`` — so this ``__init__`` is the only hook that reliably runs first. (A8 in the
+review proposed moving these into the CLI/app entry points; that was verified *unsafe* precisely
+because the entry points run after hf_hub is already imported.) These are the sole intentional
+import-time side effects; everything else is lazy.
+"""
+
+import os
+
+# Ask huggingface_hub for high-performance Xet transfer (the fast, parallel path
+# in hub >=1.21, backed by the hf_xet dependency) — big model pulls otherwise fall
+# back to slow single-stream HTTP. Set before huggingface_hub is imported anywhere.
+# (The old HF_HUB_ENABLE_HF_TRANSFER flag is deprecated/removed in 1.21.)
+os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")
+
+# Point the HF hub cache at the library drive (before huggingface_hub is imported), so
+# assets some runtimes fetch by repo id — e.g. an mlx-audio model's codec — travel with
+# the drive instead of living in ~/.cache/huggingface. Best-effort + guarded: a no-op if
+# the user set HF_HOME/HF_HUB_CACHE, or there's no real configured library (unset, or an
+# unmounted drive → falls back to the machine cache). See stabbur.hfcache.
+from stabbur import hfcache as _hfcache  # noqa: E402
+
+_hfcache.configure()
+
+# Read the version from package metadata (chapkit's pattern) rather than restating it here:
+# a second copy only drifts — this one still said 0.1.0 at the 0.4.0 release.
+try:
+    from importlib.metadata import version as _get_version
+
+    __version__ = _get_version("stabbur")
+except Exception:  # noqa: BLE001 - a source tree that was never installed has no metadata
+    __version__ = "unknown"

stabbur/agent.py

@@ -0,0 +1,332 @@
+"""The tool-calling agent loop: model ⇄ tools over the OpenAI /v1 API.
+
+The model may emit ``tool_call``s; stabbur executes each via the MCP toolset, feeds
+the results back, and repeats until the model answers with plain text.
+"""
+
+import json
+from collections.abc import Awaitable, Callable
+from inspect import isawaitable
+from typing import Any, Literal
+
+import httpx
+
+from stabbur.tools import MCPToolset, ToolResult
+
+# Prefaces the user message that carries images a tool returned (e.g. a screenshot), so a
+# vision model reads them as its own multimodal input right after the tool results.
+_TOOL_IMAGE_PREAMBLE = "Image(s) returned by the tool call(s) above:"
+
+# Callbacks: on_event(kind, detail) for tool activity; on_token(text) for streamed reply.
+# Sinks may be sync (TUI/CLI append to a buffer) or async (the /api/chat SSE path uses an
+# async sink so a full, bounded queue back-pressures generation instead of buffering the
+# whole reply in memory). ``_emit`` awaits the result only when it's awaitable.
+ToolEvent = Callable[[str, str], None | Awaitable[None]]
+TokenSink = Callable[[str], None | Awaitable[None]]
+# on_usage(usage) receives the server's token accounting for a turn (OpenAI `usage`:
+# prompt_tokens / completion_tokens / total_tokens) so a REPL can show context used.
+UsageSink = Callable[[dict[str, Any]], None]
+# on_confirm(name, args) is consulted before a gated tool call runs; it returns whether the user
+# approved the action. Async-only (the /api/chat + UI path awaits a user decision over a channel),
+# so the loop can suspend on it — a missing sink is treated as a denial (fail-safe) at the gate.
+ConfirmSink = Callable[[str, dict[str, Any]], Awaitable[bool]]
+
+
+def _needs_confirm(policy: Literal["all", "writes", "none"], toolset: MCPToolset, name: str) -> bool:
+    """Whether a tool call must be confirmed under ``policy``.
+
+    ``"all"`` gates every call; ``"writes"`` gates only tools not known read-only (fail-safe —
+    an unknown/unannotated tool is treated as a write); anything else (``"none"``) gates nothing.
+    """
+    if policy == "all":
+        return True
+    if policy == "writes":
+        return not toolset.is_readonly(name)
+    return False
+
+
+async def _emit(sink: Callable[..., None | Awaitable[None]] | None, *args: Any) -> None:
+    """Call an optional sync-or-async sink, awaiting it when it returns an awaitable."""
+    if sink is None:
+        return
+    result = sink(*args)
+    if isawaitable(result):
+        await result
+
+
+def _audio_part(data_url: str) -> dict[str, Any]:
+    """Convert an audio ``data:`` URL to an OpenAI ``input_audio`` content part.
+
+    llama-server / mlx-vlm want ``{data: <base64 (no prefix)>, format: "wav"|"mp3"}``
+    (not a data URL like images), so split the mime + payload out.
+    """
+    fmt = "wav"
+    b64 = data_url
+    if data_url.startswith("data:"):
+        header, _, b64 = data_url.partition(",")
+        mime = header[len("data:") :].split(";")[0]  # e.g. audio/wav
+        subtype = mime.split("/")[-1] or "wav"
+        fmt = {"mpeg": "mp3", "x-wav": "wav", "wave": "wav"}.get(subtype, subtype)
+    return {"type": "input_audio", "input_audio": {"data": b64, "format": fmt}}
+
+
+def user_content(
+    text: str, images: list[str] | None = None, audios: list[str] | None = None
+) -> str | list[dict[str, Any]]:
+    """Build a user message's content: plain text, or OpenAI multimodal parts.
+
+    ``images`` / ``audios`` are ``data:`` URL strings. With none, returns the plain
+    string (backward compatible); otherwise a ``content`` array of an optional text
+    part followed by ``image_url`` and ``input_audio`` parts — the format both
+    llama-server (with ``--mmproj``) and mlx-vlm accept.
+    """
+    if not images and not audios:
+        return text
+    parts: list[dict[str, Any]] = []
+    if text:
+        parts.append({"type": "text", "text": text})
+    parts += [{"type": "image_url", "image_url": {"url": u}} for u in images or []]
+    parts += [_audio_part(a) for a in audios or []]
+    return parts
+
+
+async def _stream_turn(
+    http: httpx.AsyncClient,
+    base_url: str,
+    body: dict[str, Any],
+    on_token: TokenSink | None,
+    on_reasoning: TokenSink | None = None,
+) -> tuple[str, list[dict[str, Any]], dict[str, Any] | None]:
+    """Stream one completion; return (content, tool_calls, usage). Emits content + reasoning live."""
+    content = ""
+    calls: dict[int, dict[str, str]] = {}
+    usage: dict[str, Any] | None = None
+    async with http.stream("POST", f"{base_url}/v1/chat/completions", json=body) as resp:
+        if resp.status_code >= 400:
+            # Read the body on error — llama-server puts the real cause (e.g. context overflow)
+            # in the JSON detail, which raise_for_status alone would discard.
+            detail = (await resp.aread()).decode("utf-8", errors="replace").strip()[:500]
+            raise httpx.HTTPStatusError(
+                f"runtime returned {resp.status_code}: {detail}" if detail else f"runtime returned {resp.status_code}",
+                request=resp.request,
+                response=resp,
+            )
+        async for line in resp.aiter_lines():
+            if not line.startswith("data:"):
+                continue
+            payload = line[len("data:") :].strip()
+            if payload == "[DONE]":
+                break
+            chunk = json.loads(payload)
+            # With include_usage the final chunk carries `usage` and an empty
+            # `choices` list; capture it and skip the (missing) delta.
+            if chunk.get("usage"):
+                captured = chunk["usage"]
+                usage = dict(captured) if isinstance(captured, dict) else {}
+                # llama.cpp adds its own `timings` (prompt_ms, predicted_ms,
+                # predicted_per_second) to that chunk. Pass them through: the runtime's
+                # measurement of its own decode rate beats anything a client can infer
+                # from arrival times. Absent on other servers, hence the guard.
+                timings = chunk.get("timings")
+                if isinstance(timings, dict):
+                    usage["timings"] = timings
+            if not chunk.get("choices"):
+                continue
+            delta = chunk["choices"][0]["delta"]
+            if delta.get("content"):
+                content += delta["content"]
+                await _emit(on_token, delta["content"])
+            # Reasoning models (gemma-4, Qwen3.5, …) stream their thinking here, not in
+            # content; surface it separately instead of dropping it (→ blank replies).
+            if delta.get("reasoning_content") and on_reasoning:
+                await _emit(on_reasoning, delta["reasoning_content"])
+            for tc in delta.get("tool_calls") or []:
+                slot = calls.setdefault(tc["index"], {"id": "", "name": "", "args": ""})
+                if tc.get("id"):
+                    slot["id"] = tc["id"]
+                fn = tc.get("function") or {}
+                if fn.get("name"):
+                    slot["name"] = fn["name"]
+                if fn.get("arguments"):
+                    slot["args"] += fn["arguments"]
+    ordered = [calls[i] for i in sorted(calls)]
+    return content, ordered, usage
+
+
+ReasoningLevel = Literal["off", "low", "medium", "high", "max"]
+"""Reasoning-effort levels for thinking models, mirroring the llama.cpp webui's control."""
+
+_REASONING_BUDGETS: dict[str, int] = {"low": 512, "medium": 2048, "high": 8192}
+
+
+def reasoning_fields(level: "ReasoningLevel | None") -> dict[str, Any]:
+    """The request fields for a reasoning level; empty for ``None`` (the model default).
+
+    Speaks llama-server's dialect — exactly what its own webui sends: ``chat_template_kwargs.
+    enable_thinking`` toggles thinking in the chat template (Qwen-style models), ``thinking_
+    budget_tokens`` caps the thinking length (low 512 / medium 2048 / high 8192; ``max`` sends
+    no cap), and ``reasoning_control`` marks the request as reasoning-managed. Servers without
+    reasoning support ignore the unknown fields, so sending them is always safe.
+    """
+    if level is None:
+        return {}
+    fields: dict[str, Any] = {
+        "chat_template_kwargs": {"enable_thinking": level != "off"},
+        "reasoning_control": True,
+    }
+    if level in _REASONING_BUDGETS:
+        fields["thinking_budget_tokens"] = _REASONING_BUDGETS[level]
+    return fields
+
+
+async def run(
+    base_url: str,
+    messages: list[dict[str, Any]],
+    toolset: MCPToolset,
+    max_tokens: int | None = None,
+    on_event: ToolEvent | None = None,
+    on_token: TokenSink | None = None,
+    on_reasoning: TokenSink | None = None,
+    temperature: float | None = None,
+    top_p: float | None = None,
+    top_k: int | None = None,
+    min_p: float | None = None,
+    repeat_penalty: float | None = None,
+    model: str | None = None,
+    max_rounds: int = 8,
+    on_usage: UsageSink | None = None,
+    tool_timeout: float | None = None,
+    vision: bool = False,
+    on_confirm: ConfirmSink | None = None,
+    confirm_policy: Literal["all", "writes", "none"] = "none",
+    reasoning: "ReasoningLevel | None" = None,
+) -> str:
+    """Run the agent loop against ``base_url``, streaming the reply; return its text.
+
+    ``messages`` is mutated in place (assistant/tool turns) so a REPL keeps the
+    conversation. ``on_event`` reports tool activity; ``on_token`` receives the
+    final reply's tokens; ``on_reasoning`` receives a reasoning model's thinking
+    tokens (separate channel); ``on_usage`` receives the server's token accounting
+    (prompt/completion/total) after each round. ``model`` is sent as the OpenAI
+    ``model`` field — required by mlx-vlm (which 422s without it), ignored by
+    llama-server/mlx-lm. Bounded by ``max_rounds``; each tool call is bounded by
+    ``tool_timeout`` seconds so a hung MCP server can't stall the loop forever —
+    ``None`` (default) reads ``STABBUR_TOOL_TIMEOUT`` (120s; set 0 to disable the bound).
+    ``vision`` is set when the model can see images: an image a tool returns (e.g. a
+    screenshot) is then fed back as a follow-up user image message so the model reads it
+    as multimodal input; a text-only model instead gets a note that an image was returned.
+    ``confirm_policy`` gates tool execution behind ``on_confirm``: ``"none"`` (default) runs
+    every tool as before; ``"writes"`` confirms only tools not known read-only (via each tool's
+    ``readOnlyHint`` annotation — an unannotated tool is treated as a write); ``"all"`` confirms
+    every call. When a call is gated, ``on_confirm(name, args)`` is awaited for approval; if it
+    returns falsy (or no ``on_confirm`` is supplied — fail-safe deny), the tool is NOT run and the
+    model gets a ``tool`` turn whose text is exactly ``error: user declined this action``.
+    """
+    if tool_timeout is None:
+        from stabbur.config import get_settings  # noqa: PLC0415 - lazy to keep agent import light
+
+        tool_timeout = get_settings().tool_timeout or None  # 0 → no bound
+    async with httpx.AsyncClient(timeout=600) as http:
+        for _ in range(max_rounds):
+            body: dict[str, Any] = {"messages": messages, "stream": True}
+            # Ask for a final usage chunk (prompt/completion tokens) so callers can
+            # report real context consumption; runtimes that ignore it just omit it.
+            if on_usage is not None:
+                body["stream_options"] = {"include_usage": True}
+            if model is not None:
+                body["model"] = model
+            # Omit tools entirely when there are none, so a no-tool chat is plain
+            # completion (no --jinja tool parsing / buffering).
+            if toolset.schemas:
+                body["tools"] = toolset.schemas
+                body["tool_choice"] = "auto"
+            if max_tokens is not None:
+                body["max_tokens"] = max_tokens
+            if temperature is not None:
+                body["temperature"] = temperature
+            if top_p is not None:
+                body["top_p"] = top_p
+            # top_k / min_p / repeat_penalty are OpenAI extensions supported by
+            # llama-server and the MLX servers; unknown ones are ignored upstream.
+            if top_k is not None:
+                body["top_k"] = top_k
+            if min_p is not None:
+                body["min_p"] = min_p
+            if repeat_penalty is not None:
+                body["repeat_penalty"] = repeat_penalty
+            # Reasoning effort (thinking on/off + budget) — llama-server dialect, see reasoning_fields.
+            body.update(reasoning_fields(reasoning))
+            content, calls, usage = await _stream_turn(http, base_url, body, on_token, on_reasoning)
+            if usage and on_usage:
+                on_usage(usage)
+            if not calls:
+                messages.append({"role": "assistant", "content": content})
+                return content
+
+            messages.append(
+                {
+                    "role": "assistant",
+                    "content": content or None,
+                    "tool_calls": [
+                        {"id": c["id"], "type": "function", "function": {"name": c["name"], "arguments": c["args"]}}
+                        for c in calls
+                    ],
+                }
+            )
+            round_images: list[str] = []  # images tools returned this round (fed back below)
+            for c in calls:
+                await _emit(on_event, "call", f"{c['name']}({c['args']})")
+                try:
+                    args = json.loads(c["args"] or "{}")
+                except json.JSONDecodeError as exc:
+                    # Don't run the tool with empty args on unparseable JSON — for a tool with all
+                    # optional params that silently returns a plausible-but-wrong result. Feed the
+                    # parse error back so the model resends valid arguments.
+                    result = ToolResult(
+                        text=f"error: could not parse tool arguments as JSON ({exc}); resend valid JSON."
+                    )
+                else:
+                    if _needs_confirm(confirm_policy, toolset, c["name"]):
+                        # Gated action: no confirmation channel means deny (fail-safe); otherwise ask.
+                        approved = await on_confirm(c["name"], args) if on_confirm is not None else False
+                    else:
+                        approved = True
+                    if not approved:
+                        # Declined: skip the side-effecting call but still give the model a tool turn
+                        # (mirroring the error-branch shape) so the loop continues with a clear signal.
+                        result = ToolResult(text="error: user declined this action")
+                    else:
+                        try:
+                            result = await toolset.call(c["name"], args, timeout=tool_timeout)
+                        except Exception as exc:  # noqa: BLE001 - report tool failures (incl. timeout) to the model
+                            result = ToolResult(text=f"error: {exc}")
+                display = result.text + (f"  [+{len(result.images)} image(s)]" if result.images else "")
+                await _emit(on_event, "result", display)
+                content = result.text
+                if result.images and vision:
+                    # Feed the pixels back below; leave the tool message a short marker so the
+                    # tool_call_id still has content and the model knows where the image came from.
+                    round_images.extend(result.images)
+                    content = content or "[image returned by the tool; shown in the next message]"
+                elif result.images:
+                    # Text-only model: it can't see the image, so say so rather than drop it silently
+                    # (a screenshot vanishing makes a vision-less model hallucinate what it "saw").
+                    note = "[a tool returned an image, but this model cannot view images]"
+                    content = f"{content}\n{note}" if content else note
+                messages.append({"role": "tool", "tool_call_id": c["id"], "content": content})
+
+            # A vision model reads tool-returned images as its own input: deliver them in a user
+            # message right after the tool results (the exercised multimodal path — an image_url
+            # part in a tool message isn't understood by llama-server/mlx-vlm).
+            if round_images:
+                messages.append({"role": "user", "content": user_content(_TOOL_IMAGE_PREAMBLE, images=round_images)})
+
+    # Ran out of tool rounds: surface a terminal message the same way a normal
+    # reply is delivered — stream it (so streaming clients, incl. the web UI whose
+    # /api/chat discards the return value, actually see it) and record it in history.
+    stopped = "[agent stopped: too many tool rounds]"
+    # _emit, not a bare call: the /api/chat sink is async (queue.put) — calling it
+    # unawaited would silently drop the message for exactly the clients it's for.
+    await _emit(on_token, stopped)
+    messages.append({"role": "assistant", "content": stopped})
+    return stopped