embed.py

"""Embed text chunks via the local Ollama nomic-embed-text model (BATCHED).

This is the cartridge's INPUT ENCODER (Aaron's existing Ollama pipeline, reused). Raw text
in -> 768-d vector out; the JEPA never sees the host or the tokenizer.

CACHE (binary, UE-scale): the original single-JSON cache rewrote the whole file every batch —
fine at 2k chunks, quadratic death at 350k (multi-GB JSON). Now: keys in .embed_cache_keys.json
(list of sha1s, row order) + vectors in .embed_cache_vecs.npy (float32 [N,768]); saves are
periodic sequential binary writes. A legacy .embed_cache.json is migrated once, then renamed.

ROBUSTNESS: an HTTP 400 from Ollama (oversized payload / poison chunk) bisects the batch down
to the offending chunk(s), which get a zero vector (cached, logged to stderr) instead of
killing an hours-long warm run.

Progress goes to STDERR — on a stdio MCP transport stdout is the JSONRPC wire.
"""

from __future__ import annotations

import hashlib
import json
import sys
import urllib.error
import urllib.request
from pathlib import Path

import numpy as np

EMBED_URL = "http://localhost:11434/api/embed"
MODEL = "nomic-embed-text"
_HERE = Path(__file__).resolve().parent
KEYS_PATH = _HERE / ".embed_cache_keys.json"
VECS_PATH = _HERE / ".embed_cache_vecs.npy"
LEGACY_PATH = _HERE / ".embed_cache.json"
MAX_CHARS = 6000     # guard against pathological long chunks blowing past model context
DIM = 768
SAVE_EVERY = 40      # batches between periodic saves (crash loses at most this much work)


def _log(msg: str) -> None:
    print(msg, file=sys.stderr, flush=True)


def _embed_batch(texts: list[str], model: str, keep_alive: str = "10m") -> list[list[float]]:
    payload = json.dumps(
        {"model": model, "input": [t[:MAX_CHARS] for t in texts], "keep_alive": keep_alive}
    ).encode()
    req = urllib.request.Request(
        EMBED_URL, data=payload, headers={"Content-Type": "application/json"}
    )
    with urllib.request.urlopen(req, timeout=300) as resp:
        return json.loads(resp.read())["embeddings"]


def _embed_robust(texts: list[str], model: str) -> list[list[float]]:
    """Embed a batch; on HTTP error bisect down to isolate poison chunks, which get zeros."""
    try:
        return _embed_batch(texts, model)
    except (urllib.error.HTTPError, urllib.error.URLError) as e:
        if len(texts) == 1:
            _log(f"  ! poison chunk skipped (zero vector): {e} — {texts[0][:80]!r}")
            return [[0.0] * DIM]
        mid = len(texts) // 2
        return _embed_robust(texts[:mid], model) + _embed_robust(texts[mid:], model)


class Embedder:
    def __init__(self, model: str = MODEL):
        self.model = model
        self.key2row: dict[str, int] = {}
        self.rows: list[np.ndarray] = []
        if KEYS_PATH.exists() and VECS_PATH.exists():
            keys = json.loads(KEYS_PATH.read_text())
            mat = np.load(VECS_PATH)
            self.rows = [mat[i] for i in range(mat.shape[0])]
            self.key2row = {k: i for i, k in enumerate(keys)}
        elif LEGACY_PATH.exists():
            _log("  migrating legacy JSON embed cache to binary (one-time)...")
            legacy = json.loads(LEGACY_PATH.read_text())
            for k, v in legacy.items():
                self.key2row[k] = len(self.rows)
                self.rows.append(np.asarray(v, dtype=np.float32))
            self._save()
            LEGACY_PATH.rename(LEGACY_PATH.with_suffix(".json.migrated"))
            _log(f"  migrated {len(self.rows)} cached vectors")

    def _key(self, text: str) -> str:
        return hashlib.sha1(f"{self.model}\x00{text}".encode()).hexdigest()

    def _save(self) -> None:
        keys = [None] * len(self.rows)
        for k, i in self.key2row.items():
            keys[i] = k
        tmp_v = VECS_PATH.with_name("embed_cache_vecs.tmp.npy")   # np.save appends .npy unless present
        np.save(tmp_v, np.asarray(self.rows, dtype=np.float32))
        tmp_v.replace(VECS_PATH)
        tmp_k = KEYS_PATH.with_suffix(".json.tmp")
        tmp_k.write_text(json.dumps(keys))
        tmp_k.replace(KEYS_PATH)

    def embed(self, texts: list[str], batch_size: int = 64) -> np.ndarray:
        keys = [self._key(t) for t in texts]
        missing = [i for i, k in enumerate(keys) if k not in self.key2row]
        if missing:
            _log(f"  embedding {len(missing)} new chunks ({len(texts) - len(missing)} cached)...")
        for n, start in enumerate(range(0, len(missing), batch_size)):
            idx = missing[start: start + batch_size]
            vecs = _embed_robust([texts[i] for i in idx], self.model)
            for i, v in zip(idx, vecs):
                self.key2row[keys[i]] = len(self.rows)
                self.rows.append(np.asarray(v, dtype=np.float32))
            if (n + 1) % SAVE_EVERY == 0:
                self._save()
                _log(f"    {min(start + batch_size, len(missing))}/{len(missing)} (saved)")
        if missing:
            self._save()
        return np.stack([self.rows[self.key2row[k]] for k in keys]) if keys else np.zeros((0, DIM), np.float32)


if __name__ == "__main__":
    emb = Embedder()
    v = emb.embed([
        "URiftInventoryComponent::AddItem grants an item and fires OnItemAdded.",
        "void URiftInventoryComponent::AddItem(URiftItemDefinition* Def) {}",
        "The mitochondria is the powerhouse of the cell.",
    ])
    print(f"shape={v.shape} dtype={v.dtype}")