e159b93b97* studio: improve GGUF tool calling accuracy and reliability - Add URL fetching to web_search tool so models can read full page content instead of only getting search snippets. Uses html2text for clean markdown conversion with regex fallback. - Inject current date and behavioral guidance (URL fetch workflow, no repeated queries, use code for data processing) into the tool-use system prompt. - Append error recovery nudge to tool results that indicate failure, helping small models avoid looping on the same broken call. - Strip leaked <tool_call> XML from assistant messages in conversation history and from the outgoing SSE stream. - Raise default max tool iterations from 10 to 25 across backend, model schema, and frontend defaults. - Increase _MAX_PAGE_CHARS from 4k to 16k so fetched pages contain enough content for the model to extract useful information. - Add "IMPORTANT: These are only short snippets" hint to search results so models know to fetch full pages when needed. Tested with Qwen3.5-4B-GGUF (UD-Q4_K_XL), 10 runs before/after: - XML leaks in responses: 10/10 -> 0/10 - URL fetch usage: 0 -> 4/10 runs - Runs producing actual correct answers: 0/10 -> 2/10 - Average tool calls per query: 5.5 -> 3.8 (more efficient) - Average response time: 12.3s -> 9.8s * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add tool calling benchmark results across model sizes and quants Tested 16 configurations (4 models x 2 quants x 2 KV cache types) with 10 runs each on NVIDIA B200. Best config: 27B UD-Q4_K_XL + bf16 KV -- 6/10 runs found all 4 correct songs, 0 XML leaks, 131s average response time. * Add duplicate tool-call detection and final-answer synthesis When the model repeats the exact same tool call (same name + arguments) twice in a row, skip execution and return a redirect message telling it to try a different approach. This prevents the 8x-repeated-query loops observed on 27B and 35B models. When the tool iteration cap (25) is reached, inject a "provide your final answer now" message before the final streaming pass. This lets the model synthesize a useful answer from everything it gathered instead of being silently cut off. Tested on Qwen3.5-27B UD-Q4_K_XL (10 runs): - Repeated query runs: 4/10 -> 2/10 - Cap hits: 1/10 -> 0/10 - All 4/4 accuracy: 5/10 -> 7/10 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix CodeQL alert: handle whitespace in script/style closing tags The regex fallback for HTML stripping did not match closing tags with whitespace before the angle bracket (e.g. </script >). Use \s* before > in both script and style patterns. * Address reviewer findings: SSRF, timeout crash, XML regex, dedup - SSRF: resolve hostname via getaddrinfo and reject private, loopback, link-local, multicast, and reserved addresses before fetching - Timeout: handle timeout=None (unlimited mode) in URL fetch path by defaulting to 60s instead of crashing on min(None, 60) - Download cap: read at most max_chars*4+1 bytes instead of the full response body before truncating - XML regex: match both <tool_call> and <function=...> markup in the history/stream cleanup (inference.py) - CodeQL: use [^>]* in closing script/style tags to handle any whitespace or attributes before > - Dedup: track whether each tool call failed so retries after transient errors are allowed; only block consecutive identical calls that both succeeded - Final-answer synthesis: guard on max_tool_iterations > 0 so callers who disable tools do not get a false "used all calls" turn * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix redirect SSRF, SSE streaming regression, dedup off-by-one - SSRF redirect bypass: disable auto-redirect in urllib, manually follow up to 5 hops with host validation at each step. Prevents public URLs from redirecting to loopback/private targets. - SSE streaming: track prev_text on the raw cumulative and strip XML from the delta only, so completed tool_call tags do not cause the cumulative to shrink and drop trailing real text. - Dedup off-by-one: check the immediately previous call (window=1) instead of requiring 2 matching history entries, so the second identical successful call is blocked rather than the third. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix redirect HTTPError handling and tighten error prefixes - Redirect fix: urllib raises HTTPError (not a normal response) when the redirect handler returns None. Catch HTTPError for 3xx codes and extract the Location header from the exception object. - Error prefixes: remove overly broad "No " prefix that matched "No results found." (a valid empty-search outcome, not an error). Replace with specific prefixes like "Blocked:", "No query provided", "Failed to resolve". This ensures empty search results are correctly classified as non-errors for duplicate-call tracking. * Fix SSE cross-chunk XML leaks, cleanup review findings - SSE streaming: sanitize the full cumulative text before diffing against the previous sanitized snapshot, so XML tags that span chunk boundaries are stripped correctly. The previous delta-based approach leaked split tags. - DRAINING fallback: use _strip_tool_markup() helper instead of a manual regex that only handled <tool_call> but not <function=...>. - Move hashlib import, _TOOL_XML_RE compile, and datetime import to module level per style guide. - Remove unused _hit_tool_cap variable. * Fix DNS rebinding, charset detection, HTTPError handling, dedup double-record - DNS rebinding: resolve hostname once via getaddrinfo, pin the returned IP, rewrite the URL to connect to the pinned IP with a Host header. Each redirect hop re-resolves and re-validates. Closes the TOCTOU window between validation and connection. - Charset: use resp.headers.get_content_charset() instead of hardcoding utf-8, so pages with other encodings decode correctly. - HTTPError: return descriptive "HTTP {code} {reason}" instead of re-raising into a generic "Search failed" message. - Dedup: remove redundant _record_tool_call in the duplicate branch; the single call at the end of the loop handles all cases. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
3096 lines
132 KiB
Python
3096 lines
132 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||
|
||
"""
|
||
llama-server inference backend for GGUF models.
|
||
|
||
Manages a llama-server subprocess and proxies chat completions
|
||
through its OpenAI-compatible /v1/chat/completions endpoint.
|
||
"""
|
||
|
||
import atexit
|
||
import contextlib
|
||
import hashlib
|
||
import json
|
||
import re
|
||
import struct
|
||
import structlog
|
||
from loggers import get_logger
|
||
import shutil
|
||
import socket
|
||
import subprocess
|
||
import threading
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Generator, Optional
|
||
|
||
import httpx
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
class LlamaCppBackend:
|
||
"""
|
||
Manages a llama-server subprocess for GGUF model inference.
|
||
|
||
Lifecycle:
|
||
1. load_model() — starts llama-server with the GGUF file
|
||
2. generate_chat_completion() — proxies to /v1/chat/completions, streams back
|
||
3. unload_model() — terminates llama-server subprocess
|
||
"""
|
||
|
||
def __init__(self):
|
||
self._process: Optional[subprocess.Popen] = None
|
||
self._port: Optional[int] = None
|
||
self._model_identifier: Optional[str] = None
|
||
self._gguf_path: Optional[str] = None
|
||
self._hf_repo: Optional[str] = None
|
||
self._hf_variant: Optional[str] = None
|
||
self._is_vision: bool = False
|
||
self._healthy = False
|
||
self._context_length: Optional[int] = None
|
||
self._effective_context_length: Optional[int] = None
|
||
self._max_context_length: Optional[int] = None
|
||
self._chat_template: Optional[str] = None
|
||
self._supports_reasoning: bool = False
|
||
self._reasoning_always_on: bool = False
|
||
self._supports_tools: bool = False
|
||
self._cache_type_kv: Optional[str] = None
|
||
self._reasoning_default: bool = True
|
||
# KV-cache estimation fields (populated by _read_gguf_metadata)
|
||
self._n_layers: Optional[int] = None
|
||
self._n_kv_heads: Optional[int] = None
|
||
self._n_heads: Optional[int] = None
|
||
self._embedding_length: Optional[int] = None
|
||
self._lock = threading.Lock()
|
||
self._stdout_lines: list[str] = []
|
||
self._stdout_thread: Optional[threading.Thread] = None
|
||
self._cancel_event = threading.Event()
|
||
self._api_key: Optional[str] = None
|
||
|
||
self._kill_orphaned_servers()
|
||
atexit.register(self._cleanup)
|
||
|
||
# ── Properties ────────────────────────────────────────────────
|
||
|
||
@property
|
||
def is_loaded(self) -> bool:
|
||
return self._process is not None and self._healthy
|
||
|
||
@property
|
||
def is_active(self) -> bool:
|
||
"""True if a llama-server process exists (loading or loaded)."""
|
||
return self._process is not None
|
||
|
||
@property
|
||
def base_url(self) -> str:
|
||
return f"http://127.0.0.1:{self._port}"
|
||
|
||
@property
|
||
def model_identifier(self) -> Optional[str]:
|
||
return self._model_identifier
|
||
|
||
@property
|
||
def is_vision(self) -> bool:
|
||
return self._is_vision
|
||
|
||
@property
|
||
def hf_variant(self) -> Optional[str]:
|
||
return self._hf_variant
|
||
|
||
@property
|
||
def context_length(self) -> Optional[int]:
|
||
"""Return the effective context length the server is running at."""
|
||
return self._effective_context_length or self._context_length
|
||
|
||
@property
|
||
def max_context_length(self) -> Optional[int]:
|
||
"""Return the maximum context currently available on this hardware."""
|
||
return self._max_context_length or self._context_length
|
||
|
||
@property
|
||
def chat_template(self) -> Optional[str]:
|
||
return self._chat_template
|
||
|
||
@property
|
||
def supports_reasoning(self) -> bool:
|
||
return self._supports_reasoning
|
||
|
||
@property
|
||
def reasoning_always_on(self) -> bool:
|
||
return self._reasoning_always_on
|
||
|
||
@property
|
||
def reasoning_default(self) -> bool:
|
||
return self._reasoning_default
|
||
|
||
@property
|
||
def supports_tools(self) -> bool:
|
||
return self._supports_tools
|
||
|
||
@property
|
||
def cache_type_kv(self) -> Optional[str]:
|
||
return self._cache_type_kv
|
||
|
||
# ── Binary discovery ──────────────────────────────────────────
|
||
|
||
@staticmethod
|
||
def _find_llama_server_binary() -> Optional[str]:
|
||
"""
|
||
Locate the llama-server binary.
|
||
|
||
Search order:
|
||
1. LLAMA_SERVER_PATH environment variable (direct path to binary)
|
||
1b. UNSLOTH_LLAMA_CPP_PATH env var (custom llama.cpp install dir)
|
||
2. ~/.unsloth/llama.cpp/llama-server (make build, root dir)
|
||
3. ~/.unsloth/llama.cpp/build/bin/llama-server (cmake build, Linux)
|
||
4. ~/.unsloth/llama.cpp/build/bin/Release/llama-server.exe (cmake build, Windows)
|
||
5. ./llama.cpp/llama-server (legacy: make build, root dir)
|
||
6. ./llama.cpp/build/bin/llama-server (legacy: cmake in-tree build)
|
||
7. llama-server on PATH (system install)
|
||
8. ./bin/llama-server (legacy: extracted binary)
|
||
"""
|
||
import os
|
||
import sys
|
||
|
||
binary_name = "llama-server.exe" if sys.platform == "win32" else "llama-server"
|
||
|
||
# 1. Env var — direct path to binary
|
||
env_path = os.environ.get("LLAMA_SERVER_PATH")
|
||
if env_path and Path(env_path).is_file():
|
||
return env_path
|
||
|
||
# 1b. UNSLOTH_LLAMA_CPP_PATH — custom llama.cpp install directory
|
||
custom_llama_cpp = os.environ.get("UNSLOTH_LLAMA_CPP_PATH")
|
||
if custom_llama_cpp:
|
||
custom_dir = Path(custom_llama_cpp)
|
||
# Root dir (make builds)
|
||
root_bin = custom_dir / binary_name
|
||
if root_bin.is_file():
|
||
return str(root_bin)
|
||
# build/bin/ (cmake builds on Linux)
|
||
cmake_bin = custom_dir / "build" / "bin" / binary_name
|
||
if cmake_bin.is_file():
|
||
return str(cmake_bin)
|
||
# build/bin/Release/ (cmake builds on Windows)
|
||
if sys.platform == "win32":
|
||
win_bin = custom_dir / "build" / "bin" / "Release" / binary_name
|
||
if win_bin.is_file():
|
||
return str(win_bin)
|
||
|
||
# 2–4. ~/.unsloth/llama.cpp (primary — setup.sh / setup.ps1 build here)
|
||
unsloth_home = Path.home() / ".unsloth" / "llama.cpp"
|
||
# Root dir (make builds copy binaries here)
|
||
home_root = unsloth_home / binary_name
|
||
if home_root.is_file():
|
||
return str(home_root)
|
||
# build/bin/ (cmake builds on Linux)
|
||
home_linux = unsloth_home / "build" / "bin" / binary_name
|
||
if home_linux.is_file():
|
||
return str(home_linux)
|
||
|
||
# 3. Windows MSVC build has Release subdir
|
||
if sys.platform == "win32":
|
||
home_win = unsloth_home / "build" / "bin" / "Release" / binary_name
|
||
if home_win.is_file():
|
||
return str(home_win)
|
||
|
||
# 5–6. Legacy: in-tree build (older setup.sh / setup.ps1 versions)
|
||
project_root = Path(__file__).resolve().parents[4]
|
||
# Root dir (make builds)
|
||
root_path = project_root / "llama.cpp" / binary_name
|
||
if root_path.is_file():
|
||
return str(root_path)
|
||
# build/bin/ (cmake builds)
|
||
build_path = project_root / "llama.cpp" / "build" / "bin" / binary_name
|
||
if build_path.is_file():
|
||
return str(build_path)
|
||
if sys.platform == "win32":
|
||
win_path = (
|
||
project_root / "llama.cpp" / "build" / "bin" / "Release" / binary_name
|
||
)
|
||
if win_path.is_file():
|
||
return str(win_path)
|
||
|
||
# 7. System PATH
|
||
system_path = shutil.which("llama-server")
|
||
if system_path:
|
||
return system_path
|
||
|
||
# 8. Legacy: extracted to bin/
|
||
bin_path = project_root / "bin" / binary_name
|
||
if bin_path.is_file():
|
||
return str(bin_path)
|
||
|
||
return None
|
||
|
||
# ── GPU allocation ────────────────────────────────────────────
|
||
|
||
@staticmethod
|
||
def _get_gguf_size_bytes(model_path: str) -> int:
|
||
"""Get total GGUF size in bytes, including split shards."""
|
||
import re
|
||
|
||
main = Path(model_path)
|
||
total = main.stat().st_size
|
||
|
||
# Check for split shards (e.g., model-00001-of-00003.gguf)
|
||
shard_pat = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$")
|
||
m = shard_pat.match(main.name)
|
||
if m:
|
||
prefix, _, num_total = m.group(1), m.group(2), m.group(3)
|
||
sibling_pat = re.compile(
|
||
r"^"
|
||
+ re.escape(prefix)
|
||
+ r"-\d{5}-of-"
|
||
+ re.escape(num_total)
|
||
+ r"\.gguf$"
|
||
)
|
||
for sibling in main.parent.iterdir():
|
||
if sibling != main and sibling_pat.match(sibling.name):
|
||
total += sibling.stat().st_size
|
||
|
||
return total
|
||
|
||
@staticmethod
|
||
def _get_gpu_free_memory() -> list[tuple[int, int]]:
|
||
"""Query free memory per GPU via nvidia-smi.
|
||
|
||
Returns list of (gpu_index, free_mib) sorted by index.
|
||
Respects CUDA_VISIBLE_DEVICES if set.
|
||
Returns empty list if nvidia-smi is not available.
|
||
"""
|
||
import os
|
||
|
||
try:
|
||
result = subprocess.run(
|
||
[
|
||
"nvidia-smi",
|
||
"--query-gpu=index,memory.free",
|
||
"--format=csv,noheader,nounits",
|
||
],
|
||
capture_output = True,
|
||
text = True,
|
||
timeout = 10,
|
||
)
|
||
if result.returncode != 0:
|
||
return []
|
||
|
||
# Parse which GPUs are allowed by existing CUDA_VISIBLE_DEVICES
|
||
allowed = None
|
||
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||
if cvd is not None and cvd.strip():
|
||
try:
|
||
allowed = set(int(x.strip()) for x in cvd.split(","))
|
||
except ValueError:
|
||
pass # Non-numeric (e.g., "GPU-uuid"), ignore filter
|
||
|
||
gpus = []
|
||
for line in result.stdout.strip().splitlines():
|
||
parts = line.split(",")
|
||
if len(parts) == 2:
|
||
idx = int(parts[0].strip())
|
||
free_mib = int(parts[1].strip())
|
||
if allowed is not None and idx not in allowed:
|
||
continue
|
||
gpus.append((idx, free_mib))
|
||
return gpus
|
||
except Exception as e:
|
||
logger.debug(f"Failed to query GPU free memory via nvidia-smi: {e}")
|
||
return []
|
||
|
||
@staticmethod
|
||
def _select_gpus(
|
||
model_size_bytes: int,
|
||
gpus: list[tuple[int, int]],
|
||
) -> tuple[Optional[list[int]], bool]:
|
||
"""Pick GPU(s) for a model based on estimated VRAM and free memory.
|
||
|
||
``model_size_bytes`` should include both model weights and estimated
|
||
KV cache. The 70% threshold provides headroom for compute buffers,
|
||
CUDA context, and other runtime overhead.
|
||
|
||
Returns (gpu_indices, use_fit):
|
||
- ([1], False) model fits on 1 GPU at 70% of free
|
||
- ([1, 2], False) model needs 2 GPUs
|
||
- (None, True) model too large, let --fit handle it
|
||
"""
|
||
if not gpus:
|
||
return None, True
|
||
|
||
model_size_mib = model_size_bytes / (1024 * 1024)
|
||
|
||
# Sort GPUs by free memory descending
|
||
ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
|
||
|
||
# Try fitting on 1 GPU (70% of free memory threshold)
|
||
if ranked[0][1] * 0.70 >= model_size_mib:
|
||
return [ranked[0][0]], False
|
||
|
||
# Try fitting on N GPUs (accumulate free memory from most-free)
|
||
cumulative = 0
|
||
selected = []
|
||
for idx, free_mib in ranked:
|
||
selected.append(idx)
|
||
cumulative += free_mib * 0.70
|
||
if cumulative >= model_size_mib:
|
||
return sorted(selected), False
|
||
|
||
# Model is too large even for all GPUs, let --fit handle it
|
||
logger.debug(
|
||
"Model does not fit in available GPU memory, falling back to --fit",
|
||
model_size_mib = round(model_size_mib, 2),
|
||
ranked_gpus = ranked,
|
||
)
|
||
return None, True
|
||
|
||
# ── KV cache VRAM estimation ─────────────────────────────────────
|
||
|
||
def _can_estimate_kv(self) -> bool:
|
||
"""True if we have enough GGUF metadata to estimate KV cache size."""
|
||
return (
|
||
self._n_layers is not None
|
||
and self._embedding_length is not None
|
||
and (self._n_kv_heads is not None or self._n_heads is not None)
|
||
)
|
||
|
||
def _estimate_kv_cache_bytes(
|
||
self, n_ctx: int, cache_type_kv: Optional[str] = None
|
||
) -> int:
|
||
"""Estimate KV cache VRAM for a given context length.
|
||
|
||
Returns 0 if metadata is insufficient for estimation.
|
||
"""
|
||
if not self._can_estimate_kv() or n_ctx <= 0:
|
||
return 0
|
||
|
||
n_layers = self._n_layers # type: ignore[assignment]
|
||
n_kv_heads = self._n_kv_heads or self._n_heads # type: ignore[assignment]
|
||
head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator]
|
||
|
||
# Bytes per element depends on KV cache quantization
|
||
bpe = {
|
||
"f32": 4.0,
|
||
"f16": 2.0,
|
||
"bf16": 2.0,
|
||
"q8_0": 34 / 32,
|
||
"q5_1": 0.75,
|
||
"q5_0": 0.6875,
|
||
"q4_1": 0.625,
|
||
"q4_0": 0.5625,
|
||
"iq4_nl": 0.5625,
|
||
}.get(cache_type_kv or "f16", 2.0)
|
||
|
||
# K + V caches: 2 * n_kv_heads * head_dim * n_layers * n_ctx * bpe
|
||
return int(2 * n_kv_heads * head_dim * n_layers * n_ctx * bpe)
|
||
|
||
def _fit_context_to_vram(
|
||
self,
|
||
requested_ctx: int,
|
||
available_mib: int,
|
||
model_size_bytes: int,
|
||
cache_type_kv: Optional[str] = None,
|
||
min_ctx: int = 4096,
|
||
) -> int:
|
||
"""Return the largest context length that fits in GPU VRAM.
|
||
|
||
Uses 70% of available VRAM as the budget (matching _select_gpus
|
||
threshold -- 30% reserved for compute buffers, CUDA context,
|
||
scratch space, flash-attn workspace, etc.).
|
||
If the model weights alone don't fit, returns min_ctx unchanged.
|
||
"""
|
||
if not self._can_estimate_kv():
|
||
logger.debug(
|
||
"Skipping context fit because KV cache metadata is unavailable",
|
||
requested_ctx = requested_ctx,
|
||
available_mib = available_mib,
|
||
)
|
||
return requested_ctx
|
||
|
||
budget_bytes = available_mib * 1024 * 1024 * 0.70
|
||
model_footprint = model_size_bytes
|
||
|
||
# Check if requested context already fits
|
||
kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv)
|
||
if model_footprint + kv <= budget_bytes:
|
||
return requested_ctx
|
||
|
||
# Model weights alone exceed budget -- can't help by reducing ctx.
|
||
# Return requested_ctx unchanged; --fit will handle VRAM management.
|
||
if model_footprint >= budget_bytes:
|
||
logger.debug(
|
||
"Model footprint exceeds GPU budget before KV cache",
|
||
requested_ctx = requested_ctx,
|
||
available_mib = available_mib,
|
||
model_size_gb = round(model_footprint / (1024**3), 2),
|
||
)
|
||
return requested_ctx
|
||
|
||
# Binary search for max context that fits
|
||
remaining = budget_bytes - model_footprint
|
||
effective_min = min(min_ctx, requested_ctx)
|
||
lo, hi = effective_min, requested_ctx
|
||
best = effective_min
|
||
while lo <= hi:
|
||
mid = (lo + hi) // 2
|
||
kv = self._estimate_kv_cache_bytes(mid, cache_type_kv)
|
||
if kv <= remaining:
|
||
best = mid
|
||
lo = mid + 1
|
||
else:
|
||
hi = mid - 1
|
||
|
||
# Round down to nearest 256 for alignment, but never exceed requested_ctx
|
||
best = (best // 256) * 256
|
||
best = max(effective_min, best)
|
||
best = min(best, requested_ctx)
|
||
return best
|
||
|
||
# ── Variant fallback ────────────────────────────────────────────
|
||
|
||
@staticmethod
|
||
def _find_smallest_fitting_variant(
|
||
hf_repo: str,
|
||
free_bytes: int,
|
||
hf_token: Optional[str] = None,
|
||
) -> Optional[tuple[str, int]]:
|
||
"""Find the smallest GGUF variant (including all shards) that fits.
|
||
|
||
Groups split shards by variant prefix and sums their sizes.
|
||
For example, UD-Q4_K_XL with 9 shards of 50 GB each = 450 GB total.
|
||
|
||
Returns (first_shard_filename, total_size_bytes) or None if nothing fits.
|
||
"""
|
||
import re
|
||
|
||
try:
|
||
from huggingface_hub import get_paths_info, list_repo_files
|
||
|
||
files = list_repo_files(hf_repo, token = hf_token)
|
||
gguf_files = [
|
||
f for f in files if f.endswith(".gguf") and "mmproj" not in f.lower()
|
||
]
|
||
if not gguf_files:
|
||
return None
|
||
|
||
# Get sizes for all GGUF files
|
||
path_infos = list(get_paths_info(hf_repo, gguf_files, token = hf_token))
|
||
size_map = {p.path: (p.size or 0) for p in path_infos}
|
||
|
||
# Group files by variant: shards share a prefix before -NNNNN-of-NNNNN
|
||
shard_pat = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$")
|
||
variants: dict[str, list[str]] = {}
|
||
for f in gguf_files:
|
||
m = shard_pat.match(f)
|
||
key = m.group(1) if m else f
|
||
variants.setdefault(key, []).append(f)
|
||
|
||
# Sum shard sizes per variant, track the first shard (for download)
|
||
variant_sizes: list[tuple[str, int, list[str]]] = []
|
||
for key, shard_files in variants.items():
|
||
total = sum(size_map.get(f, 0) for f in shard_files)
|
||
first = sorted(shard_files)[0]
|
||
variant_sizes.append((first, total, shard_files))
|
||
|
||
# Sort by total size ascending and pick the smallest that fits
|
||
variant_sizes.sort(key = lambda x: x[1])
|
||
for first_file, total_size, _ in variant_sizes:
|
||
if total_size > 0 and total_size <= free_bytes:
|
||
return first_file, total_size
|
||
|
||
return None
|
||
except Exception:
|
||
return None
|
||
|
||
# ── Port allocation ───────────────────────────────────────────
|
||
|
||
@staticmethod
|
||
def _find_free_port() -> int:
|
||
"""Find an available TCP port."""
|
||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||
s.bind(("127.0.0.1", 0))
|
||
return s.getsockname()[1]
|
||
|
||
# ── Stdout drain (prevents pipe deadlock on Windows) ─────────
|
||
|
||
def _drain_stdout(self):
|
||
"""
|
||
Read lines from the subprocess stdout in a background thread.
|
||
|
||
This prevents a pipe-buffer deadlock on Windows where the default
|
||
pipe buffer is only ~4 KB. Without draining, llama-server blocks
|
||
on writes and never becomes healthy.
|
||
"""
|
||
try:
|
||
for line in self._process.stdout:
|
||
line = line.rstrip()
|
||
if line:
|
||
self._stdout_lines.append(line)
|
||
logger.debug(f"[llama-server] {line}")
|
||
except (ValueError, OSError):
|
||
# Pipe closed — process is terminating
|
||
pass
|
||
|
||
# GGUF KV type sizes for fast skipping
|
||
_GGUF_TYPE_SIZE = {
|
||
0: 1,
|
||
1: 1,
|
||
2: 2,
|
||
3: 2,
|
||
4: 4,
|
||
5: 4,
|
||
6: 4,
|
||
7: 1,
|
||
10: 8,
|
||
11: 8,
|
||
12: 8,
|
||
}
|
||
|
||
@staticmethod
|
||
def _gguf_skip_value(f, vtype: int) -> None:
|
||
"""Skip a GGUF KV value without reading it."""
|
||
sz = LlamaCppBackend._GGUF_TYPE_SIZE.get(vtype)
|
||
if sz is not None:
|
||
f.seek(sz, 1)
|
||
elif vtype == 8: # STRING
|
||
slen = struct.unpack("<Q", f.read(8))[0]
|
||
f.seek(slen, 1)
|
||
elif vtype == 9: # ARRAY
|
||
atype = struct.unpack("<I", f.read(4))[0]
|
||
alen = struct.unpack("<Q", f.read(8))[0]
|
||
elem_sz = LlamaCppBackend._GGUF_TYPE_SIZE.get(atype)
|
||
if elem_sz is not None:
|
||
f.seek(elem_sz * alen, 1)
|
||
elif atype == 8:
|
||
for _ in range(alen):
|
||
slen = struct.unpack("<Q", f.read(8))[0]
|
||
f.seek(slen, 1)
|
||
else:
|
||
for _ in range(alen):
|
||
LlamaCppBackend._gguf_skip_value(f, atype)
|
||
|
||
def _read_gguf_metadata(self, gguf_path: str) -> None:
|
||
"""Read context_length, architecture params, and chat_template from a GGUF header.
|
||
|
||
Parses only the KV pairs we need (~30ms even for multi-GB files).
|
||
For split GGUFs, metadata is always in shard 1.
|
||
"""
|
||
# Reset metadata from any previously loaded model so stale flags
|
||
# (eg _supports_reasoning) do not carry over when switching models.
|
||
self._context_length = None
|
||
self._chat_template = None
|
||
self._supports_reasoning = False
|
||
self._reasoning_always_on = False
|
||
self._supports_tools = False
|
||
self._n_layers = None
|
||
self._n_kv_heads = None
|
||
self._n_heads = None
|
||
self._embedding_length = None
|
||
|
||
try:
|
||
WANTED = {"general.architecture", "tokenizer.chat_template"}
|
||
# Additional arch-specific keys are added dynamically once
|
||
# we know the architecture name.
|
||
arch_keys: dict[str, str] = {} # gguf_key -> attribute name
|
||
arch = None
|
||
|
||
with open(gguf_path, "rb") as f:
|
||
magic = struct.unpack("<I", f.read(4))[0]
|
||
if magic != 0x46554747: # b"GGUF" as little-endian u32
|
||
return
|
||
_version = struct.unpack("<I", f.read(4))[0]
|
||
_tensor_count, kv_count = struct.unpack("<QQ", f.read(16))
|
||
|
||
for _ in range(kv_count):
|
||
key_len = struct.unpack("<Q", f.read(8))[0]
|
||
key = f.read(key_len).decode("utf-8")
|
||
vtype = struct.unpack("<I", f.read(4))[0]
|
||
|
||
if key in WANTED or key in arch_keys:
|
||
# Read this value
|
||
if vtype == 8: # STRING
|
||
slen = struct.unpack("<Q", f.read(8))[0]
|
||
val_s = f.read(slen).decode("utf-8")
|
||
if key == "general.architecture":
|
||
arch = val_s
|
||
# Register arch-specific keys to look for
|
||
arch_keys = {
|
||
f"{arch}.context_length": "context_length",
|
||
f"{arch}.block_count": "n_layers",
|
||
f"{arch}.attention.head_count_kv": "n_kv_heads",
|
||
f"{arch}.attention.head_count": "n_heads",
|
||
f"{arch}.embedding_length": "embedding_length",
|
||
}
|
||
elif key == "tokenizer.chat_template":
|
||
self._chat_template = val_s
|
||
elif vtype in (4, 10): # UINT32 or UINT64
|
||
val_i = (
|
||
struct.unpack("<I", f.read(4))[0]
|
||
if vtype == 4
|
||
else struct.unpack("<Q", f.read(8))[0]
|
||
)
|
||
attr = arch_keys.get(key)
|
||
if attr:
|
||
setattr(self, f"_{attr}", val_i)
|
||
else:
|
||
self._gguf_skip_value(f, vtype)
|
||
else:
|
||
self._gguf_skip_value(f, vtype)
|
||
|
||
if self._context_length:
|
||
logger.info(f"GGUF metadata: context_length={self._context_length}")
|
||
if self._chat_template:
|
||
logger.info(
|
||
f"GGUF metadata: chat_template={len(self._chat_template)} chars"
|
||
)
|
||
# Detect thinking/reasoning support from chat template
|
||
tpl = self._chat_template
|
||
if "enable_thinking" in tpl:
|
||
self._supports_reasoning = True
|
||
logger.info(
|
||
"GGUF metadata: model supports reasoning (enable_thinking)"
|
||
)
|
||
elif "thinking" in tpl:
|
||
# DeepSeek uses 'thinking' instead of 'enable_thinking'
|
||
normalized_id = (self._model_identifier or "").lower()
|
||
if "deepseek" in normalized_id:
|
||
self._supports_reasoning = True
|
||
logger.info(
|
||
"GGUF metadata: model supports reasoning (DeepSeek thinking)"
|
||
)
|
||
# Models with hardcoded <think> tags or reasoning_content
|
||
# in their chat template always produce thinking output
|
||
# (no toggle to disable it).
|
||
if not self._supports_reasoning:
|
||
if (
|
||
"<think>" in tpl
|
||
and "</think>" in tpl
|
||
or "reasoning_content" in tpl
|
||
):
|
||
self._supports_reasoning = True
|
||
self._reasoning_always_on = True
|
||
logger.info(
|
||
"GGUF metadata: model always reasons (<think> tags in template)"
|
||
)
|
||
# Detect tool calling support from chat template
|
||
tool_markers = [
|
||
"{%- if tools %}",
|
||
"{% if tools %}",
|
||
'"role" == "tool"',
|
||
"'role' == 'tool'",
|
||
'message.role == "tool"',
|
||
"message.role == 'tool'",
|
||
]
|
||
if any(marker in tpl for marker in tool_markers):
|
||
self._supports_tools = True
|
||
logger.info("GGUF metadata: model supports tool calling")
|
||
except Exception as e:
|
||
logger.warning(f"Failed to read GGUF metadata: {e}")
|
||
|
||
# ── HF download (no lock held) ───────────────────────────────
|
||
|
||
def _download_gguf(
|
||
self,
|
||
*,
|
||
hf_repo: str,
|
||
hf_variant: Optional[str] = None,
|
||
hf_token: Optional[str] = None,
|
||
) -> str:
|
||
"""Download GGUF file(s) from HuggingFace. Returns local path.
|
||
|
||
Runs WITHOUT self._lock so that unload_model() can set
|
||
_cancel_event at any time. Checks _cancel_event between
|
||
each shard download.
|
||
"""
|
||
try:
|
||
from huggingface_hub import hf_hub_download
|
||
except ImportError:
|
||
raise RuntimeError(
|
||
"huggingface_hub is required for HF model loading. "
|
||
"Install it with: pip install huggingface_hub"
|
||
)
|
||
|
||
# Determine the filename from the variant
|
||
gguf_filename = None
|
||
gguf_extra_shards: list[str] = []
|
||
if hf_variant:
|
||
try:
|
||
import re
|
||
from huggingface_hub import list_repo_files
|
||
|
||
files = list_repo_files(hf_repo, token = hf_token)
|
||
variant_lower = hf_variant.lower()
|
||
boundary = re.compile(
|
||
r"(?<![a-zA-Z0-9])" + re.escape(variant_lower) + r"(?![a-zA-Z0-9])"
|
||
)
|
||
gguf_files = sorted(
|
||
f
|
||
for f in files
|
||
if f.endswith(".gguf") and boundary.search(f.lower())
|
||
)
|
||
if gguf_files:
|
||
gguf_filename = gguf_files[0]
|
||
shard_pat = re.compile(r"^(.*)-\d{5}-of-(\d{5})\.gguf$")
|
||
m = shard_pat.match(gguf_filename)
|
||
if m:
|
||
prefix = m.group(1)
|
||
total = m.group(2)
|
||
sibling_pat = re.compile(
|
||
r"^"
|
||
+ re.escape(prefix)
|
||
+ r"-\d{5}-of-"
|
||
+ re.escape(total)
|
||
+ r"\.gguf$"
|
||
)
|
||
gguf_extra_shards = [
|
||
f for f in gguf_files[1:] if sibling_pat.match(f)
|
||
]
|
||
except Exception as e:
|
||
logger.warning(f"Could not list repo files: {e}")
|
||
|
||
if not gguf_filename:
|
||
repo_name = hf_repo.split("/")[-1].replace("-GGUF", "")
|
||
gguf_filename = f"{repo_name}-{hf_variant}.gguf"
|
||
|
||
# Check disk space and fall back to a smaller variant if needed
|
||
all_gguf_files = [gguf_filename] + gguf_extra_shards
|
||
try:
|
||
import os
|
||
|
||
from huggingface_hub import get_paths_info
|
||
|
||
path_infos = list(get_paths_info(hf_repo, all_gguf_files, token = hf_token))
|
||
total_download_bytes = sum((p.size or 0) for p in path_infos)
|
||
|
||
if total_download_bytes > 0:
|
||
cache_dir = os.environ.get(
|
||
"HF_HUB_CACHE",
|
||
str(Path.home() / ".cache" / "huggingface" / "hub"),
|
||
)
|
||
Path(cache_dir).mkdir(parents = True, exist_ok = True)
|
||
free_bytes = shutil.disk_usage(cache_dir).free
|
||
|
||
total_gb = total_download_bytes / (1024**3)
|
||
free_gb = free_bytes / (1024**3)
|
||
|
||
logger.info(
|
||
f"GGUF download: {total_gb:.1f} GB needed, "
|
||
f"{free_gb:.1f} GB free on disk"
|
||
)
|
||
|
||
if total_download_bytes > free_bytes:
|
||
smaller = self._find_smallest_fitting_variant(
|
||
hf_repo,
|
||
free_bytes,
|
||
hf_token,
|
||
)
|
||
if smaller:
|
||
fallback_file, fallback_size = smaller
|
||
logger.info(
|
||
f"Selected variant too large ({total_gb:.1f} GB), "
|
||
f"falling back to {fallback_file} ({fallback_size / (1024**3):.1f} GB)"
|
||
)
|
||
gguf_filename = fallback_file
|
||
import re as _re
|
||
|
||
_shard_pat = _re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$")
|
||
_m = _shard_pat.match(gguf_filename)
|
||
_prefix = _m.group(1) if _m else None
|
||
if _prefix:
|
||
gguf_extra_shards = sorted(
|
||
f
|
||
for f in all_gguf_files
|
||
if f.startswith(_prefix)
|
||
and f != gguf_filename
|
||
and "mmproj" not in f.lower()
|
||
)
|
||
else:
|
||
gguf_extra_shards = []
|
||
else:
|
||
raise RuntimeError(
|
||
f"Not enough disk space to download any variant. "
|
||
f"Only {free_gb:.1f} GB free in {cache_dir}"
|
||
)
|
||
except RuntimeError:
|
||
raise
|
||
except Exception as e:
|
||
logger.warning(f"Could not check disk space: {e}")
|
||
|
||
gguf_label = f"{hf_repo}/{gguf_filename}" + (
|
||
f" (+{len(gguf_extra_shards)} shards)" if gguf_extra_shards else ""
|
||
)
|
||
logger.info(f"Resolving GGUF: {gguf_label}")
|
||
try:
|
||
if self._cancel_event.is_set():
|
||
raise RuntimeError("Cancelled")
|
||
dl_start = time.monotonic()
|
||
local_path = hf_hub_download(
|
||
repo_id = hf_repo,
|
||
filename = gguf_filename,
|
||
token = hf_token,
|
||
)
|
||
for shard in gguf_extra_shards:
|
||
if self._cancel_event.is_set():
|
||
raise RuntimeError("Cancelled")
|
||
logger.info(f"Resolving GGUF shard: {shard}")
|
||
hf_hub_download(
|
||
repo_id = hf_repo,
|
||
filename = shard,
|
||
token = hf_token,
|
||
)
|
||
except RuntimeError as e:
|
||
if "Cancelled" in str(e):
|
||
raise
|
||
raise RuntimeError(
|
||
f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}"
|
||
)
|
||
except Exception as e:
|
||
raise RuntimeError(
|
||
f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}"
|
||
)
|
||
|
||
dl_elapsed = time.monotonic() - dl_start
|
||
if dl_elapsed < 2.0:
|
||
logger.info(f"GGUF resolved from cache: {local_path}")
|
||
else:
|
||
logger.info(f"GGUF downloaded in {dl_elapsed:.1f}s: {local_path}")
|
||
return local_path
|
||
|
||
def _download_mmproj(
|
||
self,
|
||
*,
|
||
hf_repo: str,
|
||
hf_token: Optional[str] = None,
|
||
) -> Optional[str]:
|
||
"""Download the mmproj (vision projection) file from a GGUF repo.
|
||
|
||
Prefers mmproj-F16.gguf, falls back to any mmproj*.gguf file.
|
||
Returns the local path, or None if no mmproj file exists.
|
||
"""
|
||
try:
|
||
from huggingface_hub import hf_hub_download, list_repo_files
|
||
|
||
files = list_repo_files(hf_repo, token = hf_token)
|
||
mmproj_files = sorted(
|
||
f for f in files if f.endswith(".gguf") and "mmproj" in f.lower()
|
||
)
|
||
if not mmproj_files:
|
||
return None
|
||
|
||
# Prefer F16 variant
|
||
target = None
|
||
for f in mmproj_files:
|
||
if "f16" in f.lower():
|
||
target = f
|
||
break
|
||
if target is None:
|
||
target = mmproj_files[0]
|
||
|
||
logger.info(f"Downloading mmproj: {hf_repo}/{target}")
|
||
local_path = hf_hub_download(
|
||
repo_id = hf_repo,
|
||
filename = target,
|
||
token = hf_token,
|
||
)
|
||
return local_path
|
||
except Exception as e:
|
||
logger.warning(f"Could not download mmproj: {e}")
|
||
return None
|
||
|
||
# ── Lifecycle ─────────────────────────────────────────────────
|
||
|
||
def load_model(
|
||
self,
|
||
*,
|
||
# Local mode: pass a path to a .gguf file
|
||
gguf_path: Optional[str] = None,
|
||
# Vision projection (mmproj) for local vision models
|
||
mmproj_path: Optional[str] = None,
|
||
# HF mode: let llama-server download via -hf "repo:quant"
|
||
hf_repo: Optional[str] = None,
|
||
hf_variant: Optional[str] = None,
|
||
hf_token: Optional[str] = None,
|
||
# Common
|
||
model_identifier: str,
|
||
is_vision: bool = False,
|
||
n_ctx: int = 4096,
|
||
chat_template_override: Optional[str] = None,
|
||
cache_type_kv: Optional[str] = None,
|
||
n_threads: Optional[int] = None,
|
||
n_gpu_layers: Optional[int] = None, # Accepted for caller compat, unused
|
||
) -> bool:
|
||
"""
|
||
Start llama-server with a GGUF model.
|
||
|
||
Two modes:
|
||
- Local: ``gguf_path="/path/to/model.gguf"`` → uses ``-m``
|
||
- HF: ``hf_repo="unsloth/gemma-3-4b-it-GGUF", hf_variant="Q4_K_M"`` → uses ``-hf``
|
||
|
||
In HF mode, llama-server handles downloading, caching, and
|
||
auto-loading mmproj files for vision models.
|
||
|
||
Returns True if server started and health check passed.
|
||
"""
|
||
self._cancel_event.clear()
|
||
|
||
# ── Phase 1: kill old process (under lock, fast) ──────────
|
||
with self._lock:
|
||
self._kill_process()
|
||
|
||
binary = self._find_llama_server_binary()
|
||
if not binary:
|
||
raise RuntimeError(
|
||
"llama-server binary not found. "
|
||
"Run setup.sh to build it, install llama.cpp, "
|
||
"or set LLAMA_SERVER_PATH environment variable."
|
||
)
|
||
|
||
# ── Phase 2: download (NO lock held, so cancel can proceed) ──
|
||
if hf_repo:
|
||
model_path = self._download_gguf(
|
||
hf_repo = hf_repo,
|
||
hf_variant = hf_variant,
|
||
hf_token = hf_token,
|
||
)
|
||
# Auto-download mmproj for vision models
|
||
if is_vision and not mmproj_path:
|
||
mmproj_path = self._download_mmproj(
|
||
hf_repo = hf_repo,
|
||
hf_token = hf_token,
|
||
)
|
||
elif gguf_path:
|
||
if not Path(gguf_path).is_file():
|
||
raise FileNotFoundError(f"GGUF file not found: {gguf_path}")
|
||
model_path = gguf_path
|
||
else:
|
||
raise ValueError("Either gguf_path or hf_repo must be provided")
|
||
|
||
# Set identifier early so _read_gguf_metadata can use it for DeepSeek detection
|
||
self._model_identifier = model_identifier
|
||
|
||
# Read GGUF metadata (context_length, chat_template) -- fast, header only
|
||
self._read_gguf_metadata(model_path)
|
||
|
||
# Check cancel after download
|
||
if self._cancel_event.is_set():
|
||
logger.info("Load cancelled after download phase")
|
||
return False
|
||
|
||
# ── Phase 3: start llama-server (under lock) ──────────────
|
||
with self._lock:
|
||
# Re-check cancel inside lock
|
||
if self._cancel_event.is_set():
|
||
logger.info("Load cancelled before server start")
|
||
return False
|
||
|
||
self._port = self._find_free_port()
|
||
|
||
# Select GPU(s) based on model size + estimated KV cache.
|
||
# Seed safe defaults before GPU probing so the except path
|
||
# still has valid state to publish.
|
||
effective_ctx = n_ctx if n_ctx > 0 else (self._context_length or 0)
|
||
max_available_ctx = self._context_length or effective_ctx
|
||
try:
|
||
model_size = self._get_gguf_size_bytes(model_path)
|
||
gpus = self._get_gpu_free_memory()
|
||
|
||
# Resolve effective context: 0 means let llama-server use the
|
||
# model's native length. Only expand to a known native length
|
||
# if metadata is available; otherwise preserve 0 as a sentinel.
|
||
if n_ctx > 0:
|
||
effective_ctx = n_ctx
|
||
elif self._context_length is not None:
|
||
effective_ctx = self._context_length
|
||
else:
|
||
effective_ctx = 0
|
||
original_ctx = effective_ctx
|
||
# Default UI ceiling to the model's native context length.
|
||
# GPU/VRAM-fit logic below may shrink this if hardware is limited.
|
||
max_available_ctx = self._context_length or effective_ctx
|
||
|
||
# Auto-cap context to fit in GPU VRAM and select GPUs.
|
||
#
|
||
# Two policies depending on whether the user set n_ctx:
|
||
#
|
||
# Explicit n_ctx (user chose a context length):
|
||
# Honor it. Try the full requested context with _select_gpus
|
||
# (which uses as many GPUs as needed). Only cap if it doesn't
|
||
# fit on any GPU combination.
|
||
#
|
||
# Auto n_ctx=0 (model's native context):
|
||
# Prefer fewer GPUs with reduced context over more GPUs,
|
||
# since multi-GPU is slower and the user didn't ask for a
|
||
# specific context length.
|
||
gpu_indices, use_fit = None, True
|
||
explicit_ctx = n_ctx > 0
|
||
|
||
if gpus and self._can_estimate_kv() and effective_ctx > 0:
|
||
# Compute the largest hardware-aware cap from the model's
|
||
# native context across all usable GPU subsets (for UI
|
||
# bounds), independent of the currently requested context.
|
||
native_ctx_for_cap = self._context_length or effective_ctx
|
||
if native_ctx_for_cap > 0:
|
||
ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True)
|
||
best_cap = 0
|
||
for n_gpus in range(1, len(ranked_for_cap) + 1):
|
||
subset = ranked_for_cap[:n_gpus]
|
||
pool_mib = sum(free for _, free in subset)
|
||
capped = self._fit_context_to_vram(
|
||
native_ctx_for_cap,
|
||
pool_mib,
|
||
model_size,
|
||
cache_type_kv,
|
||
)
|
||
kv = self._estimate_kv_cache_bytes(capped, cache_type_kv)
|
||
total_mib = (model_size + kv) / (1024 * 1024)
|
||
if total_mib <= pool_mib * 0.70:
|
||
best_cap = max(best_cap, capped)
|
||
if best_cap > 0:
|
||
max_available_ctx = best_cap
|
||
|
||
if explicit_ctx:
|
||
# Try to honor the user's requested context exactly.
|
||
requested_total = model_size + self._estimate_kv_cache_bytes(
|
||
effective_ctx, cache_type_kv
|
||
)
|
||
gpu_indices, use_fit = self._select_gpus(requested_total, gpus)
|
||
|
||
# Full context doesn't fit anywhere -- cap it on the
|
||
# best GPU subset we can find (fewest GPUs first).
|
||
if use_fit:
|
||
ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
|
||
for n_gpus in range(1, len(ranked) + 1):
|
||
subset = ranked[:n_gpus]
|
||
pool_mib = sum(free for _, free in subset)
|
||
capped = self._fit_context_to_vram(
|
||
effective_ctx,
|
||
pool_mib,
|
||
model_size,
|
||
cache_type_kv,
|
||
)
|
||
kv = self._estimate_kv_cache_bytes(
|
||
capped, cache_type_kv
|
||
)
|
||
total_mib = (model_size + kv) / (1024 * 1024)
|
||
if total_mib <= pool_mib * 0.70:
|
||
effective_ctx = capped
|
||
gpu_indices = sorted(idx for idx, _ in subset)
|
||
use_fit = False
|
||
break
|
||
else:
|
||
# Auto context: prefer fewer GPUs, cap context to fit.
|
||
ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
|
||
for n_gpus in range(1, len(ranked) + 1):
|
||
subset = ranked[:n_gpus]
|
||
pool_mib = sum(free for _, free in subset)
|
||
capped = self._fit_context_to_vram(
|
||
effective_ctx,
|
||
pool_mib,
|
||
model_size,
|
||
cache_type_kv,
|
||
)
|
||
kv = self._estimate_kv_cache_bytes(capped, cache_type_kv)
|
||
total_mib = (model_size + kv) / (1024 * 1024)
|
||
if total_mib <= pool_mib * 0.70:
|
||
effective_ctx = capped
|
||
gpu_indices = sorted(idx for idx, _ in subset)
|
||
use_fit = False
|
||
break
|
||
|
||
elif gpus:
|
||
# Can't estimate KV -- fall back to file-size-only check.
|
||
# Without KV estimation we cannot prove a hardware cap, so
|
||
# keep the ceiling at the native context (already the default).
|
||
logger.debug(
|
||
"Falling back to file-size-only GPU selection",
|
||
model_size_gb = round(model_size / (1024**3), 2),
|
||
)
|
||
gpu_indices, use_fit = self._select_gpus(model_size, gpus)
|
||
|
||
if effective_ctx < original_ctx:
|
||
kv_est = self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv)
|
||
logger.info(
|
||
f"Context auto-reduced: {original_ctx} -> {effective_ctx} "
|
||
f"(model: {model_size / (1024**3):.1f} GB, "
|
||
f"est. KV cache: {kv_est / (1024**3):.1f} GB)"
|
||
)
|
||
|
||
kv_cache_bytes = self._estimate_kv_cache_bytes(
|
||
effective_ctx, cache_type_kv
|
||
)
|
||
logger.info(
|
||
f"GGUF size: {model_size / (1024**3):.1f} GB, "
|
||
f"est. KV cache: {kv_cache_bytes / (1024**3):.1f} GB, "
|
||
f"context: {effective_ctx}, "
|
||
f"GPUs free: {gpus}, selected: {gpu_indices}, fit: {use_fit}"
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"GPU selection failed ({e}), using --fit on")
|
||
gpu_indices, use_fit = None, True
|
||
effective_ctx = n_ctx # fall back to original
|
||
|
||
cmd = [
|
||
binary,
|
||
"-m",
|
||
model_path,
|
||
"--port",
|
||
str(self._port),
|
||
"-c",
|
||
str(effective_ctx) if effective_ctx > 0 else "0",
|
||
"--parallel",
|
||
"1", # Single-user studio, saves VRAM
|
||
"--flash-attn",
|
||
"on", # Force flash attention for speed
|
||
]
|
||
|
||
if use_fit:
|
||
cmd.extend(["--fit", "on"])
|
||
elif gpu_indices is not None:
|
||
# Model fits on selected GPU(s) -- offload all layers
|
||
cmd.extend(["-ngl", "-1"])
|
||
|
||
if n_threads is not None:
|
||
cmd.extend(["--threads", str(n_threads)])
|
||
|
||
# Always enable Jinja chat template rendering for proper template support
|
||
cmd.extend(["--jinja"])
|
||
|
||
# KV cache data type
|
||
_valid_cache_types = {
|
||
"f16",
|
||
"bf16",
|
||
"q8_0",
|
||
"q4_0",
|
||
"q4_1",
|
||
"q5_0",
|
||
"q5_1",
|
||
"iq4_nl",
|
||
"f32",
|
||
}
|
||
if cache_type_kv and cache_type_kv in _valid_cache_types:
|
||
cmd.extend(
|
||
["--cache-type-k", cache_type_kv, "--cache-type-v", cache_type_kv]
|
||
)
|
||
self._cache_type_kv = cache_type_kv
|
||
logger.info(f"KV cache type: {cache_type_kv}")
|
||
else:
|
||
self._cache_type_kv = None
|
||
|
||
# Apply custom chat template override if provided
|
||
if chat_template_override:
|
||
import tempfile
|
||
|
||
self._chat_template_file = tempfile.NamedTemporaryFile(
|
||
mode = "w",
|
||
suffix = ".jinja",
|
||
delete = False,
|
||
prefix = "unsloth_chat_template_",
|
||
)
|
||
self._chat_template_file.write(chat_template_override)
|
||
self._chat_template_file.close()
|
||
cmd.extend(["--chat-template-file", self._chat_template_file.name])
|
||
logger.info(
|
||
f"Using custom chat template file: {self._chat_template_file.name}"
|
||
)
|
||
|
||
# For reasoning models, set default thinking mode.
|
||
# Qwen3.5 models below 9B (0.8B, 2B, 4B) disable thinking by default.
|
||
# Only 9B and larger enable thinking.
|
||
if self._supports_reasoning:
|
||
import re
|
||
|
||
thinking_default = True
|
||
mid = (model_identifier or "").lower()
|
||
if "qwen3.5" in mid:
|
||
# Extract size like "0.8b", "4b", "35b" etc.
|
||
size_match = re.search(r"(\d+\.?\d*)\s*b", mid)
|
||
if size_match:
|
||
size_val = float(size_match.group(1))
|
||
if size_val < 9:
|
||
thinking_default = False
|
||
self._reasoning_default = thinking_default
|
||
cmd.extend(
|
||
[
|
||
"--chat-template-kwargs",
|
||
json.dumps({"enable_thinking": thinking_default}),
|
||
]
|
||
)
|
||
logger.info(
|
||
f"Reasoning model: enable_thinking={thinking_default} by default"
|
||
)
|
||
|
||
if mmproj_path:
|
||
if not Path(mmproj_path).is_file():
|
||
logger.warning(f"mmproj file not found: {mmproj_path}")
|
||
else:
|
||
cmd.extend(["--mmproj", mmproj_path])
|
||
logger.info(f"Using mmproj for vision: {mmproj_path}")
|
||
|
||
# Option C: add --api-key for direct client access when enabled
|
||
import os as _os
|
||
import secrets as _secrets
|
||
|
||
if _os.getenv("UNSLOTH_DIRECT_STREAM", "0") == "1":
|
||
self._api_key = _secrets.token_urlsafe(32)
|
||
cmd.extend(["--api-key", self._api_key])
|
||
logger.info("llama-server started with --api-key for direct streaming")
|
||
else:
|
||
self._api_key = None
|
||
|
||
_log_cmd = list(cmd)
|
||
if "--api-key" in _log_cmd:
|
||
_ki = _log_cmd.index("--api-key") + 1
|
||
if _ki < len(_log_cmd):
|
||
_log_cmd[_ki] = "<redacted>"
|
||
logger.info(f"Starting llama-server: {' '.join(_log_cmd)}")
|
||
|
||
# Set library paths so llama-server can find its shared libs and CUDA DLLs
|
||
import os
|
||
import sys
|
||
|
||
env = os.environ.copy()
|
||
binary_dir = str(Path(binary).parent)
|
||
|
||
if sys.platform == "win32":
|
||
# On Windows, CUDA DLLs (cublas64_12.dll, cudart64_12.dll, etc.)
|
||
# must be on PATH. Add CUDA_PATH\bin if available.
|
||
path_dirs = [binary_dir]
|
||
cuda_path = os.environ.get("CUDA_PATH", "")
|
||
if cuda_path:
|
||
cuda_bin = os.path.join(cuda_path, "bin")
|
||
if os.path.isdir(cuda_bin):
|
||
path_dirs.append(cuda_bin)
|
||
# Some CUDA installs put DLLs in bin\x64
|
||
cuda_bin_x64 = os.path.join(cuda_path, "bin", "x64")
|
||
if os.path.isdir(cuda_bin_x64):
|
||
path_dirs.append(cuda_bin_x64)
|
||
existing_path = env.get("PATH", "")
|
||
env["PATH"] = ";".join(path_dirs) + ";" + existing_path
|
||
else:
|
||
# Linux: set LD_LIBRARY_PATH for shared libs next to the binary
|
||
# and CUDA runtime libs (libcudart, libcublas, etc.)
|
||
import platform
|
||
|
||
lib_dirs = [binary_dir]
|
||
_arch = platform.machine() # x86_64, aarch64, etc.
|
||
|
||
# Pip-installed nvidia CUDA runtime libs (e.g. torch's
|
||
# bundled cuda-bindings). The prebuilt llama.cpp binary
|
||
# links against libcudart.so.13 / libcublas.so.13 which
|
||
# live here, not in /usr/local/cuda.
|
||
import glob as _glob
|
||
|
||
for _nv_pattern in [
|
||
os.path.join(
|
||
sys.prefix,
|
||
"lib",
|
||
"python*",
|
||
"site-packages",
|
||
"nvidia",
|
||
"cu*",
|
||
"lib",
|
||
),
|
||
os.path.join(
|
||
sys.prefix,
|
||
"lib",
|
||
"python*",
|
||
"site-packages",
|
||
"nvidia",
|
||
"cudnn",
|
||
"lib",
|
||
),
|
||
os.path.join(
|
||
sys.prefix,
|
||
"lib",
|
||
"python*",
|
||
"site-packages",
|
||
"nvidia",
|
||
"nvjitlink",
|
||
"lib",
|
||
),
|
||
]:
|
||
for _nv_dir in _glob.glob(_nv_pattern):
|
||
if os.path.isdir(_nv_dir):
|
||
lib_dirs.append(_nv_dir)
|
||
|
||
for cuda_lib in [
|
||
"/usr/local/cuda/lib64",
|
||
f"/usr/local/cuda/targets/{_arch}-linux/lib",
|
||
# Fallback CUDA compat paths (e.g. binary built with
|
||
# CUDA 12 on a system where default /usr/local/cuda
|
||
# points to CUDA 13+).
|
||
"/usr/local/cuda-12/lib64",
|
||
"/usr/local/cuda-12.8/lib64",
|
||
f"/usr/local/cuda-12/targets/{_arch}-linux/lib",
|
||
f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib",
|
||
]:
|
||
if os.path.isdir(cuda_lib):
|
||
lib_dirs.append(cuda_lib)
|
||
existing_ld = env.get("LD_LIBRARY_PATH", "")
|
||
new_ld = ":".join(lib_dirs)
|
||
env["LD_LIBRARY_PATH"] = (
|
||
f"{new_ld}:{existing_ld}" if existing_ld else new_ld
|
||
)
|
||
|
||
# Pin to selected GPU(s) via CUDA_VISIBLE_DEVICES
|
||
if gpu_indices is not None:
|
||
env["CUDA_VISIBLE_DEVICES"] = ",".join(str(i) for i in gpu_indices)
|
||
|
||
self._stdout_lines = []
|
||
self._process = subprocess.Popen(
|
||
cmd,
|
||
stdout = subprocess.PIPE,
|
||
stderr = subprocess.STDOUT,
|
||
text = True,
|
||
env = env,
|
||
)
|
||
|
||
# Start background thread to drain stdout and prevent pipe deadlock
|
||
self._stdout_thread = threading.Thread(
|
||
target = self._drain_stdout, daemon = True, name = "llama-stdout"
|
||
)
|
||
self._stdout_thread.start()
|
||
|
||
self._gguf_path = gguf_path
|
||
self._hf_repo = hf_repo
|
||
# For local GGUF files, extract variant from filename if not provided
|
||
if hf_variant:
|
||
self._hf_variant = hf_variant
|
||
elif gguf_path:
|
||
try:
|
||
from utils.models.model_config import _extract_quant_label
|
||
|
||
self._hf_variant = _extract_quant_label(gguf_path)
|
||
except Exception:
|
||
self._hf_variant = None
|
||
else:
|
||
self._hf_variant = None
|
||
self._is_vision = is_vision
|
||
self._model_identifier = model_identifier
|
||
|
||
# Store the effective (possibly capped) context separately.
|
||
# Do NOT overwrite _context_length -- it holds the model's native
|
||
# context length from GGUF metadata and is used for display/info.
|
||
self._effective_context_length = (
|
||
effective_ctx if effective_ctx > 0 else self._context_length
|
||
)
|
||
self._max_context_length = (
|
||
max_available_ctx
|
||
if max_available_ctx > 0
|
||
else self._effective_context_length
|
||
)
|
||
|
||
# Wait for llama-server to become healthy
|
||
if not self._wait_for_health(timeout = 600.0):
|
||
self._kill_process()
|
||
raise RuntimeError(
|
||
"llama-server failed to start. "
|
||
"Check that the GGUF file is valid and you have enough memory."
|
||
)
|
||
|
||
self._healthy = True
|
||
|
||
logger.info(
|
||
f"llama-server ready on port {self._port} "
|
||
f"for model '{model_identifier}'"
|
||
)
|
||
return True
|
||
|
||
def unload_model(self) -> bool:
|
||
"""Terminate the llama-server subprocess and cancel any in-flight download."""
|
||
self._cancel_event.set()
|
||
with self._lock:
|
||
self._kill_process()
|
||
logger.info(f"Unloaded GGUF model: {self._model_identifier}")
|
||
self._model_identifier = None
|
||
self._gguf_path = None
|
||
self._hf_repo = None
|
||
self._hf_variant = None
|
||
self._is_vision = False
|
||
self._is_audio = False
|
||
self._audio_type = None
|
||
self._port = None
|
||
self._healthy = False
|
||
self._context_length = None
|
||
self._effective_context_length = None
|
||
self._max_context_length = None
|
||
self._chat_template = None
|
||
self._supports_reasoning = False
|
||
self._reasoning_always_on = False
|
||
self._supports_tools = False
|
||
self._cache_type_kv = None
|
||
self._n_layers = None
|
||
self._n_kv_heads = None
|
||
self._n_heads = None
|
||
self._embedding_length = None
|
||
# Clean up temp chat template file
|
||
if hasattr(self, "_chat_template_file") and self._chat_template_file:
|
||
try:
|
||
import os
|
||
|
||
os.unlink(self._chat_template_file.name)
|
||
except Exception:
|
||
pass
|
||
self._chat_template_file = None
|
||
# Free audio codec GPU memory
|
||
if LlamaCppBackend._codec_mgr is not None:
|
||
LlamaCppBackend._codec_mgr.unload()
|
||
LlamaCppBackend._codec_mgr = None
|
||
import torch
|
||
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
return True
|
||
|
||
def _kill_process(self):
|
||
"""Terminate the subprocess if running."""
|
||
if self._process is None:
|
||
return
|
||
try:
|
||
self._process.terminate()
|
||
self._process.wait(timeout = 5)
|
||
except subprocess.TimeoutExpired:
|
||
logger.warning("llama-server did not exit on SIGTERM, sending SIGKILL")
|
||
self._process.kill()
|
||
self._process.wait(timeout = 5)
|
||
except Exception as e:
|
||
logger.warning(f"Error killing llama-server process: {e}")
|
||
finally:
|
||
self._process = None
|
||
if self._stdout_thread is not None:
|
||
self._stdout_thread.join(timeout = 2)
|
||
self._stdout_thread = None
|
||
|
||
@staticmethod
|
||
def _kill_orphaned_servers():
|
||
"""Kill orphaned llama-server processes started by studio.
|
||
|
||
Only kills processes whose resolved binary lives under a known
|
||
Studio install directory (or matches an exact env-var override)
|
||
to avoid terminating unrelated llama-server instances.
|
||
|
||
Mirrors every location that _find_llama_server_binary() can
|
||
return from so that orphans from any supported install path
|
||
are still cleaned up.
|
||
|
||
Uses psutil for cross-platform support (Linux, macOS, Windows).
|
||
Falls back to pgrep + /proc/<pid>/exe on Linux when psutil is
|
||
not installed.
|
||
"""
|
||
import os
|
||
import signal
|
||
import sys
|
||
|
||
try:
|
||
# -- Build the ownership allowlist --------------------------------
|
||
# Two kinds of matches:
|
||
# exact_binaries -- env var overrides (exact path match only)
|
||
# install_roots -- directory trees that are Studio-owned
|
||
# (binary must be *under* one of these)
|
||
install_roots: list[Path] = []
|
||
|
||
# Primary install dir (setup.sh / prebuilt installer)
|
||
install_roots.append(Path.home() / ".unsloth" / "llama.cpp")
|
||
|
||
# Legacy in-tree build dirs (older setup.sh versions)
|
||
project_root = Path(__file__).resolve().parents[4]
|
||
install_roots.append(project_root / "llama.cpp")
|
||
|
||
# Legacy: extracted binary
|
||
install_roots.append(project_root / "bin")
|
||
|
||
# UNSLOTH_LLAMA_CPP_PATH env var (custom install dir)
|
||
custom_dir = os.environ.get("UNSLOTH_LLAMA_CPP_PATH")
|
||
if custom_dir:
|
||
install_roots.append(Path(custom_dir))
|
||
|
||
# LLAMA_SERVER_PATH env var (exact binary path)
|
||
exact_binaries: list[Path] = []
|
||
env_binary = os.environ.get("LLAMA_SERVER_PATH")
|
||
if env_binary:
|
||
try:
|
||
exact_binaries.append(Path(env_binary).resolve())
|
||
except OSError:
|
||
pass
|
||
|
||
# Resolve all roots so is_relative_to works reliably
|
||
resolved_roots: list[Path] = []
|
||
for root in install_roots:
|
||
try:
|
||
resolved_roots.append(root.resolve())
|
||
except OSError:
|
||
pass
|
||
|
||
my_pid = os.getpid()
|
||
|
||
# -- Enumerate processes -------------------------------------------
|
||
# Prefer psutil (cross-platform). Fall back to pgrep + /proc on
|
||
# Linux when psutil is not installed.
|
||
try:
|
||
import psutil
|
||
|
||
has_psutil = True
|
||
except ImportError:
|
||
has_psutil = False
|
||
|
||
if has_psutil:
|
||
for proc in psutil.process_iter(["pid", "name", "exe"]):
|
||
try:
|
||
if proc.info["pid"] == my_pid:
|
||
continue
|
||
|
||
name = proc.info.get("name") or ""
|
||
if not name.lower().startswith("llama-server"):
|
||
continue
|
||
|
||
exe = proc.info.get("exe")
|
||
if not exe:
|
||
continue
|
||
|
||
exe_path = Path(exe).resolve()
|
||
|
||
# Check ownership: exact binary match OR binary is
|
||
# under a known install root (proper ancestry, not
|
||
# substring).
|
||
is_ours = exe_path in exact_binaries or any(
|
||
exe_path.is_relative_to(root) for root in resolved_roots
|
||
)
|
||
if not is_ours:
|
||
continue
|
||
|
||
proc.kill()
|
||
logger.info(
|
||
f"Killed orphaned llama-server process "
|
||
f"(pid={proc.info['pid']})"
|
||
)
|
||
except (
|
||
psutil.NoSuchProcess,
|
||
psutil.AccessDenied,
|
||
psutil.ZombieProcess,
|
||
):
|
||
pass
|
||
else:
|
||
# -- Fallback: pgrep + /proc/<pid>/exe (Linux only) -----------
|
||
if sys.platform != "linux":
|
||
return
|
||
result = subprocess.run(
|
||
["pgrep", "-a", "-f", "llama-server"],
|
||
capture_output = True,
|
||
text = True,
|
||
timeout = 5,
|
||
)
|
||
if result.returncode != 0:
|
||
return
|
||
|
||
for line in result.stdout.strip().splitlines():
|
||
parts = line.strip().split(None, 1)
|
||
if len(parts) < 2:
|
||
continue
|
||
pid = int(parts[0])
|
||
if pid == my_pid:
|
||
continue
|
||
|
||
# Resolve the actual executable. /proc/<pid>/exe is a
|
||
# symlink to the real binary and avoids all cmdline-
|
||
# parsing ambiguities (spaces in paths, argv rewriting).
|
||
# Fall back to the first cmdline token when /proc is
|
||
# unavailable.
|
||
proc_exe = Path(f"/proc/{pid}/exe")
|
||
try:
|
||
binary = proc_exe.resolve(strict = True)
|
||
except (OSError, ValueError):
|
||
cmdline = parts[1]
|
||
token = cmdline.split()[0] if cmdline.strip() else ""
|
||
if not token:
|
||
continue
|
||
binary = Path(token).resolve(strict = False)
|
||
|
||
owned = binary in exact_binaries or any(
|
||
binary.is_relative_to(root) for root in resolved_roots
|
||
)
|
||
if not owned:
|
||
continue
|
||
|
||
try:
|
||
os.kill(pid, signal.SIGKILL)
|
||
logger.info(f"Killed orphaned llama-server process (pid={pid})")
|
||
except ProcessLookupError:
|
||
pass
|
||
except PermissionError:
|
||
pass
|
||
except Exception:
|
||
logger.warning("Error during orphan server cleanup", exc_info = True)
|
||
|
||
def _cleanup(self):
|
||
"""atexit handler to ensure llama-server is terminated."""
|
||
self._kill_process()
|
||
|
||
def _wait_for_health(self, timeout: float = 120.0, interval: float = 0.5) -> bool:
|
||
"""
|
||
Poll llama-server's /health endpoint until it responds 200.
|
||
|
||
Also monitors subprocess for early exit/crash.
|
||
"""
|
||
deadline = time.monotonic() + timeout
|
||
url = f"http://127.0.0.1:{self._port}/health"
|
||
|
||
while time.monotonic() < deadline:
|
||
# Check if process crashed
|
||
if self._process.poll() is not None:
|
||
# Give the drain thread a moment to collect final output
|
||
if self._stdout_thread is not None:
|
||
self._stdout_thread.join(timeout = 2)
|
||
output = "\n".join(self._stdout_lines[-50:])
|
||
logger.error(
|
||
f"llama-server exited with code {self._process.returncode}. "
|
||
f"Output: {output[:2000]}"
|
||
)
|
||
return False
|
||
|
||
try:
|
||
resp = httpx.get(url, timeout = 2.0)
|
||
if resp.status_code == 200:
|
||
return True
|
||
except (httpx.ConnectError, httpx.TimeoutException):
|
||
pass
|
||
|
||
time.sleep(interval)
|
||
|
||
logger.error(f"llama-server health check timed out after {timeout}s")
|
||
return False
|
||
|
||
# ── Message building (OpenAI format) ──────────────────────────
|
||
|
||
@staticmethod
|
||
def _parse_tool_calls_from_text(content: str) -> list[dict]:
|
||
"""
|
||
Parse tool calls from XML markup in content text.
|
||
|
||
Handles formats like:
|
||
<tool_call>{"name":"web_search","arguments":{"query":"..."}}</tool_call>
|
||
<tool_call><function=web_search><parameter=query>...</parameter></function></tool_call>
|
||
Closing tags (</tool_call>, </function>, </parameter>) are all optional
|
||
since models frequently omit them.
|
||
"""
|
||
import re
|
||
|
||
tool_calls = []
|
||
|
||
# Pattern 1: JSON inside <tool_call> tags.
|
||
# Use balanced-brace extraction that skips braces inside JSON strings.
|
||
for m in re.finditer(r"<tool_call>\s*\{", content):
|
||
brace_start = m.end() - 1 # position of the opening {
|
||
depth, i = 0, brace_start
|
||
in_string = False
|
||
while i < len(content):
|
||
ch = content[i]
|
||
if in_string:
|
||
if ch == "\\" and i + 1 < len(content):
|
||
i += 2 # skip escaped character
|
||
continue
|
||
if ch == '"':
|
||
in_string = False
|
||
elif ch == '"':
|
||
in_string = True
|
||
elif ch == "{":
|
||
depth += 1
|
||
elif ch == "}":
|
||
depth -= 1
|
||
if depth == 0:
|
||
break
|
||
i += 1
|
||
if depth == 0:
|
||
json_str = content[brace_start : i + 1]
|
||
try:
|
||
obj = json.loads(json_str)
|
||
tc = {
|
||
"id": f"call_{len(tool_calls)}",
|
||
"type": "function",
|
||
"function": {
|
||
"name": obj.get("name", ""),
|
||
"arguments": obj.get("arguments", {}),
|
||
},
|
||
}
|
||
if isinstance(tc["function"]["arguments"], dict):
|
||
tc["function"]["arguments"] = json.dumps(
|
||
tc["function"]["arguments"]
|
||
)
|
||
tool_calls.append(tc)
|
||
except (json.JSONDecodeError, ValueError):
|
||
pass
|
||
|
||
# Pattern 2: XML-style <function=name><parameter=key>value</parameter></function>
|
||
# All closing tags optional -- models frequently omit </parameter>,
|
||
# </function>, and/or </tool_call>.
|
||
if not tool_calls:
|
||
# Step 1: Find all <function=name> positions and extract their bodies.
|
||
# Body boundary: use only </tool_call> or next <function= as hard
|
||
# boundaries. We avoid using </function> as a boundary because
|
||
# code parameter values can contain that literal string.
|
||
# After extracting, we trim a trailing </function> if present.
|
||
func_starts = list(re.finditer(r"<function=(\w+)>\s*", content))
|
||
for idx, fm in enumerate(func_starts):
|
||
func_name = fm.group(1)
|
||
body_start = fm.end()
|
||
# Hard boundaries: next <function= tag or </tool_call>
|
||
next_func = (
|
||
func_starts[idx + 1].start()
|
||
if idx + 1 < len(func_starts)
|
||
else len(content)
|
||
)
|
||
end_tag = re.search(r"</tool_call>", content[body_start:])
|
||
if end_tag:
|
||
body_end = body_start + end_tag.start()
|
||
else:
|
||
body_end = len(content)
|
||
body_end = min(body_end, next_func)
|
||
body = content[body_start:body_end]
|
||
# Trim trailing </function> if present (it's the real closing tag)
|
||
body = re.sub(r"\s*</function>\s*$", "", body)
|
||
|
||
# Step 2: Extract parameters from body.
|
||
# For single-parameter functions (the common case: code, command,
|
||
# query), use body end as the only boundary to avoid false matches
|
||
# on </parameter> inside code strings.
|
||
arguments = {}
|
||
param_starts = list(re.finditer(r"<parameter=(\w+)>\s*", body))
|
||
if len(param_starts) == 1:
|
||
# Single parameter: value is everything from after the tag
|
||
# to end of body, trimming any trailing </parameter>.
|
||
pm = param_starts[0]
|
||
val = body[pm.end() :]
|
||
val = re.sub(r"\s*</parameter>\s*$", "", val)
|
||
arguments[pm.group(1)] = val.strip()
|
||
else:
|
||
for pidx, pm in enumerate(param_starts):
|
||
param_name = pm.group(1)
|
||
val_start = pm.end()
|
||
# Value ends at next <parameter= or end of body
|
||
next_param = (
|
||
param_starts[pidx + 1].start()
|
||
if pidx + 1 < len(param_starts)
|
||
else len(body)
|
||
)
|
||
val = body[val_start:next_param]
|
||
# Trim trailing </parameter> if present
|
||
val = re.sub(r"\s*</parameter>\s*$", "", val)
|
||
arguments[param_name] = val.strip()
|
||
|
||
tc = {
|
||
"id": f"call_{len(tool_calls)}",
|
||
"type": "function",
|
||
"function": {
|
||
"name": func_name,
|
||
"arguments": json.dumps(arguments),
|
||
},
|
||
}
|
||
tool_calls.append(tc)
|
||
|
||
return tool_calls
|
||
|
||
@staticmethod
|
||
def _build_openai_messages(
|
||
messages: list[dict],
|
||
image_b64: Optional[str] = None,
|
||
) -> list[dict]:
|
||
"""
|
||
Build OpenAI-format messages, optionally injecting an image_url
|
||
content part into the last user message for vision models.
|
||
|
||
If no image is provided, returns messages as-is.
|
||
"""
|
||
if not image_b64:
|
||
return messages
|
||
|
||
# Find the last user message and convert to multimodal content parts
|
||
result = [msg.copy() for msg in messages]
|
||
last_user_idx = None
|
||
for i, msg in enumerate(result):
|
||
if msg["role"] == "user":
|
||
last_user_idx = i
|
||
|
||
if last_user_idx is not None:
|
||
text_content = result[last_user_idx].get("content", "")
|
||
result[last_user_idx]["content"] = [
|
||
{"type": "text", "text": text_content},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:image/png;base64,{image_b64}",
|
||
},
|
||
},
|
||
]
|
||
|
||
return result
|
||
|
||
# ── Generation (proxy to llama-server) ────────────────────────
|
||
|
||
@staticmethod
|
||
def _iter_text_cancellable(
|
||
response: "httpx.Response",
|
||
cancel_event: Optional[threading.Event] = None,
|
||
) -> Generator[str, None, None]:
|
||
"""Iterate over an httpx streaming response with cancel support.
|
||
|
||
Checks cancel_event between chunks and on ReadTimeout. The
|
||
cancel watcher in _stream_with_retry also calls response.close()
|
||
on cancel, which unblocks iter_text() once the response exists.
|
||
During normal streaming llama-server sends tokens frequently,
|
||
so the cancel check between chunks is the primary mechanism.
|
||
"""
|
||
text_iter = response.iter_text()
|
||
while True:
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
response.close()
|
||
return
|
||
try:
|
||
chunk = next(text_iter)
|
||
yield chunk
|
||
except StopIteration:
|
||
return
|
||
except httpx.ReadTimeout:
|
||
# No data within the timeout window -- just loop back
|
||
# and re-check cancel_event.
|
||
continue
|
||
|
||
@staticmethod
|
||
@contextlib.contextmanager
|
||
def _stream_with_retry(
|
||
client: "httpx.Client",
|
||
url: str,
|
||
payload: dict,
|
||
cancel_event: Optional[threading.Event] = None,
|
||
headers: Optional[dict] = None,
|
||
):
|
||
"""Open an httpx streaming POST with cancel support.
|
||
|
||
Sends the request once with a long read timeout (120 s) so
|
||
prompt processing (prefill) can finish without triggering a
|
||
retry storm. The previous 0.5 s timeout caused duplicate POST
|
||
requests every half second, forcing llama-server to restart
|
||
processing each time.
|
||
|
||
A background watcher thread provides cancel by closing the
|
||
response when cancel_event is set. Limitation: httpx does not
|
||
allow interrupting a blocked read from another thread before
|
||
the response object exists, so cancel during the initial
|
||
header wait (prefill phase) only takes effect once headers
|
||
arrive. After that, response.close() unblocks reads promptly.
|
||
In practice llama-server prefill is 1-5 s for typical prompts,
|
||
during which cancel is deferred -- still much better than the
|
||
old retry storm which made prefill slower.
|
||
"""
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
raise GeneratorExit
|
||
|
||
# Background watcher: close the response if cancel is requested.
|
||
# Only effective after response headers arrive (httpx limitation).
|
||
_cancel_closed = threading.Event()
|
||
_response_ref: list = [None]
|
||
|
||
def _cancel_watcher():
|
||
while not _cancel_closed.is_set():
|
||
if cancel_event.wait(timeout = 0.3):
|
||
# Cancel requested. Keep polling until the response object
|
||
# exists so we can close it, or until the main thread
|
||
# finishes on its own (_cancel_closed is set in finally).
|
||
while not _cancel_closed.is_set():
|
||
r = _response_ref[0]
|
||
if r is not None:
|
||
try:
|
||
r.close()
|
||
return
|
||
except Exception as e:
|
||
logger.debug(
|
||
f"Error closing response in cancel watcher: {e}"
|
||
)
|
||
# Response not created yet -- wait briefly and retry
|
||
_cancel_closed.wait(timeout = 0.1)
|
||
return
|
||
|
||
watcher = None
|
||
if cancel_event is not None:
|
||
watcher = threading.Thread(
|
||
target = _cancel_watcher, daemon = True, name = "prefill-cancel"
|
||
)
|
||
watcher.start()
|
||
|
||
try:
|
||
# Long read timeout so prefill (prompt processing) can finish
|
||
# without triggering a retry storm. Cancel during both
|
||
# prefill and streaming is handled by the watcher thread
|
||
# which closes the response, unblocking any httpx read.
|
||
prefill_timeout = httpx.Timeout(
|
||
connect = 30,
|
||
read = 120.0,
|
||
write = 10,
|
||
pool = 10,
|
||
)
|
||
with client.stream(
|
||
"POST",
|
||
url,
|
||
json = payload,
|
||
timeout = prefill_timeout,
|
||
headers = headers,
|
||
) as response:
|
||
_response_ref[0] = response
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
raise GeneratorExit
|
||
yield response
|
||
return
|
||
except (httpx.ReadError, httpx.RemoteProtocolError, httpx.CloseError):
|
||
# Response was closed by the cancel watcher
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
raise GeneratorExit
|
||
raise
|
||
finally:
|
||
_cancel_closed.set()
|
||
|
||
def generate_chat_completion(
|
||
self,
|
||
messages: list[dict],
|
||
image_b64: Optional[str] = None,
|
||
temperature: float = 0.6,
|
||
top_p: float = 0.95,
|
||
top_k: int = 20,
|
||
min_p: float = 0.01,
|
||
max_tokens: Optional[int] = None,
|
||
repetition_penalty: float = 1.0,
|
||
presence_penalty: float = 0.0,
|
||
stop: Optional[list[str]] = None,
|
||
cancel_event: Optional[threading.Event] = None,
|
||
enable_thinking: Optional[bool] = None,
|
||
) -> Generator[str | dict, None, None]:
|
||
"""
|
||
Send a chat completion request to llama-server and stream tokens back.
|
||
|
||
Uses /v1/chat/completions — llama-server handles chat template
|
||
application and vision (multimodal image_url parts) natively.
|
||
|
||
Yields cumulative text (matching InferenceBackend's convention).
|
||
"""
|
||
if not self.is_loaded:
|
||
raise RuntimeError("llama-server is not loaded")
|
||
|
||
openai_messages = self._build_openai_messages(messages, image_b64)
|
||
|
||
payload = {
|
||
"messages": openai_messages,
|
||
"stream": True,
|
||
"temperature": temperature,
|
||
"top_p": top_p,
|
||
"top_k": top_k if top_k >= 0 else 0,
|
||
"min_p": min_p,
|
||
"repeat_penalty": repetition_penalty,
|
||
"presence_penalty": presence_penalty,
|
||
}
|
||
# Pass enable_thinking per-request for reasoning models
|
||
if self._supports_reasoning and enable_thinking is not None:
|
||
payload["chat_template_kwargs"] = {"enable_thinking": enable_thinking}
|
||
if max_tokens is not None:
|
||
payload["max_tokens"] = max_tokens
|
||
if stop:
|
||
payload["stop"] = stop
|
||
payload["stream_options"] = {"include_usage": True}
|
||
|
||
url = f"{self.base_url}/v1/chat/completions"
|
||
cumulative = ""
|
||
in_thinking = False
|
||
_stream_done = False
|
||
_metadata_usage = None
|
||
_metadata_timings = None
|
||
|
||
try:
|
||
# _stream_with_retry uses a 120 s read timeout so prefill
|
||
# can finish. Cancel during streaming is handled by the
|
||
# watcher thread (closes the response on cancel_event).
|
||
stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10)
|
||
_auth_headers = (
|
||
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
||
)
|
||
with httpx.Client(timeout = stream_timeout) as client:
|
||
with self._stream_with_retry(
|
||
client,
|
||
url,
|
||
payload,
|
||
cancel_event,
|
||
headers = _auth_headers,
|
||
) as response:
|
||
if response.status_code != 200:
|
||
error_body = response.read().decode()
|
||
raise RuntimeError(
|
||
f"llama-server returned {response.status_code}: {error_body}"
|
||
)
|
||
|
||
buffer = ""
|
||
has_content_tokens = False
|
||
reasoning_text = ""
|
||
for raw_chunk in self._iter_text_cancellable(
|
||
response, cancel_event
|
||
):
|
||
buffer += raw_chunk
|
||
while "\n" in buffer:
|
||
line, buffer = buffer.split("\n", 1)
|
||
line = line.strip()
|
||
|
||
if not line:
|
||
continue
|
||
if line == "data: [DONE]":
|
||
if in_thinking:
|
||
if has_content_tokens:
|
||
# Real thinking + content: close the tag
|
||
cumulative += "</think>"
|
||
yield cumulative
|
||
else:
|
||
# Only reasoning_content, no content tokens:
|
||
# the model put its entire reply in reasoning
|
||
# (e.g. Qwen3 always-think mode). Show it
|
||
# as the main response, not as a thinking block.
|
||
cumulative = reasoning_text
|
||
yield cumulative
|
||
_stream_done = True
|
||
break # exit inner while
|
||
if not line.startswith("data: "):
|
||
continue
|
||
|
||
try:
|
||
data = json.loads(line[6:])
|
||
# Capture server timings/usage from final chunks
|
||
_chunk_timings = data.get("timings")
|
||
if _chunk_timings:
|
||
_metadata_timings = _chunk_timings
|
||
_chunk_usage = data.get("usage")
|
||
if _chunk_usage:
|
||
_metadata_usage = _chunk_usage
|
||
choices = data.get("choices", [])
|
||
if choices:
|
||
delta = choices[0].get("delta", {})
|
||
|
||
# Handle reasoning/thinking tokens
|
||
# llama-server sends these as "reasoning_content"
|
||
# Wrap in <think> tags for the frontend parser
|
||
reasoning = delta.get("reasoning_content", "")
|
||
if reasoning:
|
||
reasoning_text += reasoning
|
||
if not in_thinking:
|
||
cumulative += "<think>"
|
||
in_thinking = True
|
||
cumulative += reasoning
|
||
yield cumulative
|
||
|
||
token = delta.get("content", "")
|
||
if token:
|
||
has_content_tokens = True
|
||
if in_thinking:
|
||
cumulative += "</think>"
|
||
in_thinking = False
|
||
cumulative += token
|
||
yield cumulative
|
||
except json.JSONDecodeError:
|
||
logger.debug(
|
||
f"Skipping malformed SSE line: {line[:100]}"
|
||
)
|
||
if _stream_done:
|
||
break # exit outer for
|
||
if _metadata_usage or _metadata_timings:
|
||
yield {
|
||
"type": "metadata",
|
||
"usage": _metadata_usage,
|
||
"timings": _metadata_timings,
|
||
}
|
||
|
||
except httpx.ConnectError:
|
||
raise RuntimeError("Lost connection to llama-server")
|
||
except Exception as e:
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
return
|
||
raise
|
||
|
||
# ── Tool-calling agentic loop ──────────────────────────────
|
||
|
||
def generate_chat_completion_with_tools(
|
||
self,
|
||
messages: list[dict],
|
||
tools: list[dict],
|
||
temperature: float = 0.6,
|
||
top_p: float = 0.95,
|
||
top_k: int = 20,
|
||
min_p: float = 0.01,
|
||
max_tokens: Optional[int] = None,
|
||
repetition_penalty: float = 1.0,
|
||
presence_penalty: float = 0.0,
|
||
stop: Optional[list[str]] = None,
|
||
cancel_event: Optional[threading.Event] = None,
|
||
enable_thinking: Optional[bool] = None,
|
||
max_tool_iterations: int = 25,
|
||
auto_heal_tool_calls: bool = True,
|
||
tool_call_timeout: int = 300,
|
||
session_id: Optional[str] = None,
|
||
) -> Generator[dict, None, None]:
|
||
"""
|
||
Agentic loop: let the model call tools, execute them, and continue.
|
||
|
||
Yields dicts with:
|
||
{"type": "status", "text": "Searching: ..."} -- tool status updates
|
||
{"type": "content", "text": "token"} -- streamed content tokens (cumulative)
|
||
{"type": "reasoning", "text": "token"} -- streamed reasoning tokens (cumulative)
|
||
"""
|
||
from core.inference.tools import execute_tool
|
||
|
||
if not self.is_loaded:
|
||
raise RuntimeError("llama-server is not loaded")
|
||
|
||
conversation = list(messages)
|
||
url = f"{self.base_url}/v1/chat/completions"
|
||
_accumulated_completion_tokens = 0
|
||
_accumulated_predicted_ms = 0.0
|
||
_accumulated_predicted_n = 0
|
||
|
||
# ── Shared patterns for stripping tool XML from streamed content ──
|
||
import re as _re_tool
|
||
|
||
_TOOL_CLOSED_PATTERNS = [
|
||
_re_tool.compile(r"<tool_call>.*?</tool_call>", _re_tool.DOTALL),
|
||
_re_tool.compile(r"<function=\w+>.*?</function>", _re_tool.DOTALL),
|
||
]
|
||
_TOOL_ALL_PATTERNS = _TOOL_CLOSED_PATTERNS + [
|
||
_re_tool.compile(r"<tool_call>.*$", _re_tool.DOTALL),
|
||
_re_tool.compile(r"<function=\w+>.*$", _re_tool.DOTALL),
|
||
]
|
||
|
||
def _strip_tool_markup(text: str, *, final: bool = False) -> str:
|
||
if not auto_heal_tool_calls:
|
||
return text
|
||
patterns = _TOOL_ALL_PATTERNS if final else _TOOL_CLOSED_PATTERNS
|
||
for pat in patterns:
|
||
text = pat.sub("", text)
|
||
return text.strip() if final else text
|
||
|
||
# XML prefixes that signal a tool call in content.
|
||
# Empty when auto_heal is disabled so the buffer never
|
||
# speculatively holds content for XML detection.
|
||
_TOOL_XML_SIGNALS = (
|
||
("<tool_call>", "<function=") if auto_heal_tool_calls else ()
|
||
)
|
||
_MAX_BUFFER_CHARS = 32
|
||
|
||
# ── Duplicate tool-call detection ────────────────────────
|
||
# Track recent (tool_name, arguments) hashes to detect loops
|
||
# where the model repeats the exact same call. Retries after
|
||
# a transient failure are allowed (only block when the previous
|
||
# identical call succeeded).
|
||
_tool_call_history: list[tuple[str, bool]] = [] # (key, failed)
|
||
|
||
def _tool_call_key(name: str, args: dict) -> str:
|
||
raw = json.dumps({"t": name, "a": args}, sort_keys = True)
|
||
return hashlib.md5(raw.encode()).hexdigest()
|
||
|
||
def _is_duplicate_call(name: str, args: dict) -> bool:
|
||
"""Block if the immediately previous call was identical and succeeded."""
|
||
if not _tool_call_history:
|
||
return False
|
||
key = _tool_call_key(name, args)
|
||
last_key, last_failed = _tool_call_history[-1]
|
||
return last_key == key and not last_failed
|
||
|
||
def _record_tool_call(name: str, args: dict, failed: bool) -> None:
|
||
key = _tool_call_key(name, args)
|
||
_tool_call_history.append((key, failed))
|
||
|
||
for iteration in range(max_tool_iterations):
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
return
|
||
|
||
# Build payload -- stream: True so we detect tool signals
|
||
# in the first 1-2 chunks without a non-streaming penalty.
|
||
payload = {
|
||
"messages": conversation,
|
||
"stream": True,
|
||
"stream_options": {"include_usage": True},
|
||
"temperature": temperature,
|
||
"top_p": top_p,
|
||
"top_k": top_k if top_k >= 0 else 0,
|
||
"min_p": min_p,
|
||
"repeat_penalty": repetition_penalty,
|
||
"presence_penalty": presence_penalty,
|
||
"tools": tools,
|
||
"tool_choice": "auto",
|
||
}
|
||
if self._supports_reasoning and enable_thinking is not None:
|
||
payload["chat_template_kwargs"] = {"enable_thinking": enable_thinking}
|
||
if max_tokens is not None:
|
||
payload["max_tokens"] = max_tokens
|
||
if stop:
|
||
payload["stop"] = stop
|
||
|
||
try:
|
||
_auth_headers = (
|
||
{"Authorization": f"Bearer {self._api_key}"}
|
||
if self._api_key
|
||
else None
|
||
)
|
||
|
||
# ── Speculative buffer state machine ──────────────────
|
||
# BUFFERING: accumulating content, checking for tool signals
|
||
# STREAMING: no tool detected, yielding tokens to caller
|
||
# DRAINING: tool signal found, silently consuming rest
|
||
_S_BUFFERING = 0
|
||
_S_STREAMING = 1
|
||
_S_DRAINING = 2
|
||
|
||
detect_state = _S_BUFFERING
|
||
content_buffer = "" # Raw content held during BUFFERING
|
||
content_accum = "" # All content tokens (for tool parsing)
|
||
reasoning_accum = ""
|
||
cumulative_display = "" # Cumulative text yielded (with <think>)
|
||
in_thinking = False
|
||
has_content_tokens = False
|
||
tool_calls_acc = {} # Structured delta.tool_calls fragments
|
||
has_structured_tc = False
|
||
_iter_usage = None
|
||
_iter_timings = None
|
||
_stream_done = False
|
||
_last_emitted = ""
|
||
|
||
stream_timeout = httpx.Timeout(
|
||
connect = 10,
|
||
read = 0.5,
|
||
write = 10,
|
||
pool = 10,
|
||
)
|
||
with httpx.Client(timeout = stream_timeout) as client:
|
||
with self._stream_with_retry(
|
||
client,
|
||
url,
|
||
payload,
|
||
cancel_event,
|
||
headers = _auth_headers,
|
||
) as response:
|
||
if response.status_code != 200:
|
||
error_body = response.read().decode()
|
||
raise RuntimeError(
|
||
f"llama-server returned {response.status_code}: "
|
||
f"{error_body}"
|
||
)
|
||
|
||
raw_buf = ""
|
||
for raw_chunk in self._iter_text_cancellable(
|
||
response,
|
||
cancel_event,
|
||
):
|
||
raw_buf += raw_chunk
|
||
while "\n" in raw_buf:
|
||
line, raw_buf = raw_buf.split("\n", 1)
|
||
line = line.strip()
|
||
|
||
if not line:
|
||
continue
|
||
if line == "data: [DONE]":
|
||
# Flush thinking state for STREAMING
|
||
if detect_state == _S_STREAMING and in_thinking:
|
||
if has_content_tokens:
|
||
cumulative_display += "</think>"
|
||
yield {
|
||
"type": "content",
|
||
"text": _strip_tool_markup(
|
||
cumulative_display,
|
||
final = True,
|
||
),
|
||
}
|
||
else:
|
||
cumulative_display = reasoning_accum
|
||
yield {
|
||
"type": "content",
|
||
"text": cumulative_display,
|
||
}
|
||
_stream_done = True
|
||
break # exit inner while
|
||
if not line.startswith("data: "):
|
||
continue
|
||
|
||
try:
|
||
chunk_data = json.loads(line[6:])
|
||
_ct = chunk_data.get("timings")
|
||
if _ct:
|
||
_iter_timings = _ct
|
||
_cu = chunk_data.get("usage")
|
||
if _cu:
|
||
_iter_usage = _cu
|
||
|
||
choices = chunk_data.get("choices", [])
|
||
if not choices:
|
||
continue
|
||
|
||
delta = choices[0].get("delta", {})
|
||
|
||
# ── Structured tool_calls ──
|
||
tc_deltas = delta.get("tool_calls")
|
||
if tc_deltas:
|
||
# Once visible content has been
|
||
# emitted, do not reclassify this
|
||
# turn as a tool call.
|
||
if _last_emitted:
|
||
continue
|
||
has_structured_tc = True
|
||
detect_state = _S_DRAINING
|
||
for tc_d in tc_deltas:
|
||
idx = tc_d.get("index", 0)
|
||
if idx not in tool_calls_acc:
|
||
tool_calls_acc[idx] = {
|
||
"id": tc_d.get("id", f"call_{idx}"),
|
||
"type": "function",
|
||
"function": {
|
||
"name": "",
|
||
"arguments": "",
|
||
},
|
||
}
|
||
elif tc_d.get("id"):
|
||
# Update ID if real one
|
||
# arrives on a later delta
|
||
tool_calls_acc[idx]["id"] = tc_d["id"]
|
||
func = tc_d.get("function", {})
|
||
if func.get("name"):
|
||
tool_calls_acc[idx]["function"][
|
||
"name"
|
||
] += func["name"]
|
||
if func.get("arguments"):
|
||
tool_calls_acc[idx]["function"][
|
||
"arguments"
|
||
] += func["arguments"]
|
||
continue
|
||
|
||
# ── Reasoning tokens ──
|
||
# Only yield in STREAMING state. In BUFFERING
|
||
# and DRAINING, accumulate silently so we don't
|
||
# corrupt the consumer's prev_text tracker
|
||
# (routes/inference.py never resets prev_text
|
||
# between tool iterations).
|
||
reasoning = delta.get("reasoning_content", "")
|
||
if reasoning:
|
||
reasoning_accum += reasoning
|
||
if detect_state == _S_STREAMING:
|
||
if not in_thinking:
|
||
cumulative_display += "<think>"
|
||
in_thinking = True
|
||
cumulative_display += reasoning
|
||
yield {
|
||
"type": "content",
|
||
"text": cumulative_display,
|
||
}
|
||
|
||
# ── Content tokens ──
|
||
token = delta.get("content", "")
|
||
if token:
|
||
has_content_tokens = True
|
||
content_accum += token
|
||
|
||
if detect_state == _S_DRAINING:
|
||
pass # accumulate silently
|
||
|
||
elif detect_state == _S_STREAMING:
|
||
if in_thinking:
|
||
cumulative_display += "</think>"
|
||
in_thinking = False
|
||
cumulative_display += token
|
||
cleaned = _strip_tool_markup(
|
||
cumulative_display,
|
||
)
|
||
if len(cleaned) > len(_last_emitted):
|
||
_last_emitted = cleaned
|
||
yield {
|
||
"type": "content",
|
||
"text": cleaned,
|
||
}
|
||
|
||
elif detect_state == _S_BUFFERING:
|
||
content_buffer += token
|
||
stripped_buf = content_buffer.lstrip()
|
||
if not stripped_buf:
|
||
continue
|
||
|
||
# Check tool signal prefixes
|
||
is_prefix = False
|
||
is_match = False
|
||
for sig in _TOOL_XML_SIGNALS:
|
||
if stripped_buf.startswith(sig):
|
||
is_match = True
|
||
break
|
||
if sig.startswith(stripped_buf):
|
||
is_prefix = True
|
||
break
|
||
|
||
if is_match:
|
||
detect_state = _S_DRAINING
|
||
elif (
|
||
is_prefix
|
||
and len(stripped_buf)
|
||
< _MAX_BUFFER_CHARS
|
||
):
|
||
pass # keep buffering
|
||
else:
|
||
# Not a tool -- flush buffer
|
||
detect_state = _S_STREAMING
|
||
# Flush any reasoning accumulated
|
||
# during BUFFERING phase
|
||
if reasoning_accum:
|
||
cumulative_display += "<think>"
|
||
cumulative_display += (
|
||
reasoning_accum
|
||
)
|
||
cumulative_display += "</think>"
|
||
cumulative_display += content_buffer
|
||
cleaned = _strip_tool_markup(
|
||
cumulative_display,
|
||
)
|
||
if len(cleaned) > len(_last_emitted):
|
||
_last_emitted = cleaned
|
||
yield {
|
||
"type": "content",
|
||
"text": cleaned,
|
||
}
|
||
|
||
except json.JSONDecodeError:
|
||
logger.debug(
|
||
f"Skipping malformed SSE line: " f"{line[:100]}"
|
||
)
|
||
if _stream_done:
|
||
break # exit outer for
|
||
|
||
# ── Resolve BUFFERING at stream end ──
|
||
if detect_state == _S_BUFFERING:
|
||
stripped_buf = content_buffer.lstrip()
|
||
if (
|
||
stripped_buf
|
||
and auto_heal_tool_calls
|
||
and any(s in stripped_buf for s in _TOOL_XML_SIGNALS)
|
||
):
|
||
detect_state = _S_DRAINING
|
||
elif content_accum or reasoning_accum:
|
||
detect_state = _S_STREAMING
|
||
if content_buffer:
|
||
# Flush any reasoning accumulated first
|
||
if reasoning_accum:
|
||
cumulative_display += "<think>"
|
||
cumulative_display += reasoning_accum
|
||
cumulative_display += "</think>"
|
||
cumulative_display += content_buffer
|
||
yield {
|
||
"type": "content",
|
||
"text": _strip_tool_markup(
|
||
cumulative_display,
|
||
final = True,
|
||
),
|
||
}
|
||
elif reasoning_accum and not has_content_tokens:
|
||
# Reasoning-only response (no content tokens):
|
||
# show reasoning as plain text, matching
|
||
# the final streaming pass behavior for
|
||
# models that put everything in reasoning.
|
||
cumulative_display = reasoning_accum
|
||
yield {
|
||
"type": "content",
|
||
"text": cumulative_display,
|
||
}
|
||
else:
|
||
return
|
||
|
||
# ── STREAMING path: no tool call ──
|
||
if detect_state == _S_STREAMING:
|
||
# Safety net: check for XML tool signals in content.
|
||
# The route layer resets prev_text on tool_start, so
|
||
# post-tool synthesis streams correctly even if
|
||
# content was already emitted before the tool XML.
|
||
_safety_tc = None
|
||
if auto_heal_tool_calls and any(
|
||
s in content_accum for s in _TOOL_XML_SIGNALS
|
||
):
|
||
_safety_tc = self._parse_tool_calls_from_text(
|
||
content_accum,
|
||
)
|
||
if not _safety_tc:
|
||
# Content was already streamed. Yield metadata.
|
||
yield {"type": "status", "text": ""}
|
||
_fu = _iter_usage or {}
|
||
_fc = _fu.get("completion_tokens", 0)
|
||
_fp = _fu.get("prompt_tokens", 0)
|
||
_tc = _fc + _accumulated_completion_tokens
|
||
if (
|
||
_iter_usage
|
||
or _iter_timings
|
||
or _accumulated_completion_tokens
|
||
):
|
||
_mt = dict(_iter_timings) if _iter_timings else {}
|
||
if _accumulated_predicted_ms or _accumulated_predicted_n:
|
||
_mt["predicted_ms"] = (
|
||
_mt.get("predicted_ms", 0)
|
||
+ _accumulated_predicted_ms
|
||
)
|
||
_tn = (
|
||
_mt.get("predicted_n", 0) + _accumulated_predicted_n
|
||
)
|
||
_mt["predicted_n"] = _tn
|
||
_tms = _mt["predicted_ms"]
|
||
if _tms > 0:
|
||
_mt["predicted_per_second"] = _tn / (_tms / 1000.0)
|
||
yield {
|
||
"type": "metadata",
|
||
"usage": {
|
||
"prompt_tokens": _fp,
|
||
"completion_tokens": _tc,
|
||
"total_tokens": _fp + _tc,
|
||
},
|
||
"timings": _mt,
|
||
}
|
||
return
|
||
|
||
# Safety net caught tool XML -- treat as tool call
|
||
tool_calls = _safety_tc
|
||
content_text = _strip_tool_markup(
|
||
content_accum,
|
||
final = True,
|
||
)
|
||
logger.info(
|
||
f"Safety net: parsed {len(tool_calls)} tool call(s) "
|
||
f"from streamed content"
|
||
)
|
||
else:
|
||
# ── DRAINING path: assemble tool_calls ──
|
||
tool_calls = None
|
||
content_text = content_accum
|
||
if has_structured_tc:
|
||
# Filter out incomplete fragments (e.g. from
|
||
# truncation by max_tokens or disconnect).
|
||
tool_calls = [
|
||
tool_calls_acc[i]
|
||
for i in sorted(tool_calls_acc)
|
||
if (
|
||
tool_calls_acc[i]
|
||
.get("function", {})
|
||
.get("name", "")
|
||
.strip()
|
||
)
|
||
] or None
|
||
if (
|
||
not tool_calls
|
||
and auto_heal_tool_calls
|
||
and any(s in content_accum for s in _TOOL_XML_SIGNALS)
|
||
):
|
||
tool_calls = self._parse_tool_calls_from_text(
|
||
content_accum,
|
||
)
|
||
if tool_calls and not has_structured_tc:
|
||
content_text = _strip_tool_markup(
|
||
content_text,
|
||
final = True,
|
||
)
|
||
if tool_calls:
|
||
logger.info(
|
||
f"Parsed {len(tool_calls)} tool call(s) from "
|
||
f"{'structured delta' if has_structured_tc else 'content text'}"
|
||
)
|
||
if not tool_calls:
|
||
# DRAINING but no tool calls (false positive).
|
||
# Merge accumulated metrics from prior tool
|
||
# iterations so they are not silently dropped.
|
||
yield {"type": "status", "text": ""}
|
||
if content_accum:
|
||
# Strip leaked tool-call XML before yielding
|
||
content_accum = _strip_tool_markup(
|
||
content_accum, final = True
|
||
)
|
||
if content_accum:
|
||
yield {"type": "content", "text": content_accum}
|
||
_fu = _iter_usage or {}
|
||
_fc = _fu.get("completion_tokens", 0)
|
||
_fp = _fu.get("prompt_tokens", 0)
|
||
_tc = _fc + _accumulated_completion_tokens
|
||
if (
|
||
_iter_usage
|
||
or _iter_timings
|
||
or _accumulated_completion_tokens
|
||
):
|
||
_mt = dict(_iter_timings) if _iter_timings else {}
|
||
if _accumulated_predicted_ms or _accumulated_predicted_n:
|
||
_mt["predicted_ms"] = (
|
||
_mt.get("predicted_ms", 0)
|
||
+ _accumulated_predicted_ms
|
||
)
|
||
_tn = (
|
||
_mt.get("predicted_n", 0) + _accumulated_predicted_n
|
||
)
|
||
_mt["predicted_n"] = _tn
|
||
_tms = _mt["predicted_ms"]
|
||
if _tms > 0:
|
||
_mt["predicted_per_second"] = _tn / (_tms / 1000.0)
|
||
yield {
|
||
"type": "metadata",
|
||
"usage": {
|
||
"prompt_tokens": _fp,
|
||
"completion_tokens": _tc,
|
||
"total_tokens": _fp + _tc,
|
||
},
|
||
"timings": _mt,
|
||
}
|
||
return
|
||
|
||
# ── Execute tool calls ──
|
||
_accumulated_completion_tokens += (_iter_usage or {}).get(
|
||
"completion_tokens", 0
|
||
)
|
||
_it = _iter_timings or {}
|
||
_accumulated_predicted_ms += _it.get("predicted_ms", 0)
|
||
_accumulated_predicted_n += _it.get("predicted_n", 0)
|
||
|
||
assistant_msg = {"role": "assistant", "content": content_text}
|
||
if tool_calls:
|
||
assistant_msg["tool_calls"] = tool_calls
|
||
conversation.append(assistant_msg)
|
||
|
||
for tc in tool_calls or []:
|
||
func = tc.get("function", {})
|
||
tool_name = func.get("name", "")
|
||
raw_args = func.get("arguments", {})
|
||
|
||
if isinstance(raw_args, str):
|
||
try:
|
||
arguments = json.loads(raw_args)
|
||
except (json.JSONDecodeError, ValueError):
|
||
if auto_heal_tool_calls:
|
||
arguments = {"query": raw_args}
|
||
else:
|
||
arguments = {"raw": raw_args}
|
||
else:
|
||
arguments = raw_args
|
||
|
||
if tool_name == "web_search":
|
||
status_text = f"Searching: {arguments.get('query', '')}"
|
||
elif tool_name == "python":
|
||
preview = (
|
||
(arguments.get("code") or "").strip().split("\n")[0][:60]
|
||
)
|
||
status_text = (
|
||
f"Running Python: {preview}"
|
||
if preview
|
||
else "Running Python..."
|
||
)
|
||
elif tool_name == "terminal":
|
||
cmd_preview = (arguments.get("command") or "")[:60]
|
||
status_text = (
|
||
f"Running: {cmd_preview}"
|
||
if cmd_preview
|
||
else "Running command..."
|
||
)
|
||
else:
|
||
status_text = f"Calling: {tool_name}"
|
||
yield {"type": "status", "text": status_text}
|
||
|
||
yield {
|
||
"type": "tool_start",
|
||
"tool_name": tool_name,
|
||
"tool_call_id": tc.get("id", ""),
|
||
"arguments": arguments,
|
||
}
|
||
|
||
# ── Duplicate call detection ──────────────
|
||
if _is_duplicate_call(tool_name, arguments):
|
||
result = (
|
||
"You already made this exact call. "
|
||
"Do not repeat the same tool call. "
|
||
"Try a different approach: fetch a URL "
|
||
"from previous results, use Python to "
|
||
"process data you already have, or "
|
||
"provide your final answer now."
|
||
)
|
||
else:
|
||
_effective_timeout = (
|
||
None if tool_call_timeout >= 9999 else tool_call_timeout
|
||
)
|
||
result = execute_tool(
|
||
tool_name,
|
||
arguments,
|
||
cancel_event = cancel_event,
|
||
timeout = _effective_timeout,
|
||
session_id = session_id,
|
||
)
|
||
|
||
yield {
|
||
"type": "tool_end",
|
||
"tool_name": tool_name,
|
||
"tool_call_id": tc.get("id", ""),
|
||
"result": result,
|
||
}
|
||
|
||
# Nudge model to try a different approach on errors
|
||
_error_prefixes = (
|
||
"Error",
|
||
"Search failed",
|
||
"Execution error",
|
||
"Blocked:",
|
||
"Exit code",
|
||
"Failed to fetch",
|
||
"Failed to resolve",
|
||
"No query provided",
|
||
)
|
||
_is_error = isinstance(result, str) and result.lstrip().startswith(
|
||
_error_prefixes
|
||
)
|
||
_record_tool_call(tool_name, arguments, failed = _is_error)
|
||
_result_content = result
|
||
if _is_error:
|
||
_result_content = (
|
||
result + "\n\nThe tool call encountered an issue. "
|
||
"Please try a different approach or rephrase your request."
|
||
)
|
||
|
||
tool_msg = {
|
||
"role": "tool",
|
||
"name": tool_name,
|
||
"content": _result_content,
|
||
}
|
||
tool_call_id = tc.get("id")
|
||
if tool_call_id:
|
||
tool_msg["tool_call_id"] = tool_call_id
|
||
conversation.append(tool_msg)
|
||
|
||
# Continue the loop to let model respond with context
|
||
continue
|
||
|
||
except httpx.ConnectError:
|
||
raise RuntimeError("Lost connection to llama-server")
|
||
except Exception as e:
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
return
|
||
raise
|
||
|
||
# ── Tool iteration cap reached -- synthesize final answer ──
|
||
# The model used all iterations without producing a final text
|
||
# response. Inject a nudge so the final streaming pass produces
|
||
# a useful answer instead of continuing to request tools.
|
||
if max_tool_iterations > 0:
|
||
conversation.append(
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
"You have used all available tool calls. Based on "
|
||
"everything you have found so far, provide your final "
|
||
"answer now. Do not call any more tools."
|
||
),
|
||
}
|
||
)
|
||
|
||
# Clear status
|
||
yield {"type": "status", "text": ""}
|
||
|
||
# Final streaming pass with the full conversation context
|
||
stream_payload = {
|
||
"messages": conversation,
|
||
"stream": True,
|
||
"temperature": temperature,
|
||
"top_p": top_p,
|
||
"top_k": top_k if top_k >= 0 else 0,
|
||
"min_p": min_p,
|
||
"repeat_penalty": repetition_penalty,
|
||
"presence_penalty": presence_penalty,
|
||
}
|
||
if self._supports_reasoning and enable_thinking is not None:
|
||
stream_payload["chat_template_kwargs"] = {
|
||
"enable_thinking": enable_thinking
|
||
}
|
||
if max_tokens is not None:
|
||
stream_payload["max_tokens"] = max_tokens
|
||
if stop:
|
||
stream_payload["stop"] = stop
|
||
stream_payload["stream_options"] = {"include_usage": True}
|
||
|
||
cumulative = ""
|
||
_last_emitted = ""
|
||
in_thinking = False
|
||
has_content_tokens = False
|
||
reasoning_text = ""
|
||
_metadata_usage = None
|
||
_metadata_timings = None
|
||
_stream_done = False
|
||
|
||
try:
|
||
stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10)
|
||
_auth_headers = (
|
||
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
||
)
|
||
with httpx.Client(timeout = stream_timeout) as client:
|
||
with self._stream_with_retry(
|
||
client,
|
||
url,
|
||
stream_payload,
|
||
cancel_event,
|
||
headers = _auth_headers,
|
||
) as response:
|
||
if response.status_code != 200:
|
||
error_body = response.read().decode()
|
||
raise RuntimeError(
|
||
f"llama-server returned {response.status_code}: {error_body}"
|
||
)
|
||
|
||
buffer = ""
|
||
for raw_chunk in self._iter_text_cancellable(
|
||
response, cancel_event
|
||
):
|
||
buffer += raw_chunk
|
||
while "\n" in buffer:
|
||
line, buffer = buffer.split("\n", 1)
|
||
line = line.strip()
|
||
|
||
if not line:
|
||
continue
|
||
if line == "data: [DONE]":
|
||
if in_thinking:
|
||
if has_content_tokens:
|
||
cumulative += "</think>"
|
||
yield {
|
||
"type": "content",
|
||
"text": _strip_tool_markup(
|
||
cumulative, final = True
|
||
),
|
||
}
|
||
else:
|
||
cumulative = reasoning_text
|
||
yield {"type": "content", "text": cumulative}
|
||
_stream_done = True
|
||
break # exit inner while
|
||
if not line.startswith("data: "):
|
||
continue
|
||
|
||
try:
|
||
chunk_data = json.loads(line[6:])
|
||
# Capture server timings/usage from final chunks
|
||
_chunk_timings = chunk_data.get("timings")
|
||
if _chunk_timings:
|
||
_metadata_timings = _chunk_timings
|
||
_chunk_usage = chunk_data.get("usage")
|
||
if _chunk_usage:
|
||
_metadata_usage = _chunk_usage
|
||
choices = chunk_data.get("choices", [])
|
||
if choices:
|
||
delta = choices[0].get("delta", {})
|
||
|
||
reasoning = delta.get("reasoning_content", "")
|
||
if reasoning:
|
||
reasoning_text += reasoning
|
||
if not in_thinking:
|
||
cumulative += "<think>"
|
||
in_thinking = True
|
||
cumulative += reasoning
|
||
yield {"type": "content", "text": cumulative}
|
||
|
||
token = delta.get("content", "")
|
||
if token:
|
||
has_content_tokens = True
|
||
if in_thinking:
|
||
cumulative += "</think>"
|
||
in_thinking = False
|
||
cumulative += token
|
||
cleaned = _strip_tool_markup(cumulative)
|
||
# Only emit when cleaned text grows (monotonic).
|
||
if len(cleaned) > len(_last_emitted):
|
||
_last_emitted = cleaned
|
||
yield {"type": "content", "text": cleaned}
|
||
except json.JSONDecodeError:
|
||
logger.debug(
|
||
f"Skipping malformed SSE line: {line[:100]}"
|
||
)
|
||
if _stream_done:
|
||
break # exit outer for
|
||
_final_usage = _metadata_usage or {}
|
||
_final_completion = _final_usage.get("completion_tokens", 0)
|
||
_final_prompt = _final_usage.get("prompt_tokens", 0)
|
||
_total_completion = (
|
||
_final_completion + _accumulated_completion_tokens
|
||
)
|
||
if _metadata_usage or _metadata_timings:
|
||
_merged_timings = (
|
||
dict(_metadata_timings) if _metadata_timings else {}
|
||
)
|
||
if _accumulated_predicted_ms or _accumulated_predicted_n:
|
||
_merged_timings["predicted_ms"] = (
|
||
_merged_timings.get("predicted_ms", 0)
|
||
+ _accumulated_predicted_ms
|
||
)
|
||
_total_predicted_n = (
|
||
_merged_timings.get("predicted_n", 0)
|
||
+ _accumulated_predicted_n
|
||
)
|
||
_merged_timings["predicted_n"] = _total_predicted_n
|
||
_total_predicted_ms = _merged_timings["predicted_ms"]
|
||
if _total_predicted_ms > 0:
|
||
_merged_timings["predicted_per_second"] = (
|
||
_total_predicted_n / (_total_predicted_ms / 1000.0)
|
||
)
|
||
yield {
|
||
"type": "metadata",
|
||
"usage": {
|
||
"prompt_tokens": _final_prompt,
|
||
"completion_tokens": _total_completion,
|
||
"total_tokens": _final_prompt + _total_completion,
|
||
},
|
||
"timings": _merged_timings,
|
||
}
|
||
|
||
except httpx.ConnectError:
|
||
raise RuntimeError("Lost connection to llama-server")
|
||
except Exception as e:
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
return
|
||
raise
|
||
|
||
# ── TTS support ────────────────────────────────────────────
|
||
|
||
def detect_audio_type(self) -> Optional[str]:
|
||
"""Detect audio/TTS codec by probing the loaded model's vocabulary."""
|
||
if not self.is_loaded:
|
||
return None
|
||
try:
|
||
_auth_headers = (
|
||
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
||
)
|
||
with httpx.Client(timeout = 10, headers = _auth_headers) as client:
|
||
|
||
def _detok(tid: int) -> str:
|
||
r = client.post(
|
||
f"{self.base_url}/detokenize", json = {"tokens": [tid]}
|
||
)
|
||
return r.json().get("content", "") if r.status_code == 200 else ""
|
||
|
||
def _tok(text: str) -> list[int]:
|
||
r = client.post(
|
||
f"{self.base_url}/tokenize",
|
||
json = {"content": text, "add_special": False},
|
||
)
|
||
return r.json().get("tokens", []) if r.status_code == 200 else []
|
||
|
||
# Check codec-specific tokens (not generic ones that may exist in non-audio models)
|
||
if "<custom_token_" in _detok(128258) and "<custom_token_" in _detok(
|
||
128259
|
||
):
|
||
return "snac"
|
||
if len(_tok("<|AUDIO|>")) == 1 and len(_tok("<|audio_eos|>")) == 1:
|
||
return "csm"
|
||
if len(_tok("<|startoftranscript|>")) == 1:
|
||
return "whisper"
|
||
if (
|
||
len(_tok("<|bicodec_semantic_0|>")) == 1
|
||
and len(_tok("<|bicodec_global_0|>")) == 1
|
||
):
|
||
return "bicodec"
|
||
if len(_tok("<|c1_0|>")) == 1 and len(_tok("<|c2_0|>")) == 1:
|
||
return "dac"
|
||
except Exception as e:
|
||
logger.debug(f"Audio type detection failed: {e}")
|
||
return None
|
||
|
||
# Prompt format per codec: (template, stop_tokens, needs_token_ids)
|
||
# Matches prompts in InferenceBackend._generate_snac/bicodec/dac
|
||
_TTS_PROMPTS = {
|
||
"snac": (
|
||
"<custom_token_3>{text}<|eot_id|><custom_token_4>",
|
||
["<custom_token_2>"],
|
||
True,
|
||
),
|
||
"bicodec": (
|
||
"<|task_tts|><|start_content|>{text}<|end_content|><|start_global_token|>",
|
||
["<|im_end|>", "</s>"],
|
||
False,
|
||
),
|
||
"dac": (
|
||
"<|im_start|>\n<|text_start|>{text}<|text_end|>\n<|audio_start|><|global_features_start|>\n",
|
||
["<|im_end|>", "<|audio_end|>"],
|
||
False,
|
||
),
|
||
}
|
||
|
||
_codec_mgr = None # Shared AudioCodecManager instance
|
||
|
||
def init_audio_codec(self, audio_type: str) -> None:
|
||
"""Load the audio codec at model load time (mirrors non-GGUF path)."""
|
||
import torch
|
||
from core.inference.audio_codecs import AudioCodecManager
|
||
|
||
if LlamaCppBackend._codec_mgr is None:
|
||
LlamaCppBackend._codec_mgr = AudioCodecManager()
|
||
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
model_repo_path = None
|
||
|
||
# BiCodec needs a repo with BiCodec/ weights — download canonical SparkTTS
|
||
if audio_type == "bicodec":
|
||
from huggingface_hub import snapshot_download
|
||
import os
|
||
|
||
repo_path = snapshot_download(
|
||
"unsloth/Spark-TTS-0.5B", local_dir = "Spark-TTS-0.5B"
|
||
)
|
||
model_repo_path = os.path.abspath(repo_path)
|
||
|
||
LlamaCppBackend._codec_mgr.load_codec(
|
||
audio_type, device, model_repo_path = model_repo_path
|
||
)
|
||
logger.info(f"Loaded audio codec for GGUF TTS: {audio_type}")
|
||
|
||
def generate_audio_response(
|
||
self,
|
||
text: str,
|
||
audio_type: str,
|
||
temperature: float = 0.6,
|
||
top_p: float = 0.95,
|
||
top_k: int = 50,
|
||
min_p: float = 0.0,
|
||
max_new_tokens: int = 2048,
|
||
repetition_penalty: float = 1.1,
|
||
) -> tuple:
|
||
"""
|
||
Generate TTS audio via llama-server /completion + codec decoding.
|
||
Returns (wav_bytes, sample_rate).
|
||
"""
|
||
if audio_type not in self._TTS_PROMPTS:
|
||
raise RuntimeError(f"GGUF TTS does not support '{audio_type}' codec.")
|
||
|
||
tpl, stop, need_ids = self._TTS_PROMPTS[audio_type]
|
||
|
||
payload: dict = {
|
||
"prompt": tpl.format(text = text),
|
||
"stream": False,
|
||
"n_predict": max_new_tokens,
|
||
"temperature": temperature,
|
||
"top_p": top_p,
|
||
"top_k": top_k if top_k >= 0 else 0,
|
||
"min_p": min_p,
|
||
"repeat_penalty": repetition_penalty,
|
||
}
|
||
if stop:
|
||
payload["stop"] = stop
|
||
if need_ids:
|
||
payload["n_probs"] = 1
|
||
|
||
_auth_headers = (
|
||
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
||
)
|
||
with httpx.Client(
|
||
timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers
|
||
) as client:
|
||
resp = client.post(f"{self.base_url}/completion", json = payload)
|
||
if resp.status_code != 200:
|
||
raise RuntimeError(
|
||
f"llama-server returned {resp.status_code}: {resp.text}"
|
||
)
|
||
|
||
data = resp.json()
|
||
token_ids = (
|
||
[p["id"] for p in data.get("completion_probabilities", []) if "id" in p]
|
||
if need_ids
|
||
else None
|
||
)
|
||
|
||
import torch
|
||
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
return LlamaCppBackend._codec_mgr.decode(
|
||
audio_type, device, token_ids = token_ids, text = data.get("content", "")
|
||
)
|