diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 42cce3c4d8..8ea7e1e816 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -10,6 +10,7 @@ through its OpenAI-compatible /v1/chat/completions endpoint. import atexit import json +import struct import structlog from loggers import get_logger import shutil @@ -45,6 +46,8 @@ class LlamaCppBackend: self._hf_variant: Optional[str] = None self._is_vision: bool = False self._healthy = False + self._context_length: Optional[int] = None + self._chat_template: Optional[str] = None self._lock = threading.Lock() self._stdout_lines: list[str] = [] self._stdout_thread: Optional[threading.Thread] = None @@ -80,6 +83,14 @@ class LlamaCppBackend: def hf_variant(self) -> Optional[str]: return self._hf_variant + @property + def context_length(self) -> Optional[int]: + return self._context_length + + @property + def chat_template(self) -> Optional[str]: + return self._chat_template + # ── Binary discovery ────────────────────────────────────────── @staticmethod @@ -371,6 +382,85 @@ class LlamaCppBackend: # 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(" None: + """Read context_length and chat_template from a GGUF file's KV header. + + Parses only the KV pairs we need (~30ms even for multi-GB files). + For split GGUFs, metadata is always in shard 1. + """ + try: + WANTED = {"general.architecture", "tokenizer.chat_template"} + arch = None + ctx_key = None + + with open(gguf_path, "rb") as f: + magic = struct.unpack(" { async function autoLoadSmallestModel(): Promise { const toastId = toast("Loading a model…", { description: "Auto-selecting the smallest downloaded model.", - duration: Infinity, + duration: 5000, + closeButton: true, }); try { const [ggufRepos, modelRepos] = await Promise.all([ @@ -214,7 +215,7 @@ async function autoLoadSmallestModel(): Promise { .sort((a, b) => a.size_bytes - b.size_bytes); if (downloaded.length > 0) { const variant = downloaded[0]; - await loadModel({ + const loadResp = await loadModel({ model_path: repo.repo_id, hf_token: null, max_seq_length: 4096, @@ -223,7 +224,9 @@ async function autoLoadSmallestModel(): Promise { gguf_variant: variant.quant, trust_remote_code: false, }); - useChatRuntimeStore.getState().setCheckpoint(repo.repo_id, variant.quant); + const store = useChatRuntimeStore.getState(); + store.setCheckpoint(repo.repo_id, variant.quant); + store.setParams({ ...store.params, maxTokens: loadResp.context_length ?? 131072 }); toast.success(`Loaded ${repo.repo_id} (${variant.quant})`, { id: toastId }); return true; } @@ -247,7 +250,9 @@ async function autoLoadSmallestModel(): Promise { gguf_variant: null, trust_remote_code: false, }); - useChatRuntimeStore.getState().setCheckpoint(repo.repo_id); + const store = useChatRuntimeStore.getState(); + store.setCheckpoint(repo.repo_id); + store.setParams({ ...store.params, maxTokens: 4096 }); toast.success(`Loaded ${repo.repo_id}`, { id: toastId }); return true; } catch { diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index bb2b05f9da..302bc17335 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -340,10 +340,10 @@ export function ChatSettingsPanel({ label="Max Tokens" value={params.maxTokens} min={64} - max={isGguf ? 131072 : 32768} - step={64} + max={isGguf ? params.maxTokens : 32768} + step={isGguf ? params.maxTokens : 64} onChange={set("maxTokens")} - displayValue={isGguf && params.maxTokens >= 131072 ? "Max" : undefined} + displayValue={isGguf ? "Max" : undefined} /> diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 47bd0698a8..b2ec407d4c 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -132,9 +132,11 @@ function mergeRecommendedInference( modelId: string, ): InferenceParams { const inference = response.inference; - // GGUF: max tokens = 131072 (effectively unlimited, model decides) - // Non-GGUF: max tokens = 4096 - const defaultMaxTokens = response.is_gguf ? 131072 : 4096; + // GGUF: use actual context length from GGUF metadata, fallback to 131072 + // Non-GGUF: 4096 + const defaultMaxTokens = response.is_gguf + ? (response.context_length ?? 131072) + : 4096; return { ...current, checkpoint: modelId, diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 4046deb4c2..54ce25c1e2 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -82,6 +82,7 @@ export interface LoadModelResponse { min_p?: number; trust_remote_code?: boolean; }; + context_length?: number | null; } export interface UnloadModelRequest {