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 <noreply@anthropic.com>
This commit is contained in:
@@ -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))
|
||||
Reference in New Issue
Block a user