"""
pull — Ollama model pull/remove helpers and env-file writer.
Wraps the ``ollama pull`` and ``ollama rm`` subprocess calls and provides
``write_env``, which writes the validated model selection to a shell-sourceable
``env.sh`` and a machine-readable ``config.json`` under
``~/.best-engine-ai-helper/``.
Author
------
Warith Harchaoui <warithmetics@deraison.ai>
"""
from __future__ import annotations
import json
import os
import subprocess
import sys
import tempfile
import threading
from pathlib import Path
from typing import IO
import os_helper as osh
# User-writable runtime directory; created on first run by write_env
_USER_DIR = Path.home() / ".best-engine-ai-helper"
# File names for the two output formats
_ENV_SH = "env.sh"
_CONFIG_JSON = "config.json"
def _atomic_write_text(path: Path, content: str) -> None:
"""Write ``content`` to ``path`` atomically: no reader ever sees a partial write.
Writes to a sibling temp file in the same directory (so the final
``os.replace`` is a same-filesystem rename, not a cross-filesystem copy)
then renames it over the destination. A crash or a concurrent read mid-
write sees either the old file or the fully-written new one, never a
truncated one.
"""
fd, tmp_name = tempfile.mkstemp(dir=str(path.parent), prefix=f".{path.name}.")
try:
with os.fdopen(fd, "w", encoding="utf-8") as f:
f.write(content)
os.replace(tmp_name, path)
except BaseException:
Path(tmp_name).unlink(missing_ok=True)
raise
[docs]
def context_length(tag: str, base_url: str | None = None, *, timeout: float = 5.0) -> int | None:
"""
The context window this daemon reports for *tag*, read rather than assumed.
Not one of the 31 local catalogue rows declares a ``context_length``, and
none of them should: the number belongs to a particular quantisation of a
particular checkpoint, so a figure typed into ``models.yaml`` would be a
claim about a file nobody here has opened. The paid half gets it for free
because the price source publishes it; the local half has no such source.
It does have something better. A pulled model's GGUF metadata carries the
real figure, and Ollama hands it over on ``/api/show`` under a key named
after the architecture — ``gemma3.context_length``, ``llama.context_length``
— which is why this looks for the suffix rather than a fixed key. So the
number is measured on the machine that holds the weights, exactly like
decode speed is, and models you do not have report nothing at all instead
of a plausible guess.
Parameters
----------
tag : str
Model tag as the daemon knows it, e.g. ``gemma3:12b``.
base_url : str or None
Daemon URL; ``SPREZZATURE_LLM_BASE_URL`` or localhost by default.
timeout : float
Seconds to wait. Longer than :func:`installed_tags`'s because this runs
after a pull rather than inside resolution, and ``/api/show`` has to
read the checkpoint header.
Returns
-------
int or None
The window in tokens, or None when the daemon is absent, the model is
not pulled, or its metadata declares no context length. All three are
ordinary, and none of them is worth a guess.
Examples
--------
A daemon that is not there is not an error:
>>> context_length("gemma3:12b", "http://127.0.0.1:1") is None
True
"""
url = (
base_url or os.environ.get("SPREZZATURE_LLM_BASE_URL") or "http://localhost:11434"
).rstrip("/")
try:
import requests
response = requests.post(f"{url}/api/show", json={"model": tag}, timeout=timeout)
response.raise_for_status()
info = response.json().get("model_info") or {}
except Exception as exc: # noqa: BLE001 — no daemon, or no such model
osh.debug(f"Ollama /api/show unreachable for {tag} ({exc}); context length unknown.")
return None
for key, value in info.items():
if str(key).endswith(".context_length"):
try:
window = int(value)
except (TypeError, ValueError):
continue
if window > 0:
osh.debug(f"{tag}: {key} = {window}")
return window
return None
[docs]
def ollama_pull(tag: str, *, timeout: int = 600, out: IO[str] | None = None) -> bool:
"""
Pull a model via ``ollama pull`` and stream progress to stdout.
Ollama streams progress lines to stderr by default; this function merges
stderr into stdout so the caller sees a unified stream.
Parameters
----------
tag : str
Ollama model tag, e.g. ``"qwen3-vl:8b"`` or ``"qwen3-vl:72b"``.
timeout : int
Maximum seconds to wait for the pull to complete. A 72B model at
Q4_K_M (~52 GB) can take 30+ minutes on a slow connection; the
default of 600 seconds (10 minutes) is generous for fast links.
out : IO[str] or None
Stream to write progress lines to. Defaults to ``sys.stdout``.
Returns
-------
bool
True if ``ollama pull`` exited with code 0; False otherwise.
Raises
------
FileNotFoundError
If the ``ollama`` binary is not on PATH.
subprocess.TimeoutExpired
If the pull is still running after ``timeout`` seconds. The process
is killed before this is raised.
Examples
--------
>>> # ollama_pull("qwen3-vl:8b") # requires ollama running
>>> True # placeholder so doctest passes without Ollama
True
"""
sink = out or sys.stdout
osh.info(f"Pulling model via ollama:\n\t{tag}")
# stderr=subprocess.STDOUT merges the two streams so progress lines appear
try:
proc = subprocess.Popen(
["ollama", "pull", tag],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
)
except FileNotFoundError as exc:
osh.error("ollama binary not found. Install from https://ollama.com")
raise FileNotFoundError("ollama binary not found. Install from https://ollama.com") from exc
# `for line in proc.stdout` below blocks on each read until a line
# arrives or the pipe closes (i.e. until the process exits) -- it has no
# timeout of its own, so a stalled pull (network stall, hung server)
# would otherwise block forever regardless of `timeout`. A watchdog
# thread kills the process if it is still running after `timeout`
# seconds; the read loop then unblocks (naturally, via EOF on the now
# -closed pipe) instead of hanging.
timed_out = threading.Event()
def _kill_on_timeout() -> None:
timed_out.set()
proc.kill()
watchdog = threading.Timer(timeout, _kill_on_timeout)
watchdog.start()
try:
# Stream output line by line so the user sees progress in real time
assert proc.stdout is not None
for line in proc.stdout:
sink.write(line)
sink.flush()
proc.wait()
finally:
watchdog.cancel()
if timed_out.is_set():
osh.error(f"Pull timed out after {timeout}s (process killed):\n\t{tag}")
raise subprocess.TimeoutExpired(["ollama", "pull", tag], timeout)
ok = proc.returncode == 0
if ok:
osh.info(f"Pull succeeded:\n\t{tag}")
else:
osh.error(f"Pull failed (exit {proc.returncode}):\n\t{tag}")
return ok
[docs]
def ollama_rm(tag: str, *, timeout: int = 60) -> bool:
"""
Remove a pulled model via ``ollama rm``.
Used by the pull-and-validate loop to free disk space when a model fails
the Ralph gates, before trying the next candidate.
Parameters
----------
tag : str
Ollama model tag to remove.
timeout : int
Maximum seconds to wait. ``ollama rm`` is a local filesystem delete
(no network, unlike :func:`ollama_pull`), so 60s is generous; a hung
daemon or a locked model file would otherwise block the pull-and-
validate loop forever with no way out.
Returns
-------
bool
True if ``ollama rm`` exited with code 0; False if it exited nonzero
or timed out (the process is killed before returning in that case).
Examples
--------
>>> # ollama_rm("qwen3-vl:8b") # requires ollama running
>>> True # placeholder
True
"""
osh.info(f"Removing model via ollama:\n\t{tag}")
try:
result = subprocess.run(
["ollama", "rm", tag],
capture_output=True,
text=True,
timeout=timeout,
)
except subprocess.TimeoutExpired:
osh.error(f"Removal timed out after {timeout}s (process killed):\n\t{tag}")
return False
ok = result.returncode == 0
if not ok:
osh.warning(f"Removal failed (exit {result.returncode}):\n\t{tag}")
return ok
[docs]
def write_env(
text_model: str,
vision_model: str,
backend: str,
base_url: str,
*,
user_dir: Path | None = None,
) -> Path:
"""
Write the validated model selection to ``env.sh`` and ``config.json``.
Both files are written atomically: ``env.sh`` is shell-sourceable and
suitable for ``.envrc`` (direnv); ``config.json`` is for programmatic
consumers. The directory is created if it does not exist.
Parameters
----------
text_model : str
Ollama tag for text-only tasks (``BEST_LLM_TEXT``).
vision_model : str
Ollama tag for vision tasks (``BEST_LLM_VISION``).
backend : str
Backend name: ``"ollama"``, ``"openai"``, or ``"langchain"``.
base_url : str
Base URL of the inference server.
user_dir : Path or None
Override the default ``~/.best-engine-ai-helper/`` directory. Used
in tests to avoid touching the real user home directory.
Returns
-------
Path
Absolute path to the written ``env.sh`` file.
Examples
--------
>>> import tempfile, pathlib
>>> with tempfile.TemporaryDirectory() as td:
... p = write_env("qwen3-vl:8b", "qwen3-vl:8b", "ollama",
... "http://localhost:11434", user_dir=pathlib.Path(td))
... p.name
'env.sh'
"""
target = user_dir if user_dir is not None else _USER_DIR
osh.make_directory(str(target))
env_sh_path = target / _ENV_SH
config_path = target / _CONFIG_JSON
# Shell-sourceable env block: safe for ~/.zshrc or direnv .envrc
env_content = (
"# generated by best-engine-ai-helper — do not edit by hand\n"
f"export BEST_LLM_TEXT={text_model}\n"
f"export BEST_LLM_VISION={vision_model}\n"
f"export BEST_LLM_BACKEND={backend}\n"
f"export BEST_LLM_BASE_URL={base_url}\n"
)
_atomic_write_text(env_sh_path, env_content)
osh.info(f"Wrote shell env file:\n\t{env_sh_path}")
# JSON format for scripts that prefer structured config over shell sourcing
config = {
"BEST_LLM_TEXT": text_model,
"BEST_LLM_VISION": vision_model,
"BEST_LLM_BACKEND": backend,
"BEST_LLM_BASE_URL": base_url,
}
_atomic_write_text(config_path, json.dumps(config, indent=2) + "\n")
osh.info(f"Wrote JSON config file:\n\t{config_path}")
return env_sh_path