bvandeusen/FabledScribe
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 50 Wiki Activity

Implement optimistic streaming to eliminate intent classification latency

Browse Source 
Start the main LLM stream immediately after build_context finishes instead
of waiting for intent classification to complete. Race the two concurrently:

- Intent wins before first token → cancel stream, execute tool (tool path
  unchanged: confirmation, acknowledgment, multi-round loop all preserved)
- First token wins → discard intent, user sees output immediately

For pure chat messages (no tool needed, the common case) this eliminates
the full intent classification RTT from TTFT. For tool calls, intent
typically wins the race since it finishes before the main model produces
its first token, so tool behaviour is unchanged in practice.

Also extracts _drain_queue() as a module-level async generator helper.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Bryan Van Deusen
2026-02-25 23:03:30 -05:00
parent 8ee3649531
commit 98d3fca277
1 changed files with 189 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/fabledassistant/services/generation_task.py
+189 -59
View File
@@ -9,6 +9,7 @@ import json
import logging
import re
import time
from collections.abc import AsyncGenerator
from sqlalchemy import update
@@ -172,6 +173,26 @@ async def _update_message(
await session.commit()
async def _drain_queue(
prefetched: list[ChatChunk],
queue: asyncio.Queue,
) -> AsyncGenerator[ChatChunk, None]:
"""Yield pre-fetched chunks then drain remaining chunks from the queue.
A None sentinel in the queue signals the stream is finished.
A BaseException in the queue is re-raised so callers see the error.
"""
for chunk in prefetched:
yield chunk
while True:
item = await queue.get()
if item is None:
break
if isinstance(item, BaseException):
raise item
yield item
async def run_generation(
buf: GenerationBuffer,
history: list[dict],
@@ -202,47 +223,37 @@ async def run_generation(
)
# Phase 2: Summarize long conversation history if needed.
# For short conversations (<= threshold) this returns immediately with no LLM call.
# For long conversations it spends ~500ms to compress old exchanges, which
# more than pays for itself by reducing the prefill token count every turn.
history_to_use = history
history_summary: str | None = None
if len(history) > 20: # matches _HISTORY_SUMMARY_THRESHOLD in llm.py
buf.append_event("status", {"status": "Summarizing conversation history..."})
history_to_use, history_summary = await summarize_history_for_context(history, intent_model)
# Phase 3: Build context and classify intent in parallel — the two slow legs.
# Pass cached note IDs so build_context can reuse them, keeping the system
# prompt prefix stable and helping Ollama's KV cache stay warm.
# Phase 3: Build context. Start intent classification concurrently when tools
# are available so it runs in parallel with the embedding/DB work in build_context.
# We only block on context (need messages to stream) — intent result is consumed
# later via a race with the first streaming token.
cached_note_ids = get_conv_note_cache(conv_id) or None
pre_intent: IntentResult = IntentResult()
intent_timing_ms: int | None = None
context_task = asyncio.create_task(build_context(
user_id, history_to_use, context_note_id, user_content,
exclude_note_ids=exclude_note_ids,
history_summary=history_summary,
cached_note_ids=cached_note_ids,
))
intent_task: asyncio.Task[IntentResult] | None = None
t_intent = time.monotonic()
if tools:
intent_history = [
m for m in history_to_use
if m.get("role") in ("user", "assistant") and m.get("content")
][-6:]
buf.append_event("status", {"status": "Analyzing your request..."})
t_intent = time.monotonic()
context_task = asyncio.create_task(build_context(
user_id, history_to_use, context_note_id, user_content,
exclude_note_ids=exclude_note_ids,
history_summary=history_summary,
cached_note_ids=cached_note_ids,
))
intent_task = asyncio.create_task(
classify_intent(user_content, tools, intent_model, history=intent_history)
)
(messages, context_meta), pre_intent = await asyncio.gather(context_task, intent_task)
intent_timing_ms = int((time.monotonic() - t_intent) * 1000)
else:
messages, context_meta = await build_context(
user_id, history_to_use, context_note_id, user_content,
exclude_note_ids=exclude_note_ids,
history_summary=history_summary,
cached_note_ids=cached_note_ids,
)
messages, context_meta = await context_task
# Update the note cache with whatever notes ended up in context.
new_note_ids = context_meta.get("auto_note_ids") or []
@@ -254,7 +265,7 @@ async def run_generation(
t_start = time.monotonic()
timing: dict = {
"intent_ms": intent_timing_ms,
"intent_ms": None,
"tools": [],
"ttft_ms": None,
"generation_ms": None,
@@ -271,25 +282,73 @@ async def run_generation(
round_tool_calls: list[dict] = []
logger.info("Generation round %d started for conv %d (model=%s)", _round, conv_id, model)
# Intent routing — first round only (result pre-computed in parallel with build_context)
if _round == 0 and tools:
intent = pre_intent
if intent.should_execute:
logger.info(
"Intent router detected tool (confidence=%s): %s(%s)",
intent.confidence, intent.tool_name, json.dumps(intent.arguments)[:200],
)
elif intent.tool_name:
logger.info(
"Intent router low confidence (%s) for tool=%s — falling through to streaming",
intent.confidence, intent.tool_name,
)
if intent.should_execute:
# --- Round 0 with tools: optimistic streaming ---
# Start the main stream immediately (into a queue) while the intent
# classifier finishes in the background. Race the two:
# • Intent wins before first token → check for tool call, cancel stream if needed
# • First token wins → discard intent, stream has already started
if _round == 0 and tools and intent_task is not None:
stream_queue: asyncio.Queue[ChatChunk | BaseException | None] = asyncio.Queue(maxsize=256)
async def _fill_queue() -> None:
try:
async for c in stream_chat_with_tools(messages, model, tools=tools, think=think):
await stream_queue.put(c)
except BaseException as exc:
await stream_queue.put(exc)
finally:
await stream_queue.put(None)
stream_fill_task = asyncio.create_task(_fill_queue())
buf.append_event("status", {"status": "Generating response..."})
queue_peek = asyncio.create_task(stream_queue.get())
race_done, _ = await asyncio.wait(
{intent_task, queue_peek},
return_when=asyncio.FIRST_COMPLETED,
)
intent = IntentResult()
prefetched: list[ChatChunk] = []
use_tool_path = False
if intent_task in race_done and queue_peek not in race_done:
# Intent finished before the first streaming token arrived.
timing["intent_ms"] = int((time.monotonic() - t_intent) * 1000)
intent = intent_task.result()
if intent.should_execute:
# Cancel the optimistic stream — we'll execute a tool instead.
stream_fill_task.cancel()
queue_peek.cancel()
use_tool_path = True
else:
# No tool needed — collect the first chunk and keep streaming.
first = await queue_peek
if isinstance(first, BaseException):
raise first
if first is not None:
prefetched.append(first)
else:
# Stream produced a token before intent finished (or simultaneously).
# Discard intent — the response is already on its way.
if not intent_task.done():
intent_task.cancel()
else:
timing["intent_ms"] = int((time.monotonic() - t_intent) * 1000)
first = queue_peek.result() if queue_peek in race_done else await queue_peek
if isinstance(first, BaseException):
raise first
if first is not None:
prefetched.append(first)
if use_tool_path:
# === Tool path (intent won the race) ===
tool_name = intent.tool_name
confirmed = True # Non-write tools auto-confirm
confirmed = True
if tool_name in _WRITE_TOOLS:
# Pause and ask the user to accept or decline before executing.
loop = asyncio.get_running_loop()
confirm_future: asyncio.Future = loop.create_future()
buf.confirmation_future = confirm_future
@@ -316,7 +375,6 @@ async def run_generation(
confirmed = bool(confirm_future.result())
except Exception:
confirmed = False
# else: timeout → confirmed stays False
except Exception:
confirmed = False
finally:
@@ -326,7 +384,6 @@ async def run_generation(
if not confirmed:
if not cancelled:
# Record the declined action so the UI can show it
declined_record = {
"function": tool_name,
"arguments": intent.arguments,
@@ -335,15 +392,10 @@ async def run_generation(
}
all_tool_calls.append(declined_record)
buf.append_event("tool_call", {"tool_call": declined_record})
# Fall through to streaming without tool context
if confirmed:
buf.append_event("status", {"status": f"{_TOOL_LABELS.get(tool_name, 'Working')}..."})
# Run tool execution and acknowledgment generation in parallel.
# The acknowledgment uses the fast intent model (already in VRAM),
# so the user sees text within ~200-400ms instead of waiting for
# the full main-model TTFT (~22s).
t_tool = time.monotonic()
result, ack_text = await asyncio.gather(
execute_tool(user_id, tool_name, intent.arguments),
@@ -352,15 +404,12 @@ async def run_generation(
timing["tools"].append({"name": tool_name, "ms": int((time.monotonic() - t_tool) * 1000)})
logger.info("Intent-routed tool %s result: success=%s", tool_name, result.get("success"))
# Stream acknowledgment immediately — user sees text before main LLM starts
if ack_text:
buf.append_event("chunk", {"chunk": ack_text})
buf.content_so_far += ack_text
if timing["ttft_ms"] is None:
timing["ttft_ms"] = int((time.monotonic() - t_start) * 1000)
# Invalidate the note context cache after any successful note write
# so the next turn can pick up newly created/modified notes.
if result.get("success") and tool_name in {"create_task", "create_note", "update_note", "delete_note", "delete_task"}:
clear_conv_note_cache(conv_id)
@@ -373,8 +422,6 @@ async def run_generation(
all_tool_calls.append(tool_record)
buf.append_event("tool_call", {"tool_call": tool_record})
# Include ack as the assistant's partial response so round 1
# continues coherently from where the acknowledgment left off
messages.append({
"role": "assistant",
"content": ack_text,
@@ -386,9 +433,95 @@ async def run_generation(
"role": "tool",
"content": json.dumps(result),
})
continue # Next round: LLM streams response incorporating result
continue # Round 1: stream the response incorporating tool result
# Bail out here if cancelled during confirmation wait
# Declined write tool — fall through to stream a fresh response.
if cancelled:
break
# === Stream path for round 0 ===
# Either intent said no-tool, or a write tool was declined.
# For no-tool: drain the already-started queue (prefetched + remaining).
# For declined: start a fresh stream (queue was cancelled).
buf.append_event("status", {"status": "Generating response..."})
t_stream = time.monotonic()
if use_tool_path:
# Declined write tool — the optimistic stream was cancelled, start fresh.
stream_source: AsyncGenerator = stream_chat_with_tools(messages, model, tools=tools, think=think)
else:
stream_source = _drain_queue(prefetched, stream_queue)
async for chunk in stream_source:
if buf.cancel_event.is_set():
cancelled = True
break
if chunk.type == "content":
if timing["ttft_ms"] is None:
timing["ttft_ms"] = int((time.monotonic() - t_start) * 1000)
buf.content_so_far += chunk.content
clean = _TOOL_CALL_MARKER.sub("", chunk.content)
if clean:
buf.append_event("chunk", {"chunk": clean})
now = time.monotonic()
if now - last_flush >= DB_FLUSH_INTERVAL:
try:
await _update_message(msg_id, buf.content_so_far, "generating")
except Exception:
logger.warning("Failed periodic flush for message %d", msg_id, exc_info=True)
last_flush = now
elif chunk.type == "tool_calls" and chunk.tool_calls:
logger.info("Round %d: model returned %d tool call(s)", _round, len(chunk.tool_calls))
for tc in chunk.tool_calls:
fn = tc.get("function", {})
tool_name = fn.get("name", "")
arguments = fn.get("arguments", {})
logger.info("Executing tool: %s(%s)", tool_name, json.dumps(arguments)[:200])
buf.append_event("status", {"status": f"{_TOOL_LABELS.get(tool_name, 'Working')}..."})
t_tool = time.monotonic()
result = await execute_tool(user_id, tool_name, arguments)
timing["tools"].append({"name": tool_name, "ms": int((time.monotonic() - t_tool) * 1000)})
logger.info("Tool %s result: success=%s", tool_name, result.get("success"))
if result.get("success") and tool_name in {"create_task", "create_note", "update_note", "delete_note", "delete_task"}:
clear_conv_note_cache(conv_id)
tool_record = {
"function": tool_name,
"arguments": arguments,
"result": result,
"status": "success" if result.get("success") else "error",
}
round_tool_calls.append(tool_record)
all_tool_calls.append(tool_record)
buf.append_event("tool_call", {"tool_call": tool_record})
timing["generation_ms"] = int((time.monotonic() - t_stream) * 1000)
if cancelled:
break
if not round_tool_calls:
break
buf.content_so_far = _TOOL_CALL_MARKER.sub("", buf.content_so_far)
messages.append({
"role": "assistant",
"content": buf.content_so_far,
"tool_calls": [
{"function": {"name": tc["function"], "arguments": tc["arguments"]}}
for tc in round_tool_calls
],
})
for tc in round_tool_calls:
messages.append({"role": "tool", "content": json.dumps(tc["result"])})
buf.content_so_far = ""
continue
# --- Rounds 1+ (and round 0 with no tools) ---
if cancelled:
break
@@ -419,7 +552,6 @@ async def run_generation(
elif chunk.type == "tool_calls" and chunk.tool_calls:
logger.info("Round %d: model returned %d tool call(s)", _round, len(chunk.tool_calls))
# Process each tool call
for tc in chunk.tool_calls:
fn = tc.get("function", {})
tool_name = fn.get("name", "")
@@ -443,8 +575,6 @@ async def run_generation(
}
round_tool_calls.append(tool_record)
all_tool_calls.append(tool_record)
# Emit tool_call SSE event
buf.append_event("tool_call", {"tool_call": tool_record})
timing["generation_ms"] = int((time.monotonic() - t_stream) * 1000)