FabledScribe/src/fabledassistant/services/llm.py

import asyncio
import json
import logging
import re
import time
from collections.abc import AsyncGenerator
from dataclasses import dataclass, field
from typing import Literal

import httpx

from fabledassistant.config import Config
from fabledassistant.services.caldav import is_caldav_configured
from fabledassistant.services.notes import get_note, search_notes_for_context
from fabledassistant.services.settings import get_setting

logger = logging.getLogger(__name__)

STOP_WORDS = frozenset({
    "a", "an", "the", "is", "it", "to", "in", "for", "of", "and", "or",
    "on", "at", "by", "with", "from", "as", "be", "was", "were", "been",
    "are", "am", "do", "does", "did", "have", "has", "had", "will", "would",
    "can", "could", "shall", "should", "may", "might", "must", "that",
    "this", "these", "those", "i", "me", "my", "you", "your", "he", "she",
    "we", "they", "them", "his", "her", "its", "our", "their", "what",
    "which", "who", "whom", "how", "when", "where", "why", "not", "no",
    "but", "if", "so", "than", "too", "very", "just", "about", "up",
})


async def get_installed_models() -> set[str]:
    """Return set of installed Ollama model names (with and without :latest)."""
    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            resp = await client.get(f"{Config.OLLAMA_URL}/api/tags")
            resp.raise_for_status()
            data = resp.json()
            names: set[str] = set()
            for m in data.get("models", []):
                name = m["name"]
                names.add(name)
                if name.endswith(":latest"):
                    names.add(name.removesuffix(":latest"))
            return names
    except Exception:
        logger.warning("Failed to fetch installed Ollama models")
        return set()


async def ensure_model(model: str) -> None:
    """Check if model exists in Ollama, pull if missing."""
    try:
        installed = await get_installed_models()
        if model in installed or f"{model}:latest" in installed:
            logger.info("Model '%s' already available", model)
            return
    except Exception:
        logger.warning("Failed to check Ollama models, attempting pull anyway")

    logger.info("Pulling model '%s' from Ollama...", model)
    try:
        async with httpx.AsyncClient(timeout=1800.0) as client:
            async with client.stream(
                "POST",
                f"{Config.OLLAMA_URL}/api/pull",
                json={"name": model},
            ) as resp:
                resp.raise_for_status()
                async for line in resp.aiter_lines():
                    if line.strip():
                        status = json.loads(line)
                        if "status" in status:
                            logger.info("Pull %s: %s", model, status["status"])
        logger.info("Model '%s' pulled successfully", model)
    except Exception:
        logger.warning("Failed to pull model '%s' — chat may not work", model, exc_info=True)


async def wait_for_model_loaded(model: str, timeout: float = 90.0) -> bool:
    """Poll /api/ps every 2s until the model appears in Ollama's loaded-model list.

    Returns True when the model is loaded, False if timeout elapses first.
    Used before generation to avoid streaming 500s during cold model loads.
    """
    base = model.removesuffix(":latest")
    deadline = time.monotonic() + timeout
    while True:
        try:
            async with httpx.AsyncClient(timeout=5.0) as client:
                resp = await client.get(f"{Config.OLLAMA_URL}/api/ps")
                resp.raise_for_status()
                loaded = {m["name"] for m in resp.json().get("models", [])}
                if model in loaded or f"{base}:latest" in loaded or base in loaded:
                    return True
        except Exception:
            pass  # Ollama may still be starting up
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            return False
        await asyncio.sleep(min(2.0, remaining))


async def stream_chat(
    messages: list[dict],
    model: str,
    options: dict | None = None,
) -> AsyncGenerator[str, None]:
    """Stream chat completion from Ollama, yielding content chunks."""
    merged_options = {"num_ctx": Config.OLLAMA_NUM_CTX}
    if options:
        merged_options.update(options)
    payload: dict = {"model": model, "messages": messages, "stream": True, "options": merged_options}
    async with httpx.AsyncClient(timeout=httpx.Timeout(1800.0, connect=30.0, read=300.0)) as client:
        async with client.stream(
            "POST",
            f"{Config.OLLAMA_URL}/api/chat",
            json=payload,
        ) as resp:
            resp.raise_for_status()
            async for line in resp.aiter_lines():
                if not line.strip():
                    continue
                data = json.loads(line)
                chunk = data.get("message", {}).get("content", "")
                if chunk:
                    yield chunk
                if data.get("done"):
                    break


@dataclass
class ChatChunk:
    """A chunk yielded by stream_chat_with_tools."""
    type: Literal["content", "tool_calls", "done"]
    content: str = ""
    tool_calls: list[dict] | None = None


async def stream_chat_with_tools(
    messages: list[dict],
    model: str,
    tools: list[dict] | None = None,
    think: bool = False,
) -> AsyncGenerator[ChatChunk, None]:
    """Stream chat completion from Ollama with tool support.

    Yields ChatChunk objects. If the model returns tool_calls, a
    ChatChunk(type="tool_calls") is yielded. Always ends with
    ChatChunk(type="done").

    Set think=True to enable the model's chain-of-thought reasoning (qwen3+).
    Thinking tokens are consumed by Ollama and not forwarded to the caller;
    only the final response content is yielded. Expect higher TTFT when enabled.
    """
    options: dict = {"num_ctx": Config.OLLAMA_NUM_CTX}
    if tools:
        options["num_predict"] = 8192
    payload: dict = {
        "model": model,
        "messages": messages,
        "stream": True,
        "options": options,
        "think": think,
    }
    if tools:
        payload["tools"] = tools
    async with httpx.AsyncClient(timeout=httpx.Timeout(1800.0, connect=30.0, read=300.0)) as client:
        async with client.stream(
            "POST",
            f"{Config.OLLAMA_URL}/api/chat",
            json=payload,
        ) as resp:
            resp.raise_for_status()
            accumulated_tool_calls: list[dict] = []
            async for line in resp.aiter_lines():
                if not line.strip():
                    continue
                data = json.loads(line)
                msg = data.get("message", {})

                # Content chunks
                chunk = msg.get("content", "")
                if chunk:
                    yield ChatChunk(type="content", content=chunk)

                # Collect tool calls from any message (some models
                # emit them before the done flag)
                tc = msg.get("tool_calls")
                if tc:
                    accumulated_tool_calls.extend(tc)

                if data.get("done"):
                    if accumulated_tool_calls:
                        logger.info(
                            "Ollama returned %d tool call(s): %s",
                            len(accumulated_tool_calls),
                            json.dumps(accumulated_tool_calls)[:500],
                        )
                        yield ChatChunk(type="tool_calls", tool_calls=accumulated_tool_calls)
                    else:
                        logger.debug("Ollama done with no tool calls")
                    yield ChatChunk(type="done")
                    break


async def generate_completion(messages: list[dict], model: str, max_tokens: int = 4096) -> str:
    """Non-streaming chat completion, returns full response text.

    Retries up to 2 times on Ollama 500 errors (cold model loading race).
    """
    last_exc: Exception | None = None
    for attempt in range(3):
        if attempt > 0:
            delay = 3.0 * attempt
            logger.warning(
                "generate_completion 500 (attempt %d/3), retrying in %.0fs", attempt, delay
            )
            await asyncio.sleep(delay)
        try:
            async with httpx.AsyncClient(timeout=httpx.Timeout(1800.0, connect=30.0, read=300.0)) as client:
                resp = await client.post(
                    f"{Config.OLLAMA_URL}/api/chat",
                    json={
                        "model": model,
                        "messages": messages,
                        "stream": False,
                        "think": False,
                        "options": {"num_predict": max_tokens},
                    },
                )
                resp.raise_for_status()
                data = resp.json()
                return data.get("message", {}).get("content", "")
        except httpx.HTTPStatusError as exc:
            last_exc = exc
            if exc.response.status_code != 500:
                break
        except Exception as exc:
            last_exc = exc
            break
    raise last_exc


async def fetch_url_content(url: str) -> str:
    """Fetch a URL and return text content (HTML tags stripped)."""
    try:
        async with httpx.AsyncClient(
            timeout=15.0, follow_redirects=True, headers={"User-Agent": "FabledAssistant/1.0"}
        ) as client:
            resp = await client.get(url)
            resp.raise_for_status()
            text = resp.text
            # Strip HTML tags
            text = re.sub(r"<script[^>]*>.*?</script>", "", text, flags=re.DOTALL)
            text = re.sub(r"<style[^>]*>.*?</style>", "", text, flags=re.DOTALL)
            text = re.sub(r"<[^>]+>", " ", text)
            # Collapse whitespace
            text = re.sub(r"\s+", " ", text).strip()
            # Truncate to reasonable size
            if len(text) > 4000:
                text = text[:4000] + "..."
            return text
    except Exception as e:
        logger.warning("Failed to fetch URL %s: %s", url, e)
        return f"[Failed to fetch URL: {url}]"


def _extract_keywords(text: str) -> list[str]:
    """Extract meaningful keywords from text for note search."""
    words = re.findall(r"\b[a-zA-Z]{3,}\b", text.lower())
    keywords = [w for w in words if w not in STOP_WORDS]
    # Deduplicate while preserving order
    seen: set[str] = set()
    unique = []
    for w in keywords:
        if w not in seen:
            seen.add(w)
            unique.append(w)
    return unique[:5]


def _find_urls(text: str) -> list[str]:
    """Find URLs in text."""
    return re.findall(r"https?://[^\s<>\"')\]]+", text)


# History summarization thresholds
_HISTORY_SUMMARY_THRESHOLD = 20  # total messages before summarizing
_HISTORY_KEEP_RECENT = 6         # verbatim tail to preserve (3 exchanges)


async def summarize_history_for_context(
    history: list[dict],
    model: str,
) -> tuple[list[dict], str | None]:
    """Summarize old conversation history when it exceeds the threshold.

    Returns (recent_history, summary_text | None).
    recent_history is the verbatim tail passed to the model.
    summary_text (when not None) should be injected into the system prompt
    so the model retains the gist of earlier exchanges without the full tokens.
    For short conversations, returns (history, None) immediately with no LLM call.
    """
    if len(history) <= _HISTORY_SUMMARY_THRESHOLD:
        return history, None

    to_summarize = history[:-_HISTORY_KEEP_RECENT]
    recent = history[-_HISTORY_KEEP_RECENT:]

    lines: list[str] = []
    for m in to_summarize:
        role = m.get("role", "")
        content = (m.get("content") or "").strip()
        if role in ("user", "assistant") and content:
            label = "User" if role == "user" else "Assistant"
            lines.append(f"{label}: {content[:400]}")

    if not lines:
        return history, None

    prompt_messages = [
        {
            "role": "system",
            "content": (
                "Summarize this conversation history in 3-5 concise sentences. "
                "Capture: topics discussed, any notes/tasks/events created or modified, "
                "decisions made, and context needed to continue the conversation naturally. "
                "Be specific and factual. Output only the summary, nothing else."
            ),
        },
        {"role": "user", "content": "\n".join(lines)},
    ]

    try:
        summary = await generate_completion(prompt_messages, model, max_tokens=200)
        summary = summary.strip()
        if summary:
            logger.info(
                "Summarized %d history messages (%d chars) for context",
                len(to_summarize), len(summary),
            )
            return recent, summary
    except Exception:
        logger.warning("Failed to summarize conversation history", exc_info=True)

    return history, None


async def build_context(
    user_id: int,
    history: list[dict],
    current_note_id: int | None,
    user_message: str,
    exclude_note_ids: list[int] | None = None,
    history_summary: str | None = None,
    include_note_ids: list[int] | None = None,
) -> tuple[list[dict], dict]:
    """Build messages array for Ollama with system prompt and context.

    Returns (messages, context_meta) where context_meta contains info about
    which notes were included as context.
    """
    exclude_set = set(exclude_note_ids or [])
    from datetime import date as date_type
    assistant_name = await get_setting(user_id, "assistant_name", "Fable")
    today = date_type.today().isoformat()
    has_caldav = await is_caldav_configured(user_id)

    # Build tool usage guidance based on available integrations
    tool_lines = [
        "You have access to tool functions. You MUST use them when the user asks you to create, add, find, schedule, or search for anything.",
        "CRITICAL: Call the tool functions directly. NEVER write out function calls as text or code. NEVER describe what you would do — just do it.",
        "Available actions: create_task, create_note, update_note, delete_note, delete_task, get_note, list_notes, list_tasks, search_notes.",
    ]
    if has_caldav:
        tool_lines[-1] = (
            "Available actions: create_task, create_note, update_note, delete_note, delete_task, get_note, list_notes, list_tasks, search_notes, "
            "create_event, list_events, search_events, update_event, delete_event, "
            "list_calendars, create_todo, list_todos, search_todos, update_todo, complete_todo, delete_todo."
        )
        tool_lines.append("For calendar events, use ISO 8601 datetime format (e.g. 2026-09-30T00:00:00).")
        tool_lines.append("CalDAV todos are separate from app tasks. Use create_task for app tasks, create_todo for CalDAV todos.")
        tool_lines.append("When the user says 'remind me' with a time before an event, use the reminder_minutes parameter.")
        tool_lines.append("Use search_todos to find a specific CalDAV todo by keyword when list_todos would return too many results.")
    tool_lines.append("For relative dates like 'Friday' or 'next week', resolve them to YYYY-MM-DD format.")
    tool_lines.append("When creating notes, use the `tags` parameter — do not embed #tag text in the note body.")
    tool_lines.append(
        "Use update_note to edit/expand an existing note OR to update a task's status/priority/due_date. "
        "Use create_note ONLY for genuinely new notes with a different title. "
        "Use list_tasks to find tasks by status, priority, or due date (e.g. overdue, high priority, in progress). "
        "If a note was created earlier in the conversation and the user provides more content for it, use update_note. "
        "Use get_note to read the full content of a specific note. "
        "Use list_notes to browse notes by recency or tag. "
        "Use delete_note / delete_task only when explicitly asked to delete — these require confirmation."
    )
    tool_guidance = "\n".join(tool_lines)

    system_parts = [
        f"You are a helpful assistant named {assistant_name}, integrated into a note-taking and task-tracking app called Fabled Assistant. "
        "Help users with their notes, tasks, and general questions. "
        "When note context is provided, use it to give relevant answers. "
        f"Today's date is {today}.\n\n"
        f"{tool_guidance}"
    ]

    context_meta: dict = {
        "context_note_id": None,
        "context_note_title": None,
        "auto_notes": [],
    }

    # Include current note context if provided — full body, no truncation
    if current_note_id:
        note = await get_note(user_id, current_note_id)
        if note:
            context_meta["context_note_id"] = note.id
            context_meta["context_note_title"] = note.title
            system_parts.append(
                f"\n\n--- Current Note ---\n"
                f"Title: {note.title}\n"
                f"Content:\n{note.body}\n"
                f"--- End Note ---"
            )

    # Search for related notes to populate sidebar candidates.
    # Results are NOT injected into the system prompt automatically — this keeps
    # the system prompt prefix stable so Ollama's KV cache can reuse prefill state.
    # Users can explicitly include notes via the sidebar (include_note_ids).
    search_exclude = set(exclude_set)
    if current_note_id:
        search_exclude.add(current_note_id)

    found_notes = []
    # Try semantic search first; fall back to keyword search on failure / no results.
    try:
        from fabledassistant.services.embeddings import semantic_search_notes
        found_notes = await semantic_search_notes(
            user_id, user_message, exclude_ids=search_exclude or None, limit=3
        )
    except Exception:
        logger.warning("Semantic note search failed, falling back to keyword search", exc_info=True)

    if not found_notes:
        keywords = _extract_keywords(user_message)
        if keywords:
            try:
                found_notes = await search_notes_for_context(
                    user_id, keywords, exclude_ids=search_exclude or None, limit=3
                )
            except Exception:
                logger.warning("Failed to search notes for context", exc_info=True)

    # Populate sidebar candidates (never auto-injected).
    for n in found_notes:
        context_meta["auto_notes"].append({"id": n.id, "title": n.title})
    context_meta["auto_note_ids"] = [n.id for n in found_notes]

    # Inject explicitly included notes (user opted in via sidebar click).
    if include_note_ids:
        from fabledassistant.services.notes import get_note as _get_note
        included_snippets: list[str] = []
        for nid in include_note_ids:
            try:
                n = await _get_note(user_id, nid)
                if n:
                    body_preview = n.body[:2000] if n.body else ""
                    included_snippets.append(f"- {n.title}: {body_preview}")
            except Exception:
                logger.warning("Failed to load included note %d for context", nid, exc_info=True)
        if included_snippets:
            system_parts.append(
                "\n\n--- Included Notes ---\n"
                + "\n".join(included_snippets)
                + "\n--- End Included Notes ---"
            )

    # Fetch URL content from user message
    urls = _find_urls(user_message)
    for url in urls[:2]:  # Limit to 2 URLs
        content = await fetch_url_content(url)
        if content and not content.startswith("[Failed"):
            system_parts.append(
                f"\n\n--- Content from {url} ---\n{content}\n--- End URL Content ---"
            )

    # Inject compressed summary of older exchanges when history has been trimmed
    if history_summary:
        system_parts.append(
            f"\n\n--- Earlier Conversation ---\n{history_summary}\n--- End Earlier Conversation ---"
        )

    messages = [{"role": "system", "content": "".join(system_parts)}]
    messages.extend(history)
    messages.append({"role": "user", "content": user_message})
    return messages, context_meta