diff --git a/src/fabledassistant/services/generation_task.py b/src/fabledassistant/services/generation_task.py index 17d8185..ee2bb6a 100644 --- a/src/fabledassistant/services/generation_task.py +++ b/src/fabledassistant/services/generation_task.py @@ -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)