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"]*>.*?", "", text, flags=re.DOTALL) text = re.sub(r"]*>.*?", "", 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