"""Background asyncio task for LLM generation. Streams from Ollama into a GenerationBuffer, periodically flushing to DB. Runs independently of any HTTP connection. """ import asyncio import json import logging import re import time from collections.abc import AsyncGenerator import httpx from sqlalchemy import update from fabledassistant.config import Config from fabledassistant.models import async_session from fabledassistant.models.conversation import Message 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, wait_for_model_loaded from fabledassistant.services.chat import update_conversation_title from fabledassistant.services.logging import log_generation from fabledassistant.services.tools import get_tools_for_user, execute_tool from fabledassistant.services.research import run_research_pipeline logger = logging.getLogger(__name__) # Mistral prefixes tool-call responses with "[TOOL_CALLS]" as visible text _TOOL_CALL_MARKER = re.compile(r"^\s*\[TOOL_CALLS\]\s*", re.IGNORECASE) DB_FLUSH_INTERVAL = 5.0 # seconds between partial DB flushes # Human-readable labels for each tool, shown in the status indicator _TOOL_LABELS: dict[str, str] = { "create_task": "Creating task", "create_note": "Creating note", "update_note": "Updating note", "delete_note": "Deleting note", "delete_task": "Deleting task", "get_note": "Reading note", "list_notes": "Listing notes", "list_tasks": "Searching tasks", "search_notes": "Searching notes (semantic)", "create_event": "Creating calendar event", "list_events": "Searching calendar", "search_events": "Searching calendar", "update_event": "Updating calendar event", "delete_event": "Removing calendar event", "list_calendars": "Listing calendars", "search_web": "Searching the web", "research_topic": "Researching topic", } async def _generate_title(messages: list[dict], model: str) -> str: """Ask the LLM for a concise conversation title.""" # Build conversation text like summarize_conversation_as_note conv_lines = [] for m in messages: if m["role"] == "system": continue label = "User" if m["role"] == "user" else "Assistant" conv_lines.append(f"{label}: {m['content']}") # Keep only last 6 pairs worth of text conv_lines = conv_lines[-12:] prompt_messages = [ { "role": "system", "content": ( "Generate a concise 3-8 word title for this conversation. " "Reply with ONLY the title, no quotes or punctuation." ), }, {"role": "user", "content": "\n\n".join(conv_lines)}, ] title = await generate_completion(prompt_messages, model, max_tokens=30) title = title.strip().strip('"\'').strip() return title[:100] if title else "" async def _update_message( message_id: int, content: str, status: str, tool_calls: list[dict] | None = None, ) -> None: values: dict = {"content": content, "status": status} if tool_calls is not None: values["tool_calls"] = tool_calls async with async_session() as session: await session.execute( update(Message) .where(Message.id == message_id) .values(**values) ) await session.commit() async def _stream_with_retry( messages: list[dict], model: str, tools: list[dict], think: bool, ) -> AsyncGenerator[ChatChunk, None]: """stream_chat_with_tools with automatic retry on Ollama 500 errors. 500s occur when Ollama is still loading a model or handling a concurrent request (e.g. tag suggestions racing with round 1). Retries up to 2 times with a short delay — by which point the model is warm and other calls done. """ last_exc: BaseException | None = None for attempt in range(3): if attempt > 0: delay = 3.0 * attempt logger.warning( "Ollama stream 500 (attempt %d/3), retrying in %.0fs", attempt, delay ) await asyncio.sleep(delay) try: async for chunk in stream_chat_with_tools(messages, model, tools=tools, think=think): yield chunk return except httpx.HTTPStatusError as exc: last_exc = exc if exc.response.status_code != 500: break # non-500 is not retryable except BaseException as exc: last_exc = exc break if last_exc is not None: raise last_exc async def run_generation( buf: GenerationBuffer, history: list[dict], model: str, user_id: int, conv_id: int, conv_title: str, user_content: str, context_note_id: int | None = None, include_note_ids: list[int] | None = None, excluded_note_ids: list[int] | None = None, think: bool = False, rag_project_id: int | None = None, workspace_project_id: int | None = None, user_timezone: str | None = None, voice_mode: bool = False, ) -> None: """Stream LLM response into buffer with periodic DB flushes.""" MAX_TOOL_ROUNDS = 5 msg_id = buf.assistant_message_id buf.append_event("status", {"status": "Building context..."}) # Phase 1: Resolve the tools list for this user. tools = await get_tools_for_user(user_id) logger.info( "Starting generation for conv %d: model=%s, tools=%d", conv_id, model, len(tools), ) # Phase 2: Summarize long conversation history if needed. history_to_use = history history_summary: str | None = None if len(history) > 30: # 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, model) # Phase 3: Build context and wait for model in parallel. model_load_task = asyncio.create_task(wait_for_model_loaded(model, timeout=180.0)) # Fetch voice_speech_style from user settings when voice_mode is active. voice_speech_style = "conversational" if voice_mode: from fabledassistant.services.settings import get_setting voice_speech_style = await get_setting(user_id, "voice_speech_style", "conversational") context_task = asyncio.create_task(build_context( user_id, history_to_use, context_note_id, user_content, history_summary=history_summary, include_note_ids=include_note_ids, excluded_note_ids=excluded_note_ids, rag_project_id=rag_project_id, workspace_project_id=workspace_project_id, user_timezone=user_timezone, conv_id=conv_id, voice_mode=voice_mode, voice_speech_style=voice_speech_style, )) messages, context_meta = await context_task # Emit context event buf.append_event("context", {"context": context_meta}) # Wait for main model to be loaded before starting any generation. # If it's already loaded (common case), this returns immediately. if not model_load_task.done(): buf.append_event("status", {"status": "Loading model..."}) loaded = await model_load_task if not loaded: logger.warning("Model %s did not load within 180s — proceeding anyway", model) t_start = time.monotonic() timing: dict = { "think": think, "tools": [], "rounds": 0, "prompt_tokens": None, "output_tokens": None, "ttft_ms": None, "generation_ms": None, "total_ms": None, } last_flush = time.monotonic() all_tool_calls: list[dict] = [] new_rag_scope: object = False # sentinel; set to int|None when scope changes new_rag_scope_label: str | None = None try: cancelled = False research_completed = False for _round in range(MAX_TOOL_ROUNDS + 1): timing["rounds"] = _round + 1 round_tool_calls: list[dict] = [] logger.info("Generation round %d started for conv %d (model=%s)", _round, conv_id, model) if cancelled: break 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): if buf.cancel_event.is_set(): cancelled = True break if chunk.type == "thinking": buf.append_event("thinking_chunk", {"chunk": chunk.content}) elif 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 == "done": if chunk.prompt_tokens is not None: timing["prompt_tokens"] = (timing["prompt_tokens"] or 0) + chunk.prompt_tokens if chunk.output_tokens is not None: timing["output_tokens"] = (timing["output_tokens"] or 0) + chunk.output_tokens 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() if tool_name == "research_topic": topic = arguments.get("topic", "") try: note = await run_research_pipeline(topic, user_id, model, buf, project_id=workspace_project_id) result = { "success": True, "type": "research_note", "data": {"id": note.id, "title": note.title}, } done_text = ( f"\n\n---\n\nResearch complete! I've compiled a note: " f"**[{note.title}](/notes/{note.id})**." ) buf.append_event("chunk", {"chunk": done_text}) buf.content_so_far += done_text except Exception as e: logger.exception("Research pipeline failed for topic: %s", topic) result = {"success": False, "error": str(e)} err_text = f"\nResearch failed: {e}" buf.append_event("chunk", {"chunk": err_text}) buf.content_so_far += err_text research_completed = True else: result = await execute_tool( user_id, tool_name, arguments, conv_id=conv_id, workspace_project_id=workspace_project_id, ) # Capture RAG scope change for SSE done event if result.get("type") == "rag_scope_set" and result.get("success"): new_rag_scope = arguments.get("project_id") new_rag_scope_label = result.get("scope_label") 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")) 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: logger.info("Generation cancelled for conv %d", conv_id) break if research_completed: logger.info("Research complete for conv %d, ending generation", conv_id) break if not round_tool_calls: logger.info("Round %d: no tool calls, final content length=%d", _round, len(buf.content_so_far)) break logger.info("Round %d: %d tool call(s) executed, starting next round", _round, len(round_tool_calls)) 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 = "" # Strip model artifacts from final content buf.content_so_far = _TOOL_CALL_MARKER.sub("", buf.content_so_far) # Final save logger.info("Generation complete for conv %d: content_length=%d, tool_calls=%d", conv_id, len(buf.content_so_far), len(all_tool_calls)) await _update_message( msg_id, buf.content_so_far, "complete", tool_calls=all_tool_calls if all_tool_calls else None, ) timing["total_ms"] = int((time.monotonic() - t_start) * 1000) logger.info( "Generation timing for conv %d: total=%dms ttft=%s tools=%s generation=%s", conv_id, timing["total_ms"], timing["ttft_ms"], [(t["name"], t["ms"]) for t in timing["tools"]], timing["generation_ms"], ) try: await log_generation(user_id, conv_id, model, timing) except Exception: logger.warning("Failed to persist generation timing for conv %d", conv_id, exc_info=True) buf.state = GenerationState.COMPLETED buf.finished_at = time.monotonic() done_payload: dict = {"done": True, "message_id": msg_id, "timing": timing} if new_rag_scope is not False: done_payload["new_rag_scope"] = new_rag_scope done_payload["new_rag_scope_label"] = new_rag_scope_label buf.append_event("done", done_payload) # Fire push notification when complete (non-critical, fire-and-forget) try: from fabledassistant.services.push import send_push_notification, vapid_enabled if vapid_enabled(): text = buf.content_so_far.strip() if text: preview = text[:120].rstrip() if len(text) > 120: preview += "…" else: # Tool-only response — summarise what was done tool_names = [tc.get("function") for tc in all_tool_calls if tc.get("function")] if tool_names: preview = f"Completed: {', '.join(tool_names[:3])}" else: preview = "Action completed" asyncio.create_task(send_push_notification( user_id, title="Response ready", body=preview, url=f"/chat/{conv_id}", )) except Exception: logger.warning("Failed to schedule push notification", exc_info=True) # Title generation is non-critical — fire-and-forget so done fires immediately non_system = [m for m in messages if m["role"] != "system"] msg_count = len(non_system) should_gen_title = not conv_title or (msg_count > 0 and msg_count % 10 == 0) if should_gen_title: title_messages = messages + [ {"role": "assistant", "content": buf.content_so_far} ] async def _bg_title() -> None: try: title = await _generate_title(title_messages, model) if title: await update_conversation_title(user_id, conv_id, title) except Exception: logger.warning("Failed to generate title for conversation %d", conv_id, exc_info=True) if not conv_title: fallback = user_content[:80] if len(user_content) > 80: fallback += "..." await update_conversation_title(user_id, conv_id, fallback) asyncio.create_task(_bg_title()) except Exception as e: logger.exception("Error in generation task for conversation %d", conv_id) # Save partial content with error status try: await _update_message(msg_id, buf.content_so_far, "error") except Exception: logger.warning("Failed to save error state for message %d", msg_id, exc_info=True) buf.state = GenerationState.ERRORED buf.finished_at = time.monotonic() buf.append_event("error", {"error": str(e)}) async def run_assist_generation( buf: GenerationBuffer, messages: list[dict], model: str, ) -> None: """Stream LLM response for assist into buffer. No DB persistence. Retries up to 3 times on Ollama 500 errors (model still loading). On each retry the accumulated content is reset so the done event always reflects only the successful generation. """ input_chars = sum(len(m.get("content", "")) for m in messages) logger.info("Assist generation started: model=%s, input_chars=%d", model, input_chars) last_exc: BaseException | None = None for attempt in range(3): if attempt > 0: delay = 3.0 * attempt logger.warning( "Ollama assist stream 500 (attempt %d/3), retrying in %.0fs", attempt, delay ) await asyncio.sleep(delay) try: buf.content_so_far = "" async for chunk in stream_chat(messages, model, options={"num_predict": Config.OLLAMA_NUM_CTX}): buf.content_so_far += chunk buf.append_event("chunk", {"chunk": chunk}) output_chars = len(buf.content_so_far) logger.info( "Assist generation complete: output_chars=%d, events=%d", output_chars, len(buf.events), ) buf.state = GenerationState.COMPLETED buf.finished_at = time.monotonic() buf.append_event("done", {"done": True, "full_text": buf.content_so_far}) logger.info("Assist done event appended (event index %d)", len(buf.events) - 1) return except httpx.HTTPStatusError as exc: last_exc = exc if exc.response.status_code != 500: break except Exception as exc: last_exc = exc break logger.exception("Error in assist generation task") buf.state = GenerationState.ERRORED buf.finished_at = time.monotonic() buf.append_event("error", {"error": str(last_exc)})