FabledGrimoire/src/fabledassistant/services/llm.py

import json
import logging
import re
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 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": 32768}
    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,
) -> 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").
    """
    options: dict = {"num_ctx": 32768}
    if tools:
        options["num_predict"] = 8192
    payload: dict = {
        "model": model,
        "messages": messages,
        "stream": True,
        "options": options,
        "think": False,
    }
    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."""
    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", "")


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)


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,
) -> 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, list_tasks, search_notes.",
    ]
    if has_caldav:
        tool_lines[-1] = (
            "Available actions: create_task, create_note, update_note, list_tasks, search_notes, "
            "create_event, list_events, search_events, update_event, delete_event, "
            "list_calendars, create_todo, list_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("For relative dates like 'Friday' or 'next week', resolve them to YYYY-MM-DD format.")
    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."
    )
    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 notes by keywords from user message — single OR query
    keywords = _extract_keywords(user_message)
    if keywords:
        search_exclude = set(exclude_set)
        if current_note_id:
            search_exclude.add(current_note_id)
        try:
            notes = await search_notes_for_context(
                user_id, keywords, exclude_ids=search_exclude or None, limit=3
            )
            snippets: list[str] = []
            for n in notes:
                body_preview = n.body[:2000] if n.body else ""
                snippets.append(f"- {n.title}: {body_preview}")
                context_meta["auto_notes"].append({"id": n.id, "title": n.title})
            if snippets:
                system_parts.append(
                    "\n\n--- Related Notes ---\n"
                    + "\n".join(snippets)
                    + "\n--- End Related Notes ---"
                )
        except Exception:
            logger.warning("Failed to search notes for context", exc_info=True)

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

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