"""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 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 from fabledassistant.services.chat import update_conversation_title from fabledassistant.services.intent import classify_intent from fabledassistant.services.logging import log_generation from fabledassistant.services.settings import get_setting from fabledassistant.services.tools import get_tools_for_user, execute_tool 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", "list_tasks": "Searching tasks", "search_notes": "Searching notes", "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", "create_todo": "Creating todo", "list_todos": "Listing todos", "update_todo": "Updating todo", "complete_todo": "Completing todo", "delete_todo": "Removing todo", } 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 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, exclude_note_ids: list[int] | None = None, ) -> None: """Stream LLM response into buffer with periodic DB flushes.""" MAX_TOOL_ROUNDS = 5 msg_id = buf.assistant_message_id # Build context inside the background task so the 202 response returns immediately buf.append_event("status", {"status": "Building context..."}) messages, context_meta = await build_context( user_id, history, context_note_id, user_content, exclude_note_ids=exclude_note_ids ) # Emit context event buf.append_event("context", {"context": context_meta}) t_start = time.monotonic() timing: dict = { "intent_ms": None, "tools": [], "ttft_ms": None, "generation_ms": None, "total_ms": None, } last_flush = time.monotonic() all_tool_calls: list[dict] = [] # Resolve tools and intent model in parallel tools, intent_model_setting = await asyncio.gather( get_tools_for_user(user_id), get_setting(user_id, "intent_model", ""), ) intent_model = intent_model_setting or Config.OLLAMA_INTENT_MODEL or model logger.info( "Starting generation for conv %d: model=%s, intent_model=%s, tools=%d", conv_id, model, intent_model, len(tools), ) try: cancelled = False for _round in range(MAX_TOOL_ROUNDS + 1): 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 if _round == 0 and tools: # Pass last 3 user/assistant pairs (6 messages) for anaphora resolution. # messages = [system, *history, current_user] — exclude system and current user. intent_history = [ m for m in messages[1:-1] if m.get("role") in ("user", "assistant") and m.get("content") ][-6:] buf.append_event("status", {"status": "Analyzing your request..."}) t_intent = time.monotonic() intent = await classify_intent(user_content, tools, intent_model, history=intent_history) timing["intent_ms"] = int((time.monotonic() - t_intent) * 1000) 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: buf.append_event("status", {"status": f"{_TOOL_LABELS.get(intent.tool_name, 'Working')}..."}) t_tool = time.monotonic() result = await execute_tool(user_id, intent.tool_name, intent.arguments) timing["tools"].append({"name": intent.tool_name, "ms": int((time.monotonic() - t_tool) * 1000)}) logger.info("Intent-routed tool %s result: success=%s", intent.tool_name, result.get("success")) tool_record = { "function": intent.tool_name, "arguments": intent.arguments, "result": result, "status": "success" if result.get("success") else "error", } all_tool_calls.append(tool_record) buf.append_event("tool_call", {"tool_call": tool_record}) # Inject into messages so LLM can write a natural response messages.append({ "role": "assistant", "content": "", "tool_calls": [ {"function": {"name": intent.tool_name, "arguments": intent.arguments}} ], }) messages.append({ "role": "tool", "content": json.dumps(result), }) continue # Next round: LLM streams response incorporating result buf.append_event("status", {"status": "Generating response..." if _round == 0 else "Composing response..."}) t_stream = time.monotonic() async for chunk in stream_chat_with_tools(messages, model, tools=tools): 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 # Filter out "[TOOL_CALLS]" marker from streaming output clean = _TOOL_CALL_MARKER.sub("", chunk.content) if clean: buf.append_event("chunk", {"chunk": clean}) # Periodic DB flush 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)) # Process each tool call 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")) 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) # Emit tool_call SSE event 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 no tool calls this round, the LLM gave its final text response 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)) # Strip model artifacts like "[TOOL_CALLS]" from content buf.content_so_far = _TOOL_CALL_MARKER.sub("", buf.content_so_far) # Append assistant tool_call message and tool results to conversation # for the next round 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"]), }) # Reset content for the next round (LLM will produce a new response) 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 intent=%s tools=%s generation=%s", conv_id, timing["total_ms"], timing["ttft_ms"], timing["intent_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() buf.append_event("done", {"done": True, "message_id": msg_id, "timing": timing}) # 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.""" try: async for chunk in stream_chat(messages, model, options={"num_predict": 4096}): buf.content_so_far += chunk buf.append_event("chunk", {"chunk": chunk}) buf.state = GenerationState.COMPLETED buf.finished_at = time.monotonic() buf.append_event("done", {"done": True, "full_text": buf.content_so_far}) except Exception as e: logger.exception("Error in assist generation task") buf.state = GenerationState.ERRORED buf.finished_at = time.monotonic() buf.append_event("error", {"error": str(e)})