FabledGrimoire/src/fabledassistant/services/generation_task.py

"""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 IntentResult, 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",
}

# Action phrases used in the acknowledgment prompt — "You are about to: {action}"
_TOOL_ACTIONS: dict[str, str] = {
    "create_task": "create a task",
    "create_note": "create a new note",
    "update_note": "update an existing note",
    "list_tasks": "look up tasks",
    "search_notes": "search through notes",
    "create_event": "schedule a calendar event",
    "list_events": "check the calendar",
    "search_events": "search calendar events",
    "update_event": "update a calendar event",
    "delete_event": "remove a calendar event",
    "list_calendars": "list available calendars",
    "create_todo": "create a calendar todo",
    "list_todos": "check calendar todos",
    "update_todo": "update a calendar todo",
    "complete_todo": "mark a todo complete",
    "delete_todo": "remove a calendar todo",
}


async def _generate_acknowledgment(user_content: str, tool_name: str, model: str) -> str:
    """Generate a brief conversational acknowledgment that runs in parallel with tool execution.

    Uses the intent model (small, fast, already in VRAM) so the sentence is ready
    within ~200-400ms. Returned string includes a trailing double-newline so the
    main LLM response starts on a new paragraph.
    """
    action = _TOOL_ACTIONS.get(tool_name, "work on that")
    messages = [
        {
            "role": "system",
            "content": (
                "You are a helpful assistant. Write ONE short, natural sentence acknowledging "
                "what you are about to do. Vary your phrasing — do not always start with "
                "'Let me'. Be warm and conversational. Do not answer the question yet. "
                "Output only the sentence, nothing else."
            ),
        },
        {
            "role": "user",
            "content": f"User said: {user_content}\nYou are about to: {action}",
        },
    ]
    try:
        ack = await generate_completion(messages, model, max_tokens=40)
        ack = ack.strip()
        if ack:
            return ack + "\n\n"
    except Exception:
        logger.warning("Failed to generate acknowledgment", exc_info=True)
    return ""


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

    # Phase 1: launch all independent work in parallel so nothing waits on anything
    # unnecessarily. build_context (note search + system prompt) and the intent LLM
    # call are the two slow legs — run them concurrently.
    buf.append_event("status", {"status": "Building context..."})

    context_task = asyncio.create_task(build_context(
        user_id, history, context_note_id, user_content, exclude_note_ids=exclude_note_ids
    ))
    tools_task = asyncio.create_task(get_tools_for_user(user_id))
    intent_model_task = asyncio.create_task(get_setting(user_id, "intent_model", ""))

    # Tools + intent-model setting are fast DB calls — get them first so intent
    # can start immediately while build_context is still running.
    tools, intent_model_setting = await asyncio.gather(tools_task, intent_model_task)
    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),
    )

    # Start intent classification in parallel with remaining build_context work.
    pre_intent: IntentResult = IntentResult()
    intent_timing_ms: int | None = None
    if tools:
        intent_history = [
            m for m in history
            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_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 context_task

    # Emit context event
    buf.append_event("context", {"context": context_meta})

    t_start = time.monotonic()
    timing: dict = {
        "intent_ms": intent_timing_ms,
        "tools": [],
        "ttft_ms": None,
        "generation_ms": None,
        "total_ms": None,
    }

    last_flush = time.monotonic()
    all_tool_calls: list[dict] = []

    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 (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:
                    buf.append_event("status", {"status": f"{_TOOL_LABELS.get(intent.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, intent.tool_name, intent.arguments),
                        _generate_acknowledgment(user_content, intent.tool_name, intent_model),
                    )
                    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"))

                    # 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)

                    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})

                    # 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,
                        "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)})