Ported from downstream development. Four independent defects.
1. The memory dump was unrestorable. iterdump() serializes sqlite_vec virtual
tables as a raw INSERT INTO sqlite_master(...) followed by inserts into a
table the replaying connection cannot see, so replaying memory.db.sql died
on "no such table: vec_messages" and left ZERO tables behind. dump_db() now
loads the vec0 extension and filters the derived vec tables out of the
iterdump stream, matched on each statement's target table rather than as a
substring - a chat message whose text mentions vec_messages is an
INSERT INTO "messages" and has to survive.
compare() reported an unreadable dump as "diverged", which read like a real
verdict and made both guards refuse backup AND restore, locking the machine
out of syncing in either direction. Unreadable is now its own verdict.
_extra() compared updated_at against a "" default, but the column is REAL,
so the comparison raises TypeError on the first conversation the other side
lacks - exactly the case it counts. It tests membership first now. The
direction test declared updated_at TEXT, which is why this survived: the
test compared str to str while the field compared str to float.
2. The memory curator invented facts. It attributed the ASSISTANT's words to
the user, wrote absence claims read off the existing-memory block, and added
judgements ("favorite") the user never used. The prompt now scopes the USER
line as the only source, and two deterministic guards drop absence claims
and facts whose distinctive tokens appear nowhere in the user's message -
prompt wording alone did not hold on a 7B curator.
3. _best_vulkan_device scored Mesa's llvmpipe above an integrated GPU, pinning
Ollama to a software rasterizer advertising 31 GiB of "VRAM" - CPU inference
with Vulkan overhead on top. Software rasterizers are dropped.
4. Models.jsx compared catalog names to installed names literally, but Ollama
resolves a bare name to ":latest", so an untagged entry (nomic-embed-text)
read as missing forever and the Required gate never opened. Chatbot.jsx
fetched the model list once on mount although App keeps the page mounted
behind display:none, so a newly pulled model never appeared in the picker
until a full browser reload.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
759 lines
30 KiB
Python
759 lines
30 KiB
Python
import asyncio
|
||
import json
|
||
import logging
|
||
import subprocess
|
||
import time
|
||
import httpx
|
||
import os
|
||
import signal
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from .nexus_config import settings, DEFAULT_CHAT_MODEL, DEFAULT_EMBED_MODEL
|
||
from .memory.store import store
|
||
|
||
OLLAMA_PORT = 11434
|
||
|
||
_log = logging.getLogger("nexus.ollama")
|
||
|
||
# --- Global Singleton Instance ---
|
||
_ollama_manager = None
|
||
|
||
# Bundled binary ships alongside the project; fall back to system PATH
|
||
_BUNDLED_OLLAMA = Path(__file__).resolve().parent.parent / "ollama" / "bin" / "ollama"
|
||
|
||
# POSIX: detach the child into its own session so we can signal the whole group.
|
||
# Windows has no setsid/killpg — run the child normally and terminate() it.
|
||
_DETACH_KW = {} if os.name == "nt" else {"start_new_session": True}
|
||
|
||
|
||
def _ollama_bin() -> str:
|
||
"""Return path to the Ollama executable, preferring the bundled copy."""
|
||
if _BUNDLED_OLLAMA.exists():
|
||
return str(_BUNDLED_OLLAMA)
|
||
return "ollama"
|
||
|
||
|
||
def _best_vulkan_device() -> tuple[int, str]:
|
||
"""
|
||
Parse `vulkaninfo --summary` and return (device_index, device_name) for the
|
||
best Vulkan compute device. Prefers discrete GPUs over integrated ones, and
|
||
AMD/NVIDIA vendor IDs over Intel — so a Radeon is chosen over an Intel iGPU
|
||
even when the iGPU appears first in the device list.
|
||
|
||
Software rasterizers (llvmpipe/lavapipe, PHYSICAL_DEVICE_TYPE_CPU) are
|
||
dropped outright: Mesa always advertises one, it is CPU inference wearing a
|
||
GPU costume, and it is *slower* than the plain CPU backend because every
|
||
tensor takes a detour through Vulkan. Returns (-1, "") when no real GPU is
|
||
present so the caller falls back instead of pinning the rasterizer.
|
||
"""
|
||
try:
|
||
r = subprocess.run(
|
||
["vulkaninfo", "--summary"], capture_output=True, text=True, timeout=5,
|
||
)
|
||
if r.returncode != 0:
|
||
return -1, ""
|
||
|
||
devices: list[dict] = []
|
||
current: dict = {}
|
||
for line in r.stdout.splitlines():
|
||
line = line.strip()
|
||
if line.startswith("GPU") and line.endswith(":"):
|
||
if current:
|
||
devices.append(current)
|
||
raw_idx = line[3:-1]
|
||
current = {
|
||
"index": int(raw_idx) if raw_idx.isdigit() else len(devices),
|
||
"name": "GPU",
|
||
"type": "",
|
||
"vendor": "",
|
||
}
|
||
elif "=" in line:
|
||
key, _, val = line.partition("=")
|
||
key, val = key.strip(), val.strip()
|
||
if key == "deviceName":
|
||
current["name"] = val
|
||
elif key == "deviceType":
|
||
current["type"] = val.upper()
|
||
elif key == "vendorID":
|
||
current["vendor"] = val.lower()
|
||
|
||
if current:
|
||
devices.append(current)
|
||
|
||
devices = [d for d in devices if "CPU" not in d["type"]]
|
||
if not devices:
|
||
return -1, ""
|
||
|
||
def _score(d: dict) -> tuple:
|
||
# Discrete beats everything; integrated is last resort
|
||
type_score = 2 if "DISCRETE" in d["type"] else (0 if "INTEGRATED" in d["type"] else 1)
|
||
# AMD (0x1002) and NVIDIA (0x10de) preferred over Intel (0x8086)
|
||
vendor_score = 1 if any(v in d["vendor"] for v in ("0x1002", "0x10de")) else 0
|
||
return (type_score, vendor_score)
|
||
|
||
best = max(devices, key=_score)
|
||
return best["index"], best["name"]
|
||
|
||
except Exception:
|
||
return -1, ""
|
||
|
||
|
||
def _detect_gpu_backend() -> tuple[str, dict]:
|
||
"""
|
||
Probe available GPU compute backends.
|
||
Returns (label, env_overrides) where env_overrides is merged into the
|
||
Ollama subprocess environment before launch.
|
||
Priority: CUDA > Vulkan > ROCm > CPU.
|
||
"""
|
||
# NVIDIA CUDA — preferred when both GPU and CUDA drivers are present
|
||
try:
|
||
r = subprocess.run(
|
||
["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
|
||
capture_output=True, text=True, timeout=5,
|
||
)
|
||
if r.returncode == 0:
|
||
name = r.stdout.strip().splitlines()[0]
|
||
_log.info("GPU backend: CUDA (%s)", name)
|
||
return f"cuda ({name})", {} # Ollama auto-detects CUDA
|
||
except (FileNotFoundError, subprocess.TimeoutExpired):
|
||
pass
|
||
|
||
# Vulkan — works on AMD, Intel, and NVIDIA without a full CUDA/ROCm stack
|
||
vulkan_ok = False
|
||
try:
|
||
r = subprocess.run(
|
||
["vulkaninfo", "--summary"], capture_output=True, text=True, timeout=5,
|
||
)
|
||
vulkan_ok = r.returncode == 0
|
||
except (FileNotFoundError, subprocess.TimeoutExpired):
|
||
pass
|
||
|
||
if not vulkan_ok:
|
||
# Fall back to checking for ICD loader files without the vulkaninfo tool
|
||
icd_dirs = [
|
||
Path("/usr/share/vulkan/icd.d"),
|
||
Path("/etc/vulkan/icd.d"),
|
||
Path(os.path.expanduser("~/.local/share/vulkan/icd.d")),
|
||
]
|
||
try:
|
||
vulkan_ok = any(p.is_dir() and any(p.iterdir()) for p in icd_dirs)
|
||
except PermissionError:
|
||
pass
|
||
|
||
if vulkan_ok:
|
||
idx, name = _best_vulkan_device()
|
||
if idx >= 0:
|
||
# Always pin to the selected device — without this, Ollama may use the Intel
|
||
# iGPU's shared system RAM as "VRAM" for models that don't fit on discrete VRAM.
|
||
env_overrides: dict = {"OLLAMA_VULKAN": "1", "GGML_VK_VISIBLE_DEVICES": str(idx)}
|
||
_log.info("GPU backend: Vulkan device %d (%s)", idx, name)
|
||
return f"vulkan ({name})", env_overrides
|
||
# Vulkan loads but every device is a software rasterizer — no real GPU here.
|
||
|
||
# AMD ROCm — fallback when Vulkan ICD is absent but ROCm stack is installed
|
||
try:
|
||
r = subprocess.run(
|
||
["rocm-smi", "--showproductname"], capture_output=True, text=True, timeout=5,
|
||
)
|
||
if r.returncode == 0:
|
||
_log.info("GPU backend: ROCm")
|
||
return "rocm", {} # Ollama auto-detects ROCm
|
||
except (FileNotFoundError, subprocess.TimeoutExpired):
|
||
pass
|
||
|
||
_log.info("GPU backend: CPU (no GPU acceleration detected)")
|
||
return "cpu", {}
|
||
|
||
|
||
# Single source of model auto-selection preference. Ordered so small, GPU-fitting
|
||
# models come first (qwen2.5:3b fits a 4GB card and is a strong all-rounder); the
|
||
# tail differs by task. Prefix-matched against installed model names.
|
||
_MODEL_PREFERENCE = {
|
||
"chat": ("qwen2.5:3b", "qwen2.5", "gemma3:1b", "gemma3", "phi3", "phi-3", "gemma2", "gemma"),
|
||
"code": ("qwen2.5-coder", "qwen3-coder", "deepseek-coder", "codellama", "codegemma", "qwen2.5:3b", "qwen2.5", "gemma3", "phi3", "phi-3"),
|
||
}
|
||
|
||
|
||
def _preferred_model(models: list, preference) -> str | None:
|
||
"""First installed model whose name starts with a preference prefix."""
|
||
for prefix in preference:
|
||
for m in models:
|
||
if m.lower().startswith(prefix.lower()):
|
||
return m
|
||
return None
|
||
|
||
|
||
def _chat_options(temperature: float | None, num_gpu: int | None, num_ctx: int | None = None) -> dict:
|
||
"""Assemble the Ollama `options` block from the knobs we expose.
|
||
|
||
Returns an empty dict when nothing is set so callers can omit `options`
|
||
entirely (preserving Ollama's defaults / auto behaviour).
|
||
"""
|
||
opts: dict = {}
|
||
if temperature is not None:
|
||
opts["temperature"] = temperature
|
||
if num_gpu is not None:
|
||
opts["num_gpu"] = num_gpu
|
||
if num_ctx: # 0 / None -> let Ollama use the model default
|
||
opts["num_ctx"] = num_ctx
|
||
return opts
|
||
|
||
|
||
class OllamaManager:
|
||
def __init__(self, runtime_dir=None):
|
||
self.process = None
|
||
self.running = False
|
||
self._available = None # see is_available()
|
||
|
||
self.runtime_dir = Path(runtime_dir) if runtime_dir else Path(__file__).resolve().parent.parent / "runtime"
|
||
(self.runtime_dir / "logs").mkdir(parents=True, exist_ok=True)
|
||
|
||
self.log_file = self.runtime_dir / "logs" / "ollama.log"
|
||
|
||
# Resolved at construction so host changes in settings take effect
|
||
self._api_base = settings.ollama_host.rstrip("/")
|
||
|
||
# Model selection cache
|
||
self._model_cache: dict = {} # intent -> (model, monotonic_ts)
|
||
|
||
# Per-model offloadable layer count cache (never changes for a model)
|
||
self._layer_cache: dict[str, int] = {}
|
||
|
||
# How long Ollama keeps the model resident between requests. Applied to
|
||
# every chat/generate body so the model isn't reloaded on each message.
|
||
# Overridden from persisted settings at startup. Falsy → omit (Ollama's
|
||
# 5-minute default).
|
||
self.keep_alive: str | None = "30m"
|
||
|
||
def _apply_keep_alive(self, body: dict) -> dict:
|
||
"""Add `keep_alive` (a top-level Ollama field) to a request body when set."""
|
||
if self.keep_alive:
|
||
body["keep_alive"] = self.keep_alive
|
||
return body
|
||
|
||
async def warm(self, model: str | None = None, num_gpu: int | None = None) -> None:
|
||
"""Preload a model so the first request doesn't pay a cold load.
|
||
An empty-prompt /api/generate is Ollama's documented preload. Pass the
|
||
same `num_gpu` the real requests use, or the preloaded copy is placed
|
||
differently and gets reloaded on first use. Best-effort: never raises,
|
||
so a missing model or down server can't break startup."""
|
||
try:
|
||
model = model or await self.select_best_model()
|
||
if not model:
|
||
return
|
||
body = {"model": model, "prompt": "", "stream": False}
|
||
if num_gpu is not None:
|
||
body["options"] = {"num_gpu": num_gpu}
|
||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||
await client.post(
|
||
f"{self._api_base}/api/generate", json=self._apply_keep_alive(body),
|
||
)
|
||
_log.info("warm: preloaded model=%s num_gpu=%s keep_alive=%s", model, num_gpu, self.keep_alive)
|
||
except Exception as e:
|
||
_log.warning("warm: preload failed: %s", e)
|
||
|
||
def _serve_env(self, gpu_env: dict) -> dict:
|
||
"""Environment for a `serve` we spawn. One definition - this was
|
||
duplicated verbatim in two methods, so the Windows carve-out below had
|
||
to be fixed in both places or the two paths would disagree."""
|
||
env = os.environ.copy()
|
||
env["OLLAMA_HOST"] = self._api_base
|
||
# Every platform uses the project's own model store. Windows used to be
|
||
# exempt, because the installer pulled with a bare `ollama pull` into
|
||
# %USERPROFILE%\.ollama and a NexusOS-spawned serve pointed elsewhere
|
||
# would not have seen those models. The installer sets OLLAMA_MODELS for
|
||
# its pulls now (and migrates an existing store), so the exemption just
|
||
# meant Windows kept models somewhere `ncp models` could not see.
|
||
env["OLLAMA_MODELS"] = str(settings.models_dir)
|
||
env.update(gpu_env)
|
||
return env
|
||
|
||
def is_available(self):
|
||
# ponytail: cached for the life of the process. This spawns a subprocess,
|
||
# and /status calls it on every poll - the frontend polls continuously,
|
||
# so it was a process spawn per tick to answer a question whose answer
|
||
# cannot change without someone installing Ollama. Restart to re-detect.
|
||
if self._available is None:
|
||
bin_path = _ollama_bin()
|
||
try:
|
||
subprocess.run([bin_path, "--version"], capture_output=True,
|
||
check=True, timeout=5)
|
||
self._available = True
|
||
except Exception:
|
||
self._available = False
|
||
return self._available
|
||
|
||
def is_running(self):
|
||
# 0.5s, not 2s: this is a loopback request to a server that is either
|
||
# listening or is not. Ollama ships OFF (the user presses Start AI), so
|
||
# the "not running" path is the common one and every /status paid the
|
||
# full 2s for it - which pushed /status past the 2s client timeout in
|
||
# bin/nexus_window.py and made a healthy backend look dead.
|
||
try:
|
||
r = httpx.get(f"{self._api_base}/api/tags", timeout=0.5)
|
||
return r.status_code == 200
|
||
except httpx.RequestError:
|
||
return False
|
||
|
||
def start(self):
|
||
if not self.is_available():
|
||
_log.warning("Ollama not found at %s; skipping startup", _ollama_bin())
|
||
return False
|
||
|
||
if self.is_running():
|
||
_log.info("Ollama already running (%s)", self._api_base)
|
||
self.running = True
|
||
return True
|
||
|
||
try:
|
||
backend, gpu_env = _detect_gpu_backend()
|
||
_log.info("Starting Ollama service via %s (backend: %s)...", _ollama_bin(), backend)
|
||
env = self._serve_env(gpu_env)
|
||
|
||
with open(self.log_file, "w") as log:
|
||
self.process = subprocess.Popen(
|
||
[_ollama_bin(), "serve"],
|
||
stdout=log,
|
||
stderr=subprocess.STDOUT,
|
||
env=env,
|
||
**_DETACH_KW,
|
||
)
|
||
|
||
for attempt in range(30):
|
||
if self.is_running():
|
||
_log.info("Ollama service started (%s)", self._api_base)
|
||
self.running = True
|
||
return True
|
||
time.sleep(1)
|
||
if attempt % 5 == 0:
|
||
_log.info("Waiting for Ollama... (%ds)", attempt)
|
||
|
||
_log.error("Ollama failed to start: timeout")
|
||
return False
|
||
|
||
except Exception as e:
|
||
_log.exception("Ollama failed to start: %s", e)
|
||
return False
|
||
|
||
async def start_async(self):
|
||
"""Async-safe version of start() for use inside async startup handlers."""
|
||
if not self.is_available():
|
||
_log.warning("Ollama not found at %s; skipping startup", _ollama_bin())
|
||
return False
|
||
|
||
if self.is_running():
|
||
_log.info("Ollama already running (%s)", self._api_base)
|
||
self.running = True
|
||
return True
|
||
|
||
try:
|
||
backend, gpu_env = _detect_gpu_backend()
|
||
_log.info("Starting Ollama service via %s (backend: %s)...", _ollama_bin(), backend)
|
||
env = self._serve_env(gpu_env)
|
||
|
||
with open(self.log_file, "w") as log:
|
||
self.process = subprocess.Popen(
|
||
[_ollama_bin(), "serve"],
|
||
stdout=log,
|
||
stderr=subprocess.STDOUT,
|
||
env=env,
|
||
**_DETACH_KW,
|
||
)
|
||
|
||
for attempt in range(30):
|
||
if self.is_running():
|
||
_log.info("Ollama service started (%s)", self._api_base)
|
||
self.running = True
|
||
return True
|
||
await asyncio.sleep(1)
|
||
if attempt % 5 == 0:
|
||
_log.info("Waiting for Ollama... (%ds)", attempt)
|
||
|
||
_log.error("Ollama failed to start: timeout")
|
||
return False
|
||
|
||
except Exception as e:
|
||
_log.exception("Ollama failed to start: %s", e)
|
||
return False
|
||
|
||
def stop(self):
|
||
# Terminate a server we spawned ourselves.
|
||
if self.process:
|
||
try:
|
||
_log.info("Stopping Ollama service (owned process)...")
|
||
if os.name == "nt":
|
||
self.process.terminate()
|
||
else:
|
||
os.killpg(os.getpgid(self.process.pid), signal.SIGTERM)
|
||
self.process.wait(timeout=10)
|
||
_log.info("Ollama service stopped")
|
||
except Exception:
|
||
try:
|
||
if os.name == "nt":
|
||
self.process.kill()
|
||
else:
|
||
os.killpg(os.getpgid(self.process.pid), signal.SIGKILL)
|
||
except Exception:
|
||
pass
|
||
finally:
|
||
self.process = None
|
||
|
||
# Ollama may still be up because something else started it (e.g. the
|
||
# Windows desktop app autostarts one). The manual Stop button should
|
||
# still stop the AI, so kill any remaining server by name. Best-effort:
|
||
# an elevated instance can't be killed from a user-level process, so log
|
||
# rather than raise.
|
||
if self.is_running():
|
||
_log.info("Stopping externally-started Ollama...")
|
||
try:
|
||
if os.name == "nt":
|
||
for image in ("ollama app.exe", "ollama.exe"):
|
||
subprocess.run(
|
||
["taskkill", "/F", "/T", "/IM", image],
|
||
capture_output=True, timeout=10,
|
||
)
|
||
else:
|
||
subprocess.run(
|
||
["pkill", "-f", "ollama serve"], capture_output=True, timeout=10,
|
||
)
|
||
except Exception as e:
|
||
_log.warning("external Ollama stop failed: %s", e)
|
||
|
||
self.running = False
|
||
|
||
def get_status(self):
|
||
if self.is_running():
|
||
return "running"
|
||
elif self.is_available():
|
||
return "available"
|
||
else:
|
||
return "unavailable"
|
||
|
||
async def generate(
|
||
self,
|
||
prompt: str,
|
||
model: str = DEFAULT_CHAT_MODEL,
|
||
stream: bool = False,
|
||
system: str = "",
|
||
**kwargs
|
||
):
|
||
start = time.perf_counter()
|
||
|
||
_log.debug("generate model=%s system=%r prompt=%.120s", model, system, prompt)
|
||
|
||
try:
|
||
if stream:
|
||
return self._stream(prompt=prompt, model=model, system=system, start=start)
|
||
else:
|
||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||
r = await client.post(
|
||
f"{self._api_base}/api/generate",
|
||
json=self._apply_keep_alive({
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"system": system,
|
||
"stream": False,
|
||
}),
|
||
)
|
||
|
||
elapsed = time.perf_counter() - start
|
||
_log.info("generate completed model=%s status=%d elapsed=%.3fs", model, r.status_code, elapsed)
|
||
r.raise_for_status()
|
||
return r.json().get("response", "")
|
||
|
||
except Exception as e:
|
||
elapsed = time.perf_counter() - start
|
||
_log.exception("generate error after %.3fs: %s", elapsed, e)
|
||
return None
|
||
|
||
async def _stream(self, prompt: str, model: str, system: str, start: float):
|
||
"""
|
||
Async generator that streams token chunks from Ollama.
|
||
Yields string chunks as they arrive.
|
||
"""
|
||
try:
|
||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||
async with client.stream(
|
||
"POST",
|
||
f"{self._api_base}/api/generate",
|
||
json=self._apply_keep_alive({
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"system": system,
|
||
"stream": True,
|
||
}),
|
||
) as response:
|
||
response.raise_for_status()
|
||
async for line in response.aiter_lines():
|
||
if not line.strip():
|
||
continue
|
||
try:
|
||
data = json.loads(line)
|
||
token = data.get("response", "")
|
||
if token:
|
||
yield token
|
||
if data.get("done", False):
|
||
elapsed = time.perf_counter() - start
|
||
_log.info("generate stream completed model=%s elapsed=%.3fs", model, elapsed)
|
||
break
|
||
except Exception:
|
||
continue
|
||
|
||
except Exception as e:
|
||
elapsed = time.perf_counter() - start
|
||
_log.exception("generate stream error after %.3fs: %s", elapsed, e)
|
||
return
|
||
|
||
|
||
async def chat(
|
||
self,
|
||
messages: list,
|
||
model: str = DEFAULT_CHAT_MODEL,
|
||
stream: bool = False,
|
||
temperature: float | None = None,
|
||
num_gpu: int | None = None,
|
||
think: bool = False,
|
||
tools: list | None = None,
|
||
num_ctx: int | None = None,
|
||
**kwargs,
|
||
):
|
||
"""Multi-turn chat via /api/chat (accepts a messages array with roles).
|
||
|
||
When `tools` is given (non-stream only), the request advertises them and
|
||
the FULL message dict is returned (so the caller sees `tool_calls`);
|
||
otherwise the response content string is returned as before.
|
||
|
||
`think` toggles Qwen3-style reasoning. Default off: the hidden <think>
|
||
block is pure latency for chat/memory. Ollama ignores it for models that
|
||
don't support thinking.
|
||
"""
|
||
start = time.perf_counter()
|
||
try:
|
||
if stream:
|
||
return self._chat_stream(
|
||
messages=messages, model=model, temperature=temperature,
|
||
num_gpu=num_gpu, think=think, start=start, num_ctx=num_ctx,
|
||
)
|
||
else:
|
||
body: dict = {"model": model, "messages": messages, "stream": False}
|
||
body["think"] = think
|
||
if tools:
|
||
body["tools"] = tools
|
||
opts = _chat_options(temperature, num_gpu, num_ctx)
|
||
if opts:
|
||
body["options"] = opts
|
||
self._apply_keep_alive(body)
|
||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||
r = await client.post(f"{self._api_base}/api/chat", json=body)
|
||
elapsed = time.perf_counter() - start
|
||
r.raise_for_status()
|
||
message = r.json().get("message", {})
|
||
# Tool callers need the whole message (tool_calls); others want content.
|
||
return message if tools else message.get("content", "")
|
||
except Exception as e:
|
||
elapsed = time.perf_counter() - start
|
||
_log.exception("chat error after %.3fs: %s", elapsed, e)
|
||
return None
|
||
|
||
async def embed(self, text: str, model: str = DEFAULT_EMBED_MODEL) -> list[float] | None:
|
||
"""Return an embedding vector for `text` via /api/embeddings.
|
||
|
||
Returns None on any failure so callers can fall back to lexical search —
|
||
a missing embedding model should never break chat or recall.
|
||
"""
|
||
text = (text or "").strip()
|
||
if not text:
|
||
return None
|
||
try:
|
||
gpu_offload = store.get_settings().get("memory_gpu_offload", 0)
|
||
num_gpu = await self.resolve_num_gpu(gpu_offload, model)
|
||
body: dict[str, Any] = {"model": model, "prompt": text}
|
||
if num_gpu is not None:
|
||
body["options"] = {"num_gpu": num_gpu}
|
||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||
r = await client.post(
|
||
f"{self._api_base}/api/embeddings",
|
||
json=body,
|
||
)
|
||
r.raise_for_status()
|
||
vec = r.json().get("embedding")
|
||
return vec if vec else None
|
||
except Exception as e:
|
||
_log.debug("embed failed (model=%s): %s", model, e)
|
||
return None
|
||
|
||
async def list_models(self) -> list[str]:
|
||
"""Return names of all locally installed Ollama models."""
|
||
try:
|
||
async with httpx.AsyncClient(timeout=5.0) as client:
|
||
r = await client.get(f"{self._api_base}/api/tags")
|
||
r.raise_for_status()
|
||
return [m["name"] for m in r.json().get("models", [])]
|
||
except Exception:
|
||
return []
|
||
|
||
async def select_best_model(self, intent: str = "chat") -> str:
|
||
"""Pick the preferred installed model for `intent` ('chat' or 'code'),
|
||
falling back to any installed model, then DEFAULT_CHAT_MODEL. Cached ~60s per
|
||
intent so rapid requests don't rebuild the model list each time.
|
||
"""
|
||
now = time.monotonic()
|
||
cached = self._model_cache.get(intent)
|
||
if cached and (now - cached[1]) < 60:
|
||
return cached[0]
|
||
|
||
models = await self.list_models()
|
||
pref = _MODEL_PREFERENCE.get(intent, _MODEL_PREFERENCE["chat"])
|
||
best = _preferred_model(models, pref) or (models[0] if models else DEFAULT_CHAT_MODEL)
|
||
|
||
self._model_cache[intent] = (best, now)
|
||
return best
|
||
|
||
def invalidate_model_cache(self):
|
||
"""Force next select_best_model() to re-query (e.g. after pull/delete)."""
|
||
self._model_cache = {}
|
||
|
||
async def get_model_layers(self, model: str) -> int | None:
|
||
"""Total offloadable layer count for `model` (repeating blocks + output layer).
|
||
|
||
Used to turn a CPU/GPU offload percentage into an Ollama `num_gpu`
|
||
value. Reads `<arch>.block_count` from /api/show and adds 1 for the
|
||
non-repeating output layer (Ollama reports e.g. 33 layers for a model
|
||
with block_count=32). Cached per-model since it never changes.
|
||
Returns None if the count can't be determined, so callers fall back
|
||
to Auto (no num_gpu override).
|
||
"""
|
||
if model in self._layer_cache:
|
||
return self._layer_cache[model]
|
||
try:
|
||
async with httpx.AsyncClient(timeout=10.0) as client:
|
||
r = await client.post(f"{self._api_base}/api/show", json={"model": model})
|
||
r.raise_for_status()
|
||
info = r.json().get("model_info", {}) or {}
|
||
block_count = next(
|
||
(v for k, v in info.items() if k.endswith(".block_count")), None
|
||
)
|
||
layers = int(block_count) + 1 if block_count is not None else None
|
||
except Exception as e:
|
||
_log.warning("get_model_layers(%s) failed: %s", model, e)
|
||
layers = None
|
||
if layers:
|
||
self._layer_cache[model] = layers
|
||
return layers
|
||
|
||
async def resolve_num_gpu(self, gpu_offload, model: str) -> int | None:
|
||
"""Convert a stored gpu_offload setting into an Ollama `num_gpu` value.
|
||
|
||
`gpu_offload` is -1 for Auto (returns None → no override, Ollama auto-fits)
|
||
or 0–100 for the percent of the model's layers to force onto the GPU
|
||
(0 = all CPU/RAM). Layer count is model-specific, resolved from the live
|
||
model. Returns None on anything unexpected so callers fall back to Auto.
|
||
"""
|
||
try:
|
||
pct = int(gpu_offload)
|
||
except (TypeError, ValueError):
|
||
return None
|
||
if pct < 0:
|
||
return None
|
||
pct = min(pct, 100)
|
||
layers = await self.get_model_layers(model)
|
||
if not layers:
|
||
return None
|
||
return max(0, round(pct / 100 * layers))
|
||
|
||
async def _chat_stream(self, messages: list, model: str, start: float,
|
||
temperature: float | None = None, num_gpu: int | None = None,
|
||
think: bool = False, num_ctx: int | None = None):
|
||
"""Async generator streaming tokens, then a final __meta__ stats sentinel."""
|
||
try:
|
||
body: dict = {"model": model, "messages": messages, "stream": True}
|
||
body["think"] = think # see chat(): reasoning off by default for speed
|
||
opts = _chat_options(temperature, num_gpu, num_ctx)
|
||
if opts:
|
||
body["options"] = opts
|
||
self._apply_keep_alive(body)
|
||
async with httpx.AsyncClient(timeout=300.0) as client:
|
||
async with client.stream(
|
||
"POST",
|
||
f"{self._api_base}/api/chat",
|
||
json=body,
|
||
) as response:
|
||
response.raise_for_status()
|
||
async for line in response.aiter_lines():
|
||
if not line.strip():
|
||
continue
|
||
try:
|
||
data = json.loads(line)
|
||
token = data.get("message", {}).get("content", "")
|
||
if token:
|
||
yield token
|
||
if data.get("done", False):
|
||
elapsed = time.perf_counter() - start
|
||
_log.info("chat stream completed model=%s elapsed=%.3fs", model, elapsed)
|
||
eval_count = data.get("eval_count", 0)
|
||
eval_ns = data.get("eval_duration", 0)
|
||
tokens_per_s = round(eval_count / (eval_ns / 1e9), 1) if eval_ns else 0
|
||
stats = json.dumps({
|
||
"model": model,
|
||
"tokens": eval_count,
|
||
"elapsed_s": round(elapsed, 2),
|
||
"tokens_per_s": tokens_per_s,
|
||
})
|
||
yield f"__meta__{stats}"
|
||
break
|
||
except Exception:
|
||
continue
|
||
except Exception as e:
|
||
elapsed = time.perf_counter() - start
|
||
_log.exception("chat stream error after %.3fs: %s", elapsed, e)
|
||
raise
|
||
|
||
|
||
def initialize_ollama() -> OllamaManager:
|
||
global _ollama_manager
|
||
|
||
if _ollama_manager is None:
|
||
manager = OllamaManager()
|
||
|
||
if not manager.is_running():
|
||
manager.start()
|
||
|
||
if not manager.is_running():
|
||
raise RuntimeError("Ollama API is not reachable after start().")
|
||
|
||
_ollama_manager = manager
|
||
|
||
return _ollama_manager
|
||
|
||
|
||
async def initialize_ollama_async() -> OllamaManager:
|
||
"""Async-safe initializer — uses asyncio.sleep so the event loop stays live."""
|
||
global _ollama_manager
|
||
|
||
if _ollama_manager is None:
|
||
manager = OllamaManager()
|
||
|
||
if not manager.is_running():
|
||
await manager.start_async()
|
||
|
||
if not manager.is_running():
|
||
raise RuntimeError("Ollama API is not reachable after start().")
|
||
|
||
_ollama_manager = manager
|
||
|
||
return _ollama_manager
|
||
|
||
|
||
def get_ollama_manager() -> OllamaManager:
|
||
global _ollama_manager
|
||
if _ollama_manager is None:
|
||
_ollama_manager = OllamaManager()
|
||
return _ollama_manager
|
||
|
||
|
||
def shutdown_ollama() -> None:
|
||
global _ollama_manager
|
||
if _ollama_manager is not None:
|
||
_ollama_manager.stop()
|
||
_ollama_manager = None |