0d9ebc9cc7
Without this, every `second-brain process` run prints couldn't stop thread 'pool-1-worker-N' within 5.0 seconds during interpreter shutdown — psycopg_pool's background workers don't get a clean stop signal before Python's threading shutdown deadline. Registering close_pool via atexit the first time we open the pool fixes it without changing any API. `close_pool` is already idempotent, so the explicit teardown path in the smoke test (which calls it directly) and the atexit path coexist safely.
258 lines
8.2 KiB
Python
258 lines
8.2 KiB
Python
"""Embed second-brain extractions into the shared public.embeddings table.
|
|
|
|
This is a near-verbatim copy of the vault-mcp recipe (vault_mcp.core.embeddings),
|
|
adapted for our `source_schema='second_brain' / source_table='extractions'`
|
|
key tuple and our embed payload (the LLM-generated summary).
|
|
|
|
Behaviour:
|
|
- Reuses the shared `embedding_chunking` library for 512-token chunking with
|
|
64-token overlap.
|
|
- Calls Ollama's /api/embeddings with `num_ctx=8192` to embed nomic-embed-text
|
|
vectors (768-d).
|
|
- Writes vector literals as `%s::vector` to avoid taking a hard dependency on
|
|
pgvector's Python adapter.
|
|
- Delete-before-insert on the (schema, table, source_id, embedding_model)
|
|
tuple so a re-embed of an extraction that previously had N chunks but now
|
|
has M<N doesn't leave orphans.
|
|
|
|
Graceful degradation: any failure (no DB URL, no Ollama, DB error, empty
|
|
summary) logs at WARNING and returns None. The relational `extractions` row
|
|
and the file-based vault remain the source of truth — embeddings are an
|
|
index, not authoritative state.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import atexit
|
|
import logging
|
|
import os
|
|
from typing import Any
|
|
|
|
import httpx
|
|
from embedding_chunking import Chunk, chunk_text
|
|
from psycopg_pool import ConnectionPool
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
SOURCE_SCHEMA = "second_brain"
|
|
SOURCE_TABLE = "extractions"
|
|
|
|
_pool: ConnectionPool | None = None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pool / URL plumbing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _database_url() -> str | None:
|
|
"""Resolve a raw psycopg DSN (no SQLAlchemy `+psycopg` prefix)."""
|
|
for var in ("SECOND_BRAIN_DATABASE_URL", "HERBYLAB_DATABASE_URL"):
|
|
v = os.environ.get(var)
|
|
if v:
|
|
return v.replace("postgresql+psycopg://", "postgresql://", 1)
|
|
|
|
# Fallback: pull from settings.toml so a host-side `process` works even
|
|
# when the operator forgot to export the env var.
|
|
try:
|
|
from second_brain.config import load_config
|
|
|
|
url = load_config().database.get("url")
|
|
if url:
|
|
return url.replace("postgresql+psycopg://", "postgresql://", 1)
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
def _get_pool() -> ConnectionPool | None:
|
|
"""Lazy-init a small connection pool. Returns None if no DB URL is configured."""
|
|
global _pool
|
|
if _pool is not None:
|
|
return _pool
|
|
dsn = _database_url()
|
|
if not dsn:
|
|
return None
|
|
_pool = ConnectionPool(dsn, min_size=1, max_size=4, open=True, timeout=10)
|
|
# Make sure the pool's worker threads get a clean shutdown when the
|
|
# interpreter exits; without this psycopg_pool prints
|
|
# "couldn't stop thread pool-1-worker-N within 5.0 seconds" on every
|
|
# `second-brain process` run.
|
|
atexit.register(close_pool)
|
|
return _pool
|
|
|
|
|
|
def close_pool() -> None:
|
|
"""Close the embedding-side pool. Safe to call multiple times."""
|
|
global _pool
|
|
if _pool is not None:
|
|
_pool.close()
|
|
_pool = None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Embedding model client
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _ollama_url() -> str:
|
|
"""Env var takes precedence over settings.toml so containerized runs can
|
|
point at a different Ollama instance without rebuilding the image."""
|
|
env = os.environ.get("OLLAMA_URL")
|
|
if env:
|
|
return env
|
|
try:
|
|
from second_brain.config import load_config
|
|
|
|
return load_config().embeddings.get("ollama_url", "http://ollama:11434")
|
|
except Exception:
|
|
return "http://ollama:11434"
|
|
|
|
|
|
def _embedding_model() -> str:
|
|
env = os.environ.get("EMBEDDING_MODEL")
|
|
if env:
|
|
return env
|
|
try:
|
|
from second_brain.config import load_config
|
|
|
|
return load_config().embeddings.get("model", "nomic-embed-text")
|
|
except Exception:
|
|
return "nomic-embed-text"
|
|
|
|
|
|
def embed_text(text: str) -> list[float]:
|
|
"""Call Ollama's embeddings endpoint. Returns the model-native vector."""
|
|
resp = httpx.post(
|
|
f"{_ollama_url()}/api/embeddings",
|
|
json={
|
|
"model": _embedding_model(),
|
|
"prompt": text,
|
|
# nomic-embed-text supports 8192-token context. Ollama defaults
|
|
# the runner to 2048, which 500s on long summaries.
|
|
"options": {"num_ctx": 8192},
|
|
},
|
|
timeout=120.0,
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
vec = data.get("embedding")
|
|
if not vec or not isinstance(vec, list):
|
|
raise RuntimeError(
|
|
f"Ollama returned no embedding for model {_embedding_model()}"
|
|
)
|
|
return vec
|
|
|
|
|
|
def _vector_literal(vec: list[float]) -> str:
|
|
"""pgvector accepts text-literal vectors of form '[v1,v2,...]'."""
|
|
return "[" + ",".join(repr(float(v)) for v in vec) + "]"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Public API
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def embed_extraction(extraction_id: int, summary: str | None) -> int | None:
|
|
"""Embed an extraction summary into public.embeddings.
|
|
|
|
Args:
|
|
extraction_id: PK of the row in second_brain.extractions.
|
|
summary: the LLM-generated summary string; we embed this and
|
|
nothing else for the current round (key_points / claims
|
|
are not embedded yet).
|
|
|
|
Returns the number of chunk rows written, or None on skip/failure.
|
|
Best-effort: failures are logged, not raised.
|
|
"""
|
|
if not summary or not summary.strip():
|
|
logger.debug("embed_extraction: skip — extraction %s has no summary", extraction_id)
|
|
return None
|
|
|
|
pool = _get_pool()
|
|
if pool is None:
|
|
logger.info(
|
|
"embed_extraction: skip — no SECOND_BRAIN_DATABASE_URL configured"
|
|
)
|
|
return None
|
|
|
|
model = _embedding_model()
|
|
|
|
try:
|
|
chunks: list[Chunk] = chunk_text(summary)
|
|
chunk_embeddings: list[tuple[Chunk, list[float]]] = [
|
|
(c, embed_text(c.text)) for c in chunks
|
|
]
|
|
except Exception as exc:
|
|
logger.warning(
|
|
"embed_extraction: chunk/embed failed for extraction %s: %s",
|
|
extraction_id,
|
|
exc,
|
|
)
|
|
return None
|
|
|
|
try:
|
|
with pool.connection() as conn:
|
|
n = _upsert_embeddings(
|
|
conn, extraction_id, chunk_embeddings, model
|
|
)
|
|
conn.commit()
|
|
logger.info(
|
|
"embed_extraction: extraction=%s chunks=%d model=%s",
|
|
extraction_id,
|
|
n,
|
|
model,
|
|
)
|
|
return n
|
|
except Exception as exc:
|
|
logger.warning(
|
|
"embed_extraction: DB write failed for extraction %s: %s",
|
|
extraction_id,
|
|
exc,
|
|
)
|
|
return None
|
|
|
|
|
|
def _upsert_embeddings(
|
|
conn: Any,
|
|
extraction_id: int,
|
|
chunk_embeddings: list[tuple[Chunk, list[float]]],
|
|
model: str,
|
|
) -> int:
|
|
"""Replace all public.embeddings rows for this (extraction, model) pair."""
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
DELETE FROM public.embeddings
|
|
WHERE source_schema = %s
|
|
AND source_table = %s
|
|
AND source_id = %s
|
|
AND embedding_model = %s
|
|
""",
|
|
(SOURCE_SCHEMA, SOURCE_TABLE, extraction_id, model),
|
|
)
|
|
for chunk, vec in chunk_embeddings:
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO public.embeddings
|
|
(source_schema, source_table, source_id, chunk_index,
|
|
chunk_text, chunk_token_count, embedding, embedding_model)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s::vector, %s)
|
|
""",
|
|
(
|
|
SOURCE_SCHEMA,
|
|
SOURCE_TABLE,
|
|
extraction_id,
|
|
chunk.index,
|
|
chunk.text,
|
|
chunk.token_count,
|
|
_vector_literal(vec),
|
|
model,
|
|
),
|
|
)
|
|
return len(chunk_embeddings)
|
|
|
|
|
|
__all__ = ["embed_extraction", "embed_text", "close_pool"]
|