diff --git a/src/fabledassistant/app.py b/src/fabledassistant/app.py index 985d56e..3e5d07c 100644 --- a/src/fabledassistant/app.py +++ b/src/fabledassistant/app.py @@ -183,16 +183,18 @@ def create_app() -> Quart: """Send a minimal chat request to prime Ollama's KV cache with the user's system prompt. This ensures the next real user message only needs to process its own tokens - rather than the full ~5,600-token system prompt, cutting TTFT from ~25s to <1s. + rather than the full ~4,650-token system prompt, cutting TTFT from ~25s to <1s. + The num_ctx must match what real requests will use so Ollama doesn't reload. """ try: - from fabledassistant.services.llm import build_context + from fabledassistant.services.llm import build_context, pick_num_ctx messages, _ = await build_context( user_id=user_id, history=[], current_note_id=None, user_message=" ", ) + num_ctx = pick_num_ctx(messages) async with httpx.AsyncClient(timeout=120.0) as client: await client.post( f"{Config.OLLAMA_URL}/api/chat", @@ -200,11 +202,11 @@ def create_app() -> Quart: "model": model, "messages": messages, "stream": False, - "options": {"num_predict": 1}, + "options": {"num_predict": 1, "num_ctx": num_ctx}, "keep_alive": "2h", }, ) - logger.info("Primed KV cache for user %d with model '%s'", user_id, model) + logger.info("Primed KV cache for user %d with model '%s' (num_ctx=%d)", user_id, model, num_ctx) except Exception: logger.warning("Failed to prime KV cache for user %d", user_id, exc_info=True) diff --git a/src/fabledassistant/routes/settings.py b/src/fabledassistant/routes/settings.py index a180fb2..979075e 100644 --- a/src/fabledassistant/routes/settings.py +++ b/src/fabledassistant/routes/settings.py @@ -15,7 +15,7 @@ logger = logging.getLogger(__name__) async def _prime_kv_cache_bg(user_id: int, model: str) -> None: """Fire-and-forget: prime Ollama's KV cache with the user's system prompt.""" import httpx - from fabledassistant.services.llm import build_context + from fabledassistant.services.llm import build_context, pick_num_ctx try: messages, _ = await build_context( user_id=user_id, @@ -23,6 +23,7 @@ async def _prime_kv_cache_bg(user_id: int, model: str) -> None: current_note_id=None, user_message=" ", ) + num_ctx = pick_num_ctx(messages) async with httpx.AsyncClient(timeout=120.0) as client: await client.post( f"{Config.OLLAMA_URL}/api/chat", @@ -30,11 +31,11 @@ async def _prime_kv_cache_bg(user_id: int, model: str) -> None: "model": model, "messages": messages, "stream": False, - "options": {"num_predict": 1}, + "options": {"num_predict": 1, "num_ctx": num_ctx}, "keep_alive": "2h", }, ) - logger.info("Primed KV cache for user %d with model '%s'", user_id, model) + logger.info("Primed KV cache for user %d with model '%s' (num_ctx=%d)", user_id, model, num_ctx) except Exception: logger.warning("Failed to prime KV cache for user %d", user_id, exc_info=True) diff --git a/src/fabledassistant/services/generation_task.py b/src/fabledassistant/services/generation_task.py index f757b97..048fafd 100644 --- a/src/fabledassistant/services/generation_task.py +++ b/src/fabledassistant/services/generation_task.py @@ -22,7 +22,7 @@ from fabledassistant.services.generation_buffer import ( GenerationBuffer, GenerationState, ) -from fabledassistant.services.llm import ChatChunk, build_context, generate_completion, stream_chat, stream_chat_with_tools, summarize_history_for_context +from fabledassistant.services.llm import ChatChunk, build_context, generate_completion, pick_num_ctx, stream_chat, stream_chat_with_tools, summarize_history_for_context from fabledassistant.services.chat import update_conversation_title from fabledassistant.services.settings import get_setting from fabledassistant.services.logging import log_generation @@ -168,6 +168,7 @@ async def _stream_with_retry( model: str, tools: list[dict], think: bool, + num_ctx: int | None = None, ) -> AsyncGenerator[ChatChunk, None]: """stream_chat_with_tools with automatic retry on Ollama 500 errors. @@ -184,7 +185,7 @@ async def _stream_with_retry( ) await asyncio.sleep(delay) try: - async for chunk in stream_chat_with_tools(messages, model, tools=tools, think=think): + async for chunk in stream_chat_with_tools(messages, model, tools=tools, think=think, num_ctx=num_ctx): yield chunk return except httpx.HTTPStatusError as exc: @@ -260,6 +261,11 @@ async def run_generation( voice_speech_style=voice_speech_style, ) + # Pick the smallest context tier that fits the current messages. + # Using the minimum needed tier reduces KV cache size and speeds up prefill. + num_ctx = pick_num_ctx(messages) + logger.debug("Adaptive num_ctx=%d for conv %d", num_ctx, conv_id) + # Emit context event buf.append_event("context", {"context": context_meta}) @@ -269,6 +275,7 @@ async def run_generation( t_start = time.monotonic() timing: dict = { "think": think, + "num_ctx": num_ctx, "tools": [], "rounds": 0, "prompt_tokens": None, @@ -298,7 +305,7 @@ async def run_generation( buf.append_event("status", {"status": "Generating response..." if _round == 0 else "Composing response..."}) t_stream = time.monotonic() - async for chunk in _stream_with_retry(messages, model, tools, think): + async for chunk in _stream_with_retry(messages, model, tools, think, num_ctx=num_ctx): if buf.cancel_event.is_set(): cancelled = True break @@ -535,7 +542,7 @@ async def run_assist_generation( await asyncio.sleep(delay) try: buf.content_so_far = "" - async for chunk in stream_chat(messages, model, options={"num_predict": Config.OLLAMA_NUM_CTX}): + async for chunk in stream_chat(messages, model, options={"num_predict": num_ctx}, num_ctx=num_ctx): buf.content_so_far += chunk buf.append_event("chunk", {"chunk": chunk}) diff --git a/src/fabledassistant/services/llm.py b/src/fabledassistant/services/llm.py index fdb924c..67abd71 100644 --- a/src/fabledassistant/services/llm.py +++ b/src/fabledassistant/services/llm.py @@ -19,6 +19,28 @@ from fabledassistant.services.settings import get_setting logger = logging.getLogger(__name__) +# Context window tiers. The smallest tier that fits the current input is used +# so Ollama allocates a smaller KV cache, reducing prefill time and VRAM usage. +# Requests using the same tier hit Ollama's prefix cache; a tier upgrade causes +# a one-time model reload but then the larger cache stays warm. +_CTX_TIERS = (8192, 16384, 32768) + + +def pick_num_ctx(messages: list[dict]) -> int: + """Return the smallest context tier that fits *messages* with 25% headroom. + + Stays at or below Config.OLLAMA_NUM_CTX (the configured ceiling). + """ + total_chars = sum(len(m.get("content") or "") for m in messages) + estimated_tokens = int(total_chars / 3.5) + needed = int(estimated_tokens * 1.25) + 256 # 25% headroom + output buffer + cap = Config.OLLAMA_NUM_CTX + for tier in _CTX_TIERS: + if tier >= needed and tier <= cap: + return tier + return cap + + STOP_WORDS = frozenset({ "a", "an", "the", "is", "it", "to", "in", "for", "of", "and", "or", "on", "at", "by", "with", "from", "as", "be", "was", "were", "been", @@ -112,6 +134,7 @@ async def stream_chat( model: str, options: dict | None = None, think: bool = False, + num_ctx: int | None = None, ) -> AsyncGenerator[str, None]: """Stream chat completion from Ollama, yielding content chunks. @@ -119,7 +142,7 @@ async def stream_chat( Thinking tokens are silently discarded anyway, but disabling avoids the multi-minute delay before the first content token arrives. """ - merged_options = {"num_ctx": Config.OLLAMA_NUM_CTX} + merged_options = {"num_ctx": num_ctx or Config.OLLAMA_NUM_CTX} if options: merged_options.update(options) payload: dict = {"model": model, "messages": messages, "stream": True, "options": merged_options, "think": think, "keep_alive": "2h"} @@ -159,6 +182,7 @@ async def stream_chat_with_tools( model: str, tools: list[dict] | None = None, think: bool = False, + num_ctx: int | None = None, ) -> AsyncGenerator[ChatChunk, None]: """Stream chat completion from Ollama with tool support. @@ -170,7 +194,8 @@ async def stream_chat_with_tools( Thinking tokens are consumed by Ollama and not forwarded to the caller; only the final response content is yielded. Expect higher TTFT when enabled. """ - options: dict = {"num_ctx": Config.OLLAMA_NUM_CTX} + resolved_ctx = num_ctx or Config.OLLAMA_NUM_CTX + options: dict = {"num_ctx": resolved_ctx} if tools: options["num_predict"] = 8192 payload: dict = {