Athena
2ebe93b4f7
Separate bind and client addresses, include Ollama's response body in HTTP failures, and strip inline <think> blocks from complete and streamed replies.
905 lines
36 KiB
Python
905 lines
36 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
|
||
|
||
|
||
_THINK_OPEN = "<think>"
|
||
_THINK_CLOSE = "</think>"
|
||
|
||
|
||
def _partial_tag_tail(text: str, tag: str) -> int:
|
||
"""Length of the longest suffix of `text` that could be the start of `tag`.
|
||
|
||
A tag can arrive split across stream chunks ("<thi" then "nk>"), so that much
|
||
of the tail has to be held back rather than emitted.
|
||
"""
|
||
for n in range(min(len(tag) - 1, len(text)), 0, -1):
|
||
if text.endswith(tag[:n]):
|
||
return n
|
||
return 0
|
||
|
||
|
||
class ThinkStripper:
|
||
"""Removes a reasoning model's <think> spans from a token stream.
|
||
|
||
Ollama routes reasoning into `message.thinking` only when asked to think.
|
||
We ask for `think: False` because the reasoning is pure latency here — but
|
||
deepseek-r1 and friends emit `<think>` inline in `content` anyway, so the
|
||
whole internal monologue reached the chat window, closing tags and all.
|
||
|
||
Two shapes show up in practice. A well-formed span is dropped whole. A
|
||
*stray* closing tag with no opening — which is what actually shipped — is at
|
||
least removed, so the user does not see a literal `</think>` in the reply.
|
||
Text already streamed before it cannot be recalled; `strip_think()` handles
|
||
that case properly for callers that have the complete message.
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
self._buf = ""
|
||
self._inside = False
|
||
|
||
def feed(self, chunk: str) -> str:
|
||
self._buf += chunk
|
||
out: list[str] = []
|
||
while True:
|
||
if self._inside:
|
||
end = self._buf.find(_THINK_CLOSE)
|
||
if end == -1:
|
||
keep = _partial_tag_tail(self._buf, _THINK_CLOSE)
|
||
self._buf = self._buf[len(self._buf) - keep:] if keep else ""
|
||
break
|
||
self._buf = self._buf[end + len(_THINK_CLOSE):]
|
||
self._inside = False
|
||
continue
|
||
|
||
start = self._buf.find(_THINK_OPEN)
|
||
stray = self._buf.find(_THINK_CLOSE)
|
||
# A stray close before any open: drop the tag, keep going.
|
||
if stray != -1 and (start == -1 or stray < start):
|
||
out.append(self._buf[:stray])
|
||
self._buf = self._buf[stray + len(_THINK_CLOSE):]
|
||
continue
|
||
if start == -1:
|
||
keep = max(
|
||
_partial_tag_tail(self._buf, _THINK_OPEN),
|
||
_partial_tag_tail(self._buf, _THINK_CLOSE),
|
||
)
|
||
if keep:
|
||
out.append(self._buf[:len(self._buf) - keep])
|
||
self._buf = self._buf[len(self._buf) - keep:]
|
||
else:
|
||
out.append(self._buf)
|
||
self._buf = ""
|
||
break
|
||
out.append(self._buf[:start])
|
||
self._buf = self._buf[start + len(_THINK_OPEN):]
|
||
self._inside = True
|
||
return "".join(out)
|
||
|
||
def flush(self) -> str:
|
||
"""Whatever is left once the stream ends. An unterminated <think> is
|
||
reasoning that never closed, so it is dropped rather than shown."""
|
||
rest = "" if self._inside else self._buf
|
||
self._buf = ""
|
||
return rest
|
||
|
||
|
||
def strip_think(text: str) -> str:
|
||
"""Remove reasoning from a complete message.
|
||
|
||
Unlike the streaming case this sees everything, so an unmatched closing tag
|
||
can be handled the way it was meant: every token before it was reasoning,
|
||
and the answer is what follows.
|
||
"""
|
||
if not text or _THINK_CLOSE not in text and _THINK_OPEN not in text:
|
||
return text
|
||
import re
|
||
cleaned = re.sub(r"<think>.*?</think>", "", text, flags=re.S)
|
||
if _THINK_CLOSE in cleaned: # stray close: the answer follows it
|
||
cleaned = cleaned.rsplit(_THINK_CLOSE, 1)[1]
|
||
cleaned = re.sub(r"<think>.*\Z", "", cleaned, flags=re.S) # never closed
|
||
return cleaned.strip()
|
||
|
||
|
||
async def _raise_for_ollama(r: httpx.Response) -> None:
|
||
"""raise_for_status(), but say what Ollama actually said.
|
||
|
||
Ollama answers every failure with {"error": "..."} — a model that isn't
|
||
pulled, a request too large for VRAM, a cloud model retired upstream — and
|
||
httpx's default message discards the body, leaving the user with:
|
||
|
||
Client error '410 Gone' for url 'http://127.0.0.1:11434/api/chat'
|
||
|
||
when the body held "glm-4.6 was retired at 2026-06-16". Same exception type
|
||
as before so existing handlers are unaffected; only the message improves.
|
||
"""
|
||
if r.is_success:
|
||
return
|
||
# A streamed response has no body loaded yet; reading it is what makes the
|
||
# error message available at all.
|
||
try:
|
||
await r.aread()
|
||
except Exception:
|
||
pass
|
||
detail = ""
|
||
try:
|
||
body = r.json()
|
||
if isinstance(body, dict):
|
||
detail = str(body.get("error") or "").strip()
|
||
except Exception:
|
||
detail = (r.text or "").strip()
|
||
if not detail:
|
||
r.raise_for_status() # nothing to add — keep httpx's wording
|
||
raise httpx.HTTPStatusError(
|
||
f"Ollama {r.status_code} from {r.request.url.path}: {detail[:400]}",
|
||
request=r.request,
|
||
response=r,
|
||
)
|
||
|
||
|
||
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()
|
||
# The bind value, not the connect value: a user who set 0.0.0.0 to reach
|
||
# Ollama from another machine must still get a server that listens there.
|
||
env["OLLAMA_HOST"] = settings.ollama_bind
|
||
# 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)
|
||
await _raise_for_ollama(r)
|
||
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:
|
||
await _raise_for_ollama(response)
|
||
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
|
||
await _raise_for_ollama(r)
|
||
message = r.json().get("message", {})
|
||
# A reasoning model puts its monologue in `content` even with
|
||
# think off, so strip it before anyone reads the answer.
|
||
if isinstance(message, dict) and message.get("content"):
|
||
message["content"] = strip_think(message["content"])
|
||
# 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,
|
||
)
|
||
await _raise_for_ollama(r)
|
||
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")
|
||
await _raise_for_ollama(r)
|
||
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})
|
||
await _raise_for_ollama(r)
|
||
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:
|
||
await _raise_for_ollama(response)
|
||
thinking = ThinkStripper()
|
||
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:
|
||
token = thinking.feed(token)
|
||
if token:
|
||
yield token
|
||
if data.get("done", False):
|
||
tail = thinking.flush() # held-back partial tag
|
||
if tail:
|
||
yield tail
|
||
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 |
GitNexus · part of the Nexus stack