FabledScribe/src/fabledassistant/services/llm.py

import asyncio
import ipaddress
import json
import logging
import re
import socket
import time
from collections.abc import AsyncGenerator
from dataclasses import dataclass, field
from typing import Literal
from urllib.parse import urlparse

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

RAG_AUTO_THRESHOLD = 0.60
RAG_AUTO_LIMIT = 3
RAG_AUTO_SNIPPET = 4000


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,
    think: bool = False,
) -> AsyncGenerator[str, None]:
    """Stream chat completion from Ollama, yielding content chunks.

    Set think=False (default) to disable chain-of-thought on qwen3+ models.
    Thinking tokens are silently discarded anyway, but disabling avoids the
    multi-minute delay before the first content token arrives.
    """
    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, "think": think}
    # read=None: no per-chunk timeout — Ollama may pause for any duration while
    # processing a large input context before the first token arrives.
    async with httpx.AsyncClient(timeout=httpx.Timeout(connect=30.0, read=None, write=None, pool=30.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", "thinking", "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
    # read=None: no per-chunk timeout for the same reason as stream_chat.
    async with httpx.AsyncClient(timeout=httpx.Timeout(connect=30.0, read=None, write=None, pool=30.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", {})

                # Thinking chunks (qwen3 chain-of-thought, only when think=True)
                thinking = msg.get("thinking", "")
                if thinking:
                    yield ChatChunk(type="thinking", content=thinking)

                # 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,
    num_ctx: int | None = None,
) -> str:
    """Non-streaming chat completion, returns full response text.

    Retries up to 2 times on Ollama 500 errors (cold model loading race).
    num_ctx overrides the model's context window for this call only.
    """
    last_exc: Exception | None = None
    options: dict = {"num_predict": max_tokens}
    if num_ctx is not None:
        options["num_ctx"] = num_ctx
    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": options,
                    },
                )
                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


def _is_private_url(url: str) -> bool:
    """Return True if the URL resolves to a private/loopback/link-local address (SSRF guard)."""
    try:
        host = urlparse(url).hostname
        if not host:
            return True
        if host.lower() in ("localhost", "::1"):
            return True
        addr_info = socket.getaddrinfo(host, None, proto=socket.IPPROTO_TCP)
        for entry in addr_info:
            ip = ipaddress.ip_address(entry[4][0])
            if ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_reserved or ip.is_multicast:
                return True
        return False
    except Exception:
        return True  # Block on resolution failure


async def fetch_url_content(url: str) -> str:
    """Fetch a URL and return text content (HTML tags stripped)."""
    if _is_private_url(url):
        logger.warning("Blocked fetch of private/internal URL: %s", url)
        return "[URL blocked: internal network access not permitted]"
    try:
        async with httpx.AsyncClient(
            timeout=15.0, follow_redirects=False, 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 = 30  # total messages before summarizing
_HISTORY_KEEP_RECENT = 8         # verbatim tail to preserve (4 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:]

    # Two-pass for very long histories: summarize first half, combine with second half
    if len(to_summarize) > 50:
        mid = len(to_summarize) // 2
        first_half = to_summarize[:mid]
        second_half = to_summarize[mid:]

        # Summarize first half
        first_lines = []
        for m in first_half:
            role = m.get("role", "")
            content = (m.get("content") or "").strip()
            if role in ("user", "assistant") and content:
                label = "User" if role == "user" else "Assistant"
                first_lines.append(f"{label}: {content[:400]}")

        if first_lines:
            try:
                first_summary_messages = [
                    {"role": "system", "content": "Summarize this conversation in 3-4 sentences covering topics, notes/tasks created, and key decisions."},
                    {"role": "user", "content": "\n".join(first_lines)},
                ]
                summary_a = await generate_completion(first_summary_messages, model, max_tokens=300)
                summary_a = summary_a.strip()
            except Exception:
                summary_a = ""
        else:
            summary_a = ""

        # Build lines for final pass from second half
        second_lines = []
        for m in second_half:
            role = m.get("role", "")
            content = (m.get("content") or "").strip()
            if role in ("user", "assistant") and content:
                label = "User" if role == "user" else "Assistant"
                second_lines.append(f"{label}: {content[:400]}")

        if summary_a:
            lines = [f"[Earlier summary: {summary_a}]"] + second_lines
        else:
            lines = second_lines
    else:
        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. Capture: "
                "(1) All notes, tasks, and projects created or modified — include their exact names. "
                "(2) Key decisions made and conclusions reached. "
                "(3) Open questions and next steps mentioned. "
                "(4) The overall topic arc so the conversation can continue naturally. "
                "Be specific and factual. Output 4-8 concise sentences. Nothing else."
            ),
        },
        {"role": "user", "content": "\n".join(lines)},
    ]

    try:
        summary = await generate_completion(prompt_messages, model, max_tokens=400)
        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,
    excluded_note_ids: list[int] | None = None,
    rag_project_id: 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."
        )
        tool_lines.append("For calendar events, use ISO 8601 datetime format (e.g. 2026-09-30T00:00:00).")
        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("When creating notes, use the `tags` parameter — do not embed #tag text in the note body.")
    tool_lines.append(
        "When search_images returns results, embed each image directly in your response by writing "
        "the 'embed' field verbatim (e.g. ![Cat](/api/images/3)), then the 'citation' field on the "
        "next line. Never describe images as text or list their URLs — always render them as markdown images."
    )
    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 search_notes for conceptual/semantic queries — e.g. 'what notes do I have about X' or "
        "'find notes related to Y' — it uses semantic understanding to find thematically related content "
        "even when exact words don't match. Pass project= to scope the search to a specific project. "
        "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": [],
        "auto_injected_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. High-confidence results (>=0.60) are auto-injected
    # into the system prompt; lower-confidence results populate the sidebar only.
    # Users can also explicitly include notes via the sidebar (include_note_ids).
    search_exclude = set(exclude_set)
    if current_note_id:
        search_exclude.add(current_note_id)

    # (score, note) pairs — score is float for semantic results, None for keyword fallback.
    found_scored: list[tuple[float | None, object]] = []

    # Try semantic search first; fall back to keyword search on failure / no results.
    try:
        from fabledassistant.services.embeddings import semantic_search_notes
        for score, note in await semantic_search_notes(
            user_id, user_message, exclude_ids=search_exclude or None, limit=8,
            project_id=rag_project_id,
        ):
            found_scored.append((score, note))
    except Exception:
        logger.warning("Semantic note search failed, falling back to keyword search", exc_info=True)

    if not found_scored:
        keywords = _extract_keywords(user_message)
        if keywords:
            try:
                for note in await search_notes_for_context(
                    user_id, keywords, exclude_ids=search_exclude or None, limit=8,
                    project_id=rag_project_id,
                ):
                    found_scored.append((None, note))
            except Exception:
                logger.warning("Failed to search notes for context", exc_info=True)

    # Separate high-confidence results for auto-injection vs sidebar display
    excluded_inject_set = set(excluded_note_ids or [])
    auto_inject: list[tuple[float, object]] = []
    sidebar_only: list[tuple[float | None, object]] = []

    for score, n in found_scored:
        if (
            score is not None
            and score >= RAG_AUTO_THRESHOLD
            and len(auto_inject) < RAG_AUTO_LIMIT
            and n.id not in excluded_inject_set
        ):
            auto_inject.append((score, n))
        else:
            sidebar_only.append((score, n))

    # Inject high-scoring notes into system prompt
    if auto_inject:
        snippets = []
        for score, n in auto_inject:
            body_snippet = (n.body or "")[:RAG_AUTO_SNIPPET]
            snippets.append(f"**{n.title}** (relevance: {round(score * 100)}%)\n{body_snippet}")
            context_meta["auto_injected_notes"].append({
                "id": n.id,
                "title": n.title,
                "score": round(score, 2),
            })
        system_parts.append(
            "\n\n--- Relevant Notes ---\n"
            + "\n\n".join(snippets)
            + "\n--- End Relevant Notes ---"
        )

    # Populate sidebar candidates (auto-injected notes also appear here for reference,
    # but sidebar_only are the ones not yet in the prompt)
    for score, n in auto_inject:
        context_meta["auto_notes"].append({
            "id": n.id,
            "title": n.title,
            "score": round(score, 2) if score is not None else None,
            "auto_injected": True,
        })
    for score, n in sidebar_only:
        context_meta["auto_notes"].append({
            "id": n.id,
            "title": n.title,
            "score": round(score, 2) if score is not None else None,
            "auto_injected": False,
        })
    context_meta["auto_note_ids"] = [n.id for _, n in found_scored]

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