bvandeusenandClaude Sonnet 4.6
38697f2614
Auto-notes (keyword-matched): 2000 → 800 chars each (×3 max = 6000 → 2400 chars). Pinned note (explicit context): was unbounded → capped at 4000 chars with [truncated] marker. The main post-GPU bottleneck is TTFT caused by the prefill phase — the model processing the full input before generating any tokens. Shorter context = faster prefill. Users can ask follow-up questions for more detail. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
348 lines
14 KiB
Python
348 lines
14 KiB
Python
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 (cap at 4000 chars — explicitly selected
|
|
# so more generous, but very long notes would blow up the prefill phase)
|
|
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
|
|
body = note.body or ""
|
|
truncated = body[:4000]
|
|
suffix = "\n[...truncated]" if len(body) > 4000 else ""
|
|
system_parts.append(
|
|
f"\n\n--- Current Note ---\n"
|
|
f"Title: {note.title}\n"
|
|
f"Content:\n{truncated}{suffix}\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:
|
|
# 800 chars keeps context meaningful without bloating the prefill phase.
|
|
# The user can ask follow-up questions for more detail.
|
|
body_preview = n.body[:800] 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
|