From d6f4a6dbb6a93af17f83c74ff502900eaf1bd302 Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Wed, 18 Feb 2026 21:44:58 -0500 Subject: [PATCH] Add semantic note search (nomic-embed-text) and per-conversation note cache - New NoteEmbedding model + migration 0014 stores float embeddings (JSONB) - services/embeddings.py: get_embedding, upsert_note_embedding, semantic_search_notes (cosine similarity), backfill_note_embeddings - build_context() now tries semantic search first, falls back to keyword search; accepts cached_note_ids to reuse last-turn notes and stabilise the system prompt prefix for Ollama's KV cache - generation_buffer.py: per-conversation note ID cache (get/set/clear) - generation_task.py: passes cached IDs into build_context, updates cache after each turn, and invalidates it after create_note/update_note/create_task - app.py: pulls nomic-embed-text at startup and launches a background backfill to embed all existing notes (30 s delay so Ollama has time to load the model) - routes/notes.py + services/tools.py: fire-and-forget embedding update on every note create or update via the API or LLM tool calls Co-Authored-By: Claude Sonnet 4.6 --- alembic/versions/0014_add_note_embeddings.py | 24 +++ src/fabledassistant/app.py | 14 ++ src/fabledassistant/config.py | 3 + src/fabledassistant/models/__init__.py | 1 + src/fabledassistant/models/embedding.py | 25 +++ src/fabledassistant/routes/notes.py | 8 + src/fabledassistant/services/embeddings.py | 158 ++++++++++++++++++ .../services/generation_buffer.py | 21 +++ .../services/generation_task.py | 27 ++- src/fabledassistant/services/llm.py | 77 ++++++--- src/fabledassistant/services/tools.py | 13 ++ 11 files changed, 349 insertions(+), 22 deletions(-) create mode 100644 alembic/versions/0014_add_note_embeddings.py create mode 100644 src/fabledassistant/models/embedding.py create mode 100644 src/fabledassistant/services/embeddings.py diff --git a/alembic/versions/0014_add_note_embeddings.py b/alembic/versions/0014_add_note_embeddings.py new file mode 100644 index 0000000..5659558 --- /dev/null +++ b/alembic/versions/0014_add_note_embeddings.py @@ -0,0 +1,24 @@ +"""Add note_embeddings table for semantic note search.""" + +from alembic import op + +revision = "0014" +down_revision = "0013" + + +def upgrade() -> None: + op.execute(""" + CREATE TABLE IF NOT EXISTS note_embeddings ( + note_id INTEGER PRIMARY KEY REFERENCES notes(id) ON DELETE CASCADE, + user_id INTEGER NOT NULL, + embedding JSONB NOT NULL, + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + ) + """) + op.execute( + "CREATE INDEX IF NOT EXISTS ix_note_embeddings_user_id ON note_embeddings(user_id)" + ) + + +def downgrade() -> None: + op.execute("DROP TABLE IF EXISTS note_embeddings") diff --git a/src/fabledassistant/app.py b/src/fabledassistant/app.py index 9c9be67..1af6c92 100644 --- a/src/fabledassistant/app.py +++ b/src/fabledassistant/app.py @@ -84,6 +84,7 @@ def create_app() -> Quart: async def startup(): import asyncio + from fabledassistant.services.embeddings import backfill_note_embeddings from fabledassistant.services.generation_buffer import start_cleanup_loop from fabledassistant.services.llm import ensure_model from fabledassistant.services.logging import start_log_retention_loop @@ -108,9 +109,22 @@ def create_app() -> Quart: models_to_pull = {Config.OLLAMA_MODEL} if Config.OLLAMA_INTENT_MODEL and Config.OLLAMA_INTENT_MODEL != Config.OLLAMA_MODEL: models_to_pull.add(Config.OLLAMA_INTENT_MODEL) + # Also pull the embedding model (nomic-embed-text by default). + models_to_pull.add(Config.EMBEDDING_MODEL) for _model in models_to_pull: asyncio.create_task(_pull_model(_model)) + # After models are pulled, backfill embeddings for existing notes. + # Runs in the background so it never blocks the server from accepting requests. + async def _delayed_backfill() -> None: + await asyncio.sleep(30) # Give Ollama time to load the embedding model + try: + await backfill_note_embeddings() + except Exception: + logger.warning("Embedding backfill failed", exc_info=True) + + asyncio.create_task(_delayed_backfill()) + @app.route("/") async def serve_index(): resp = await make_response( diff --git a/src/fabledassistant/config.py b/src/fabledassistant/config.py index 10bcb80..4b562ad 100644 --- a/src/fabledassistant/config.py +++ b/src/fabledassistant/config.py @@ -32,6 +32,9 @@ class Config: LOG_LEVEL: str = os.environ.get("LOG_LEVEL", "INFO") LOG_RETENTION_DAYS: int = int(os.environ.get("LOG_RETENTION_DAYS", "90")) + # Embedding model for semantic note search (served by Ollama) + EMBEDDING_MODEL: str = os.environ.get("EMBEDDING_MODEL", "nomic-embed-text") + # SMTP defaults (overridden by DB settings when configured via admin UI) SMTP_HOST: str = os.environ.get("SMTP_HOST", "") SMTP_PORT: int = int(os.environ.get("SMTP_PORT", "587")) diff --git a/src/fabledassistant/models/__init__.py b/src/fabledassistant/models/__init__.py index 9669f57..4c77e2a 100644 --- a/src/fabledassistant/models/__init__.py +++ b/src/fabledassistant/models/__init__.py @@ -18,3 +18,4 @@ from fabledassistant.models.user import User # noqa: E402, F401 from fabledassistant.models.app_log import AppLog # noqa: E402, F401 from fabledassistant.models.password_reset import PasswordResetToken # noqa: E402, F401 from fabledassistant.models.invitation import InvitationToken # noqa: E402, F401 +from fabledassistant.models.embedding import NoteEmbedding # noqa: E402, F401 diff --git a/src/fabledassistant/models/embedding.py b/src/fabledassistant/models/embedding.py new file mode 100644 index 0000000..7819566 --- /dev/null +++ b/src/fabledassistant/models/embedding.py @@ -0,0 +1,25 @@ +from datetime import datetime, timezone + +from sqlalchemy import DateTime, ForeignKey, Integer +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.orm import Mapped, mapped_column + +from fabledassistant.models import Base + + +class NoteEmbedding(Base): + """Stores the embedding vector for a note, used for semantic search.""" + + __tablename__ = "note_embeddings" + + note_id: Mapped[int] = mapped_column( + Integer, + ForeignKey("notes.id", ondelete="CASCADE"), + primary_key=True, + ) + user_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True) + embedding: Mapped[list] = mapped_column(JSONB, nullable=False) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(timezone.utc), + ) diff --git a/src/fabledassistant/routes/notes.py b/src/fabledassistant/routes/notes.py index d04f3be..e6dacf3 100644 --- a/src/fabledassistant/routes/notes.py +++ b/src/fabledassistant/routes/notes.py @@ -2,6 +2,8 @@ import asyncio import logging from datetime import date +from fabledassistant.services.embeddings import upsert_note_embedding + from quart import Blueprint, Response, jsonify, request from fabledassistant.auth import login_required, get_current_user_id @@ -84,6 +86,9 @@ async def create_note_route(): priority=priority, due_date=due_date, ) + text = f"{note.title}\n{note.body}".strip() if note.body else (note.title or "") + if text: + asyncio.create_task(upsert_note_embedding(note.id, uid, text)) return jsonify(note.to_dict()), 201 @@ -182,6 +187,9 @@ async def update_note_route(note_id: int): note = await update_note(uid, note_id, **fields) if note is None: return jsonify({"error": "Note not found"}), 404 + text = f"{note.title}\n{note.body}".strip() if note.body else (note.title or "") + if text: + asyncio.create_task(upsert_note_embedding(note.id, uid, text)) return jsonify(note.to_dict()) diff --git a/src/fabledassistant/services/embeddings.py b/src/fabledassistant/services/embeddings.py new file mode 100644 index 0000000..bd7a6a2 --- /dev/null +++ b/src/fabledassistant/services/embeddings.py @@ -0,0 +1,158 @@ +"""Semantic note search via Ollama embedding model (nomic-embed-text). + +Embeddings are stored in the note_embeddings table (one row per note). +All search operations degrade gracefully — if the embedding model is +unavailable the callers fall back to keyword search. +""" + +import asyncio +import logging +import math + +import httpx +from sqlalchemy import delete, select + +from fabledassistant.config import Config +from fabledassistant.models import async_session +from fabledassistant.models.embedding import NoteEmbedding +from fabledassistant.models.note import Note + +logger = logging.getLogger(__name__) + +# Minimum cosine similarity to include a note in context results. +# nomic-embed-text produces unit-normalized vectors, so range is [-1, 1]. +_SIMILARITY_THRESHOLD = 0.30 + + +async def get_embedding(text: str, model: str | None = None) -> list[float]: + """Get an embedding vector from Ollama for the given text. + + Raises httpx.HTTPError on failure — callers should handle this. + """ + m = model or Config.EMBEDDING_MODEL + async with httpx.AsyncClient(timeout=30.0) as client: + resp = await client.post( + f"{Config.OLLAMA_URL}/api/embed", + json={"model": m, "input": text}, + ) + resp.raise_for_status() + data = resp.json() + # Ollama /api/embed → {"embeddings": [[float, ...]]} + return data["embeddings"][0] + + +def _cosine_similarity(a: list[float], b: list[float]) -> float: + """Cosine similarity between two vectors. Returns 0 for zero-length vectors.""" + dot = sum(x * y for x, y in zip(a, b)) + mag_a = math.sqrt(sum(x * x for x in a)) + mag_b = math.sqrt(sum(x * x for x in b)) + if mag_a == 0.0 or mag_b == 0.0: + return 0.0 + return dot / (mag_a * mag_b) + + +async def upsert_note_embedding(note_id: int, user_id: int, text: str) -> None: + """Generate and persist an embedding for a note. Safe to fire-and-forget.""" + try: + embedding = await get_embedding(text) + except Exception: + logger.debug("Skipping embedding for note %d — model unavailable", note_id) + return + + try: + async with async_session() as session: + await session.execute( + delete(NoteEmbedding).where(NoteEmbedding.note_id == note_id) + ) + session.add(NoteEmbedding(note_id=note_id, user_id=user_id, embedding=embedding)) + await session.commit() + logger.debug("Upserted embedding for note %d", note_id) + except Exception: + logger.warning("Failed to persist embedding for note %d", note_id, exc_info=True) + + +async def semantic_search_notes( + user_id: int, + query: str, + exclude_ids: set[int] | None = None, + limit: int = 3, +) -> list[Note]: + """Return up to *limit* notes most relevant to *query* using cosine similarity. + + Returns an empty list if the embedding model is unavailable or on any error. + """ + try: + query_vec = await get_embedding(query) + except Exception: + logger.debug("Semantic search skipped — embedding model unavailable") + return [] + + try: + async with async_session() as session: + stmt = ( + select(NoteEmbedding, Note) + .join(Note, NoteEmbedding.note_id == Note.id) + .where(NoteEmbedding.user_id == user_id) + ) + if exclude_ids: + stmt = stmt.where(NoteEmbedding.note_id.notin_(exclude_ids)) + rows = list((await session.execute(stmt)).all()) + except Exception: + logger.warning("Failed to query note embeddings", exc_info=True) + return [] + + if not rows: + return [] + + scored: list[tuple[float, Note]] = [] + for ne, note in rows: + try: + sim = _cosine_similarity(query_vec, ne.embedding) + except Exception: + continue + if sim >= _SIMILARITY_THRESHOLD: + scored.append((sim, note)) + + scored.sort(key=lambda x: x[0], reverse=True) + return [note for _, note in scored[:limit]] + + +async def backfill_note_embeddings() -> None: + """Generate embeddings for all notes that don't have one yet. + + Runs as a background task at startup. Adds a small sleep between notes + to avoid overwhelming Ollama. + """ + try: + async with async_session() as session: + existing = { + row[0] + for row in ( + await session.execute(select(NoteEmbedding.note_id)) + ).fetchall() + } + result = await session.execute( + select(Note.id, Note.user_id, Note.title, Note.body) + ) + notes_to_embed = [ + row for row in result.fetchall() if row[0] not in existing + ] + except Exception: + logger.warning("Embedding backfill: failed to query notes", exc_info=True) + return + + if not notes_to_embed: + logger.info("Embedding backfill: all notes already have embeddings") + return + + logger.info("Embedding backfill: generating embeddings for %d notes", len(notes_to_embed)) + success = 0 + for note_id, user_id, title, body in notes_to_embed: + text = f"{title}\n{body}".strip() if body else (title or "") + if not text: + continue + await upsert_note_embedding(note_id, user_id, text) + success += 1 + await asyncio.sleep(0.05) # gentle pacing + + logger.info("Embedding backfill complete: %d/%d notes embedded", success, len(notes_to_embed)) diff --git a/src/fabledassistant/services/generation_buffer.py b/src/fabledassistant/services/generation_buffer.py index ac3cccc..2c59dc3 100644 --- a/src/fabledassistant/services/generation_buffer.py +++ b/src/fabledassistant/services/generation_buffer.py @@ -74,6 +74,27 @@ class GenerationBuffer: # Module-level singleton registry _buffers: dict[int | str, GenerationBuffer] = {} + +# Per-conversation note context cache — maps conv_id → sorted list of note IDs. +# Stores the note IDs that were last included in the system prompt so that +# subsequent turns in the same conversation can reuse them, stabilizing the +# system prompt prefix and improving Ollama's KV cache hit rate. +_conv_note_cache: dict[int, list[int]] = {} + + +def get_conv_note_cache(conv_id: int) -> list[int]: + """Return cached note IDs for a conversation (empty list if none).""" + return list(_conv_note_cache.get(conv_id, [])) + + +def set_conv_note_cache(conv_id: int, note_ids: list[int]) -> None: + """Store note IDs to reuse on the next turn of this conversation.""" + _conv_note_cache[conv_id] = list(note_ids) + + +def clear_conv_note_cache(conv_id: int) -> None: + """Invalidate the note cache for a conversation (e.g. after a note write).""" + _conv_note_cache.pop(conv_id, None) _cleanup_task: asyncio.Task | None = None _GRACE_PERIOD = 60.0 # seconds to keep completed buffers diff --git a/src/fabledassistant/services/generation_task.py b/src/fabledassistant/services/generation_task.py index 7ff4ff8..a34e5c4 100644 --- a/src/fabledassistant/services/generation_task.py +++ b/src/fabledassistant/services/generation_task.py @@ -15,7 +15,13 @@ from sqlalchemy import update from fabledassistant.config import Config from fabledassistant.models import async_session from fabledassistant.models.conversation import Message -from fabledassistant.services.generation_buffer import GenerationBuffer, GenerationState +from fabledassistant.services.generation_buffer import ( + GenerationBuffer, + GenerationState, + clear_conv_note_cache, + get_conv_note_cache, + set_conv_note_cache, +) from fabledassistant.services.llm import ChatChunk, build_context, generate_completion, stream_chat, stream_chat_with_tools, summarize_history_for_context from fabledassistant.services.chat import update_conversation_title from fabledassistant.services.intent import IntentResult, classify_intent @@ -196,6 +202,10 @@ async def run_generation( history_to_use, history_summary = await summarize_history_for_context(history, intent_model) # Phase 3: Build context and classify intent in parallel — the two slow legs. + # Pass cached note IDs so build_context can reuse them, keeping the system + # prompt prefix stable and helping Ollama's KV cache stay warm. + cached_note_ids = get_conv_note_cache(conv_id) or None + pre_intent: IntentResult = IntentResult() intent_timing_ms: int | None = None if tools: @@ -209,6 +219,7 @@ async def run_generation( user_id, history_to_use, context_note_id, user_content, exclude_note_ids=exclude_note_ids, history_summary=history_summary, + cached_note_ids=cached_note_ids, )) intent_task = asyncio.create_task( classify_intent(user_content, tools, intent_model, history=intent_history) @@ -220,8 +231,14 @@ async def run_generation( user_id, history_to_use, context_note_id, user_content, exclude_note_ids=exclude_note_ids, history_summary=history_summary, + cached_note_ids=cached_note_ids, ) + # Update the note cache with whatever notes ended up in context. + new_note_ids = context_meta.get("auto_note_ids") or [] + if new_note_ids: + set_conv_note_cache(conv_id, new_note_ids) + # Emit context event buf.append_event("context", {"context": context_meta}) @@ -332,6 +349,11 @@ async def run_generation( if timing["ttft_ms"] is None: timing["ttft_ms"] = int((time.monotonic() - t_start) * 1000) + # Invalidate the note context cache after any successful note write + # so the next turn can pick up newly created/modified notes. + if result.get("success") and tool_name in {"create_task", "create_note", "update_note"}: + clear_conv_note_cache(conv_id) + tool_record = { "function": tool_name, "arguments": intent.arguments, @@ -400,6 +422,9 @@ async def run_generation( timing["tools"].append({"name": tool_name, "ms": int((time.monotonic() - t_tool) * 1000)}) logger.info("Tool %s result: success=%s", tool_name, result.get("success")) + if result.get("success") and tool_name in {"create_task", "create_note", "update_note"}: + clear_conv_note_cache(conv_id) + tool_record = { "function": tool_name, "arguments": arguments, diff --git a/src/fabledassistant/services/llm.py b/src/fabledassistant/services/llm.py index 52035ab..ddd4bc6 100644 --- a/src/fabledassistant/services/llm.py +++ b/src/fabledassistant/services/llm.py @@ -307,6 +307,7 @@ async def build_context( user_message: str, exclude_note_ids: list[int] | None = None, history_summary: str | None = None, + cached_note_ids: list[int] | None = None, ) -> tuple[list[dict], dict]: """Build messages array for Ollama with system prompt and context. @@ -370,29 +371,63 @@ async def build_context( 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) + # Find related notes to inject into context. + # Priority: (1) use cached note IDs from a previous turn in this conversation + # (2) try semantic search via nomic-embed-text + # (3) fall back to keyword search + # The cache stabilises the system prompt prefix so Ollama's KV cache stays warm. + search_exclude = set(exclude_set) + if current_note_id: + search_exclude.add(current_note_id) + + found_notes = [] + if cached_note_ids: + # Load the same notes as last turn — keeps system prompt prefix identical. 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 ---" - ) + from fabledassistant.services.notes import get_note as _get_note + for nid in cached_note_ids: + if nid not in search_exclude: + n = await _get_note(user_id, nid) + if n: + found_notes.append(n) except Exception: - logger.warning("Failed to search notes for context", exc_info=True) + logger.warning("Failed to load cached notes for context", exc_info=True) + found_notes = [] + + if not 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) + + if found_notes: + snippets: list[str] = [] + for n in found_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}) + system_parts.append( + "\n\n--- Related Notes ---\n" + + "\n".join(snippets) + + "\n--- End Related Notes ---" + ) + + # Expose note IDs so the caller can update the per-conversation cache. + context_meta["auto_note_ids"] = [n.id for n in found_notes] # Fetch URL content from user message urls = _find_urls(user_message) diff --git a/src/fabledassistant/services/tools.py b/src/fabledassistant/services/tools.py index a5d1ba1..1c29e84 100644 --- a/src/fabledassistant/services/tools.py +++ b/src/fabledassistant/services/tools.py @@ -1,5 +1,6 @@ """Tool definitions and executor for LLM tool calling.""" +import asyncio import logging from datetime import date, datetime @@ -22,6 +23,15 @@ from fabledassistant.services.tag_suggestions import suggest_tags logger = logging.getLogger(__name__) + +def _schedule_embedding(note_id: int, user_id: int, title: str, body: str) -> None: + """Fire-and-forget: update the embedding for a note after it's created/modified.""" + from fabledassistant.services.embeddings import upsert_note_embedding + text = f"{title}\n{body}".strip() if body else (title or "") + if text: + asyncio.create_task(upsert_note_embedding(note_id, user_id, text)) + + # Core tools — always available _CORE_TOOLS = [ { @@ -546,6 +556,7 @@ async def execute_tool(user_id: int, tool_name: str, arguments: dict) -> dict: due_date=_parse_due_date(arguments.get("due_date")), ) suggested = await suggest_tags(user_id, task_title, task_body) + _schedule_embedding(note.id, user_id, task_title, task_body) return { "success": True, "type": "task", @@ -568,6 +579,7 @@ async def execute_tool(user_id: int, tool_name: str, arguments: dict) -> dict: body=note_body, ) suggested = await suggest_tags(user_id, note_title, note_body) + _schedule_embedding(note.id, user_id, note_title, note_body) return { "success": True, "type": "note", @@ -616,6 +628,7 @@ async def execute_tool(user_id: int, tool_name: str, arguments: dict) -> dict: return {"success": False, "error": "Failed to update note."} suggested = await suggest_tags(user_id, updated.title, updated.body or "") + _schedule_embedding(updated.id, user_id, updated.title, updated.body or "") return { "success": True, "type": "note_updated",