| """Local GGUF LLM client (llama.cpp). Lazy-loaded + cached; no network at generate time. |
| |
| ``llama_cpp`` is imported INSIDE the loader (not at module top) so importing this module — and running the |
| gates — needs no compiled backend or downloaded weights. Models cache under ``data/artifacts/models/`` via |
| an ``HF_HOME`` redirect, so once fetched the run is fully offline. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import os |
| import re |
| import threading |
| import time |
| from functools import lru_cache |
| from typing import TYPE_CHECKING, Any |
|
|
| from common.logging import get_logger |
| from common.paths import REPO_ROOT |
|
|
| if TYPE_CHECKING: |
| from config import Settings |
|
|
| MODEL_DIR = REPO_ROOT / "data" / "artifacts" / "models" |
|
|
|
|
| def cached_hf_download(repo_id: str, filename: str) -> str: |
| """Resolve ``filename`` from ``repo_id``: local cache first (zero network), download only if missing. |
| |
| Plain ``hf_hub_download`` fires a live HEAD request on every call — even for a fully-cached file — to |
| validate the etag, which is a real network attempt the "no network during ranking" contract shouldn't |
| have to tolerate. ``try_to_load_from_cache`` gives a tri-state local-only answer (found / confirmed |
| absent via a ``.no_exist`` marker from a prior check / never checked) — only the last state should ever |
| reach the network, otherwise a known-missing optional file (e.g. a model without a sibling |
| ``.onnx_data``) would retry the Hub on every single load forever.""" |
| from huggingface_hub import _CACHED_NO_EXIST, hf_hub_download, try_to_load_from_cache |
|
|
| cached = try_to_load_from_cache(repo_id, filename) |
| if isinstance(cached, str): |
| return cached |
| if cached is _CACHED_NO_EXIST: |
| raise FileNotFoundError(f"{filename!r} confirmed absent from {repo_id!r} (cached .no_exist marker)") |
| return hf_hub_download(repo_id=repo_id, filename=filename) |
|
|
| |
| |
| |
| _THINK_BLOCK = re.compile(r"<think>.*?</think>", re.DOTALL) |
| _THINK_DANGLING = re.compile(r"<think>.*$", re.DOTALL) |
|
|
|
|
| def _strip_thinking(text: str) -> str: |
| """Remove a Qwen3 <think> block (closed or truncated) and surrounding whitespace.""" |
| return _THINK_DANGLING.sub("", _THINK_BLOCK.sub("", text)).strip() |
|
|
|
|
| _HEARTBEAT_SECONDS = 5.0 |
|
|
|
|
| def _generate_streaming(model: Any, messages: list, temperature: float, top_p: float, max_tokens: int, log: Any) -> str: |
| """Run ``create_chat_completion(stream=True)`` and log a heartbeat every ``_HEARTBEAT_SECONDS`` — gives |
| visibility into whether a long call is actively progressing (chars accumulating) or genuinely stuck |
| (heartbeats stop appearing), which a single before/after log pair can't distinguish. Used for EVERY call, |
| not just timed ones — JD parse has no timeout by design, so this is the only way to tell "slow but working" |
| from "hung" without cutting anything off.""" |
| chunks: list[str] = [] |
| start = time.perf_counter() |
| last_heartbeat = start |
| for chunk in model.create_chat_completion( |
| messages=messages, temperature=temperature, top_p=top_p, max_tokens=max_tokens, stream=True |
| ): |
| piece = chunk["choices"][0].get("delta", {}).get("content") or "" |
| if piece: |
| chunks.append(piece) |
| now = time.perf_counter() |
| if now - last_heartbeat >= _HEARTBEAT_SECONDS: |
| log.info("llm.generate_heartbeat", elapsed=round(now - start, 1), chars_so_far=sum(len(c) for c in chunks)) |
| last_heartbeat = now |
| return "".join(chunks) |
|
|
|
|
| def _run_with_timeout(fn: Any, timeout: float) -> Any: |
| """Run ``fn()`` on a daemon thread and wait up to ``timeout`` seconds; raise ``TimeoutError`` if it's still |
| running. llama.cpp offers no cooperative cancellation mid-generation, so a genuinely stuck/runaway call |
| can't be killed — the thread is simply abandoned (daemon=True so it can't block process exit) and the |
| caller moves on. This exists because a single generation was observed taking 1996s (33 min) on real |
| deployment hardware with NOTHING to stop it — the between-call budget check in tiers.py can't help since |
| it only gates the NEXT call, not one already in flight.""" |
| result: dict[str, Any] = {} |
|
|
| def target() -> None: |
| try: |
| result["value"] = fn() |
| except Exception as exc: |
| result["error"] = exc |
|
|
| thread = threading.Thread(target=target, daemon=True) |
| thread.start() |
| thread.join(timeout) |
| if thread.is_alive(): |
| raise TimeoutError(f"generation exceeded {timeout}s — abandoned (still running in the background)") |
| if "error" in result: |
| raise result["error"] |
| return result["value"] |
|
|
|
|
| @lru_cache(maxsize=2) |
| def _load_llama(repo: str, gguf: str, n_ctx: int, n_threads: int) -> Any: |
| """Load a GGUF model into a cached llama.cpp handle. |
| |
| An EXACT ``gguf`` filename is resolved with ``cached_hf_download`` (cache-only lookup first, no Hub API |
| call once baked — falls back to a real fetch only if missing). ``Llama.from_pretrained`` is NOT used for |
| exact names because it always lists the repo over the network. A glob ``gguf`` (the rare fallback model) |
| still routes through ``from_pretrained`` (needs network).""" |
| from llama_cpp import Llama |
|
|
| MODEL_DIR.mkdir(parents=True, exist_ok=True) |
| os.environ.setdefault("HF_HOME", str(MODEL_DIR)) |
| log = get_logger("models") |
| log.info("llm.load", repo=repo, gguf=gguf, n_ctx=n_ctx, n_threads=n_threads) |
| start = time.perf_counter() |
| try: |
| if "*" in gguf or "?" in gguf: |
| model = Llama.from_pretrained(repo_id=repo, filename=gguf, n_ctx=n_ctx, |
| n_threads=(n_threads or None), verbose=False) |
| else: |
| model_path = cached_hf_download(repo, gguf) |
| model = Llama(model_path=model_path, n_ctx=n_ctx, n_threads=(n_threads or None), verbose=False) |
| except Exception as exc: |
| log.error("llm.load_failed", repo=repo, gguf=gguf, seconds=round(time.perf_counter() - start, 3), error=str(exc)) |
| raise |
| log.info("llm.load_complete", repo=repo, seconds=round(time.perf_counter() - start, 3)) |
| return model |
|
|
|
|
| class LlmClient: |
| """Thin wrapper over a cached GGUF model exposing a single ``generate`` call.""" |
|
|
| def __init__(self, repo: str, gguf: str, n_ctx: int = 4096, n_threads: int = 0) -> None: |
| self._repo, self._gguf, self._n_ctx, self._n_threads = repo, gguf, n_ctx, n_threads |
| |
| |
| self._soft_no_think = "qwen" in repo.lower() |
|
|
| def ensure(self) -> None: |
| """Force the (down)load now — used by setup so the rank/JD run doesn't pay it on the clock.""" |
| _load_llama(self._repo, self._gguf, self._n_ctx, self._n_threads) |
|
|
| def generate( |
| self, prompt: str, system: str | None = None, |
| temperature: float = 0.0, top_p: float = 0.95, max_tokens: int = 512, think: bool = False, |
| timeout: float | None = None, |
| ) -> str: |
| """Chat-complete ``prompt`` (temperature=0 ⇒ greedy). Reasoning suppressed + stripped by default. |
| |
| Always streams internally with a periodic heartbeat log (``llm.generate_heartbeat`` every ~5s) — for a |
| call with no ``timeout`` (JD parse, by design), that heartbeat is the ONLY way to tell "slow but still |
| producing tokens" from "genuinely hung" during a long silence; without it there's nothing between |
| ``generate_start`` and either ``generate_complete`` or nothing at all. |
| |
| ``timeout``, if given, is a hard wall-clock cap enforced via ``_run_with_timeout`` — a real generation |
| was observed taking 1996s with nothing to stop it, so callers on a wall-clock budget (Step 11 polish) |
| should always pass one; JD parse intentionally does NOT (its accuracy matters more than its wall-clock). |
| """ |
| model = _load_llama(self._repo, self._gguf, self._n_ctx, self._n_threads) |
| |
| user = prompt + " /no_think" if self._soft_no_think and not think else prompt |
| messages = ([{"role": "system", "content": system}] if system else []) + [ |
| {"role": "user", "content": user} |
| ] |
| log = get_logger("models") |
| log.info("llm.generate_start", prompt_chars=len(user), max_tokens=max_tokens, timeout=timeout) |
| start = time.perf_counter() |
| call = lambda: _generate_streaming(model, messages, temperature, top_p, max_tokens, log) |
| try: |
| raw_text = _run_with_timeout(call, timeout) if timeout else call() |
| except Exception as exc: |
| log.error("llm.generate_failed", seconds=round(time.perf_counter() - start, 3), error=str(exc)) |
| raise |
| text = _strip_thinking(raw_text) |
| log.info("llm.generate_complete", seconds=round(time.perf_counter() - start, 3), chars_out=len(text)) |
| return text |
|
|
|
|
| def index_gen_client(settings: Settings) -> LlmClient: |
| """The JD-parse LLM client (Qwen3-1.7B GGUF, standard instruct) built from central config. |
| |
| ``models.n_threads`` is llama.cpp's OWN explicit override; when it's left at its default (0), the shared |
| ``compute.cpu_threads`` budget applies instead, so one config knob caps the whole engine consistently. |
| """ |
| from common.runtime import resolve_cpu_threads |
|
|
| models = settings.models |
| threads = models.n_threads or resolve_cpu_threads(settings.compute.cpu_threads) |
| return LlmClient(models.index_gen_llm_repo, models.index_gen_llm_gguf, models.n_ctx, threads) |
|
|
