diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md index 16f257cb..b4dfc93d 100644 --- a/python/CHANGELOG.md +++ b/python/CHANGELOG.md @@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- **Semantic search works in the Python client (Hermes).** Harrier, the embedding model, takes the last token's state as the sentence embedding, so each input must end with ``. The ONNX repo's tokenizer config leaves out `add_eos_token`, and `transformers` then drops ``, so every Python embedding was the state of the text's last word. Recall leaned on near-random similarity: on a 13-fact test vault, the right fact ranked first for 4 of 12 questions ("vegetarian" put "User is vegetarian." 11th of 13); with `` it ranks first for all 12. Near-duplicate and contradiction checks use the same embeddings, so they get better too. The TS clients (OpenClaw plugin, MCP server) already appended ``; Python embeddings now match theirs exactly (cosine 1.0000, was about 0.4), so facts written by one client are found by the other. Facts stored earlier keep their old embedding: with fixed queries they rank first for 9 of 12 questions in the same test, and every newly stored fact gets a correct embedding. - **Asking the agent to update TotalReclaw ends with "send `/restart`", not shell commands.** The usage skill treated "upgrade" as buying Pro and said nothing about installing a newer version, so an agent asked "upgrade TotalReclaw to rc6" improvised: it installed the package, tried to restart the gateway itself (Hermes blocks that from agent sessions), wrote a restart script and told the user to run it over SSH. The skill now separates a version update (follow the guide's Update prompt) from Pro, and says how to restart: the agent can't, so in a chat app it asks the user to send `/restart`; only a terminal-CLI user gets `hermes gateway restart`. - **Automatic recall before each turn is fast enough for Hermes again.** Hermes gives the memory provider 8 seconds to fetch context for a turn, then goes ahead without it. TotalReclaw split the user's message into trapdoor batches (about 11 for a typical chat message) and queried them one after another, and each batch carried the full encrypted fact and its embedding, so the same facts were downloaded many times: about 10 MB and 6–20 s per turn on a real vault. Most Telegram turns ran without recalled memories. Batches now return fact ids only and run in parallel (up to 6 at a time), each matching fact's payload is fetched once, and automatic recall skips embeddings it can't use (it ranks without a query vector). On the same vault, results are identical and a typical message now takes about 3 s and 1.5 MB. A slow prefetch (3 s or more) logs its duration. - **Hermes keeps TotalReclaw on for the turn right after pairing.** Pairing writes the credentials from a separate process. Hermes asks the provider whether it is available at the start of the next turn, before the plugin re-reads its credentials, and on "no" it disabled external memory for the rest of that session with "Memory provider 'totalreclaw' is selected but reports unavailable" (until `/new`). The availability check now re-reads the credentials file first (a local read, no network). diff --git a/python/src/totalreclaw/embedding.py b/python/src/totalreclaw/embedding.py index b4cc3eb4..1263e57f 100644 --- a/python/src/totalreclaw/embedding.py +++ b/python/src/totalreclaw/embedding.py @@ -6,9 +6,13 @@ Model: onnx-community/harrier-oss-v1-270m-ONNX (q4 quantized) Output: pre-pooled sentence_embedding (no manual pooling needed) +Pooling: last token, so every input must end with . The ONNX repo's +tokenizer_config.json omits ``add_eos_token``, and transformers' Gemma +tokenizer then drops , which leaves the embedding as the last word's +state. We force it on; the output then matches the TS clients +(transformers.js reads tokenizer.json, which appends ). Normalization: already L2-normalized Dimensions: 640 -No instruction prefix needed. """ from __future__ import annotations @@ -124,7 +128,7 @@ def _ensure_loaded(): sess_options = ort.SessionOptions() sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL _session = ort.InferenceSession(model_path, sess_options) - _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) + _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, add_eos_token=True) if needs_download: _emit_download_done() @@ -144,9 +148,20 @@ def get_embedding(text: str) -> list[float]: max_length=512, ) + input_ids = inputs["input_ids"].astype(np.int64) + attention_mask = inputs["attention_mask"].astype(np.int64) + eos_id = tokenizer.eos_token_id + if eos_id is not None and input_ids[0, -1] != eos_id: + # A transformers version that ignores add_eos_token: append it here. + if input_ids.shape[1] >= 512: + input_ids[0, -1] = eos_id + else: + input_ids = np.concatenate([input_ids, [[eos_id]]], axis=1).astype(np.int64) + attention_mask = np.concatenate([attention_mask, [[1]]], axis=1).astype(np.int64) + feed: dict[str, np.ndarray] = { - "input_ids": inputs["input_ids"].astype(np.int64), - "attention_mask": inputs["attention_mask"].astype(np.int64), + "input_ids": input_ids, + "attention_mask": attention_mask, } output_names = [o.name for o in session.get_outputs()] diff --git a/python/tests/test_embedding.py b/python/tests/test_embedding.py index 36dab38d..afde39ec 100644 --- a/python/tests/test_embedding.py +++ b/python/tests/test_embedding.py @@ -69,3 +69,27 @@ def test_batch(self, embed, dims): results = get_embeddings_batch(["hello", "world"]) assert len(results) == 2 assert all(len(e) == dims for e in results) + + +class TestRetrievalQuality: + """Harrier pools the last token: without these rank the wrong fact first.""" + + DOCS = [ + "User moved from Lisbon to Porto and works at Tidewater.", + "User has a 1:1 with Ines every Monday at 10:00.", + "User uses a ThinkPad X1 Carbon running Fedora.", + "User goes bouldering on weekends.", + "User uses Bitwarden as a password manager.", + "User is vegetarian.", + ] + + @pytest.mark.parametrize( + "query,want", + [("vegetarian", "vegetarian"), ("what laptop do I use", "ThinkPad"), ("what do I do on weekends", "bouldering")], + ) + def test_query_finds_the_right_fact(self, query, want): + from totalreclaw.embedding import get_embedding + + q = get_embedding(query) + scored = sorted(self.DOCS, key=lambda d: -sum(a * b for a, b in zip(q, get_embedding(d)))) + assert want in scored[0] diff --git a/python/tests/test_embedding_eos.py b/python/tests/test_embedding_eos.py new file mode 100644 index 00000000..27427491 --- /dev/null +++ b/python/tests/test_embedding_eos.py @@ -0,0 +1,97 @@ +"""Harrier pools the last token, so every input must end with . + +The ONNX repo's tokenizer config omits ``add_eos_token`` and transformers' +Gemma tokenizer then drops . These tests run without the model: a fake +tokenizer and session record what reaches the ONNX graph. +""" +from __future__ import annotations + +import sys +import types + +import numpy as np + +from totalreclaw import embedding + +EOS = 1 + + +class _Tok: + eos_token_id = EOS + + def __init__(self, ids): + self._ids = ids + + def __call__(self, text, **kw): + ids = np.array([self._ids], dtype=np.int64) + return {"input_ids": ids, "attention_mask": np.ones_like(ids)} + + +class _Out: + def __init__(self, name): + self.name = name + + +class _Session: + def __init__(self): + self.feed = None + + def get_outputs(self): + return [_Out("sentence_embedding")] + + def run(self, names, feed): + self.feed = feed + return [np.ones((1, embedding.EMBEDDING_DIMS), dtype=np.float32)] + + +def _run(monkeypatch, ids): + sess = _Session() + monkeypatch.setattr(embedding, "_session", sess) + monkeypatch.setattr(embedding, "_tokenizer", _Tok(ids)) + embedding.get_embedding("hello world") + return sess.feed + + +def test_appends_eos_when_tokenizer_drops_it(monkeypatch): + feed = _run(monkeypatch, [2, 23391, 1902]) + assert feed["input_ids"].tolist() == [[2, 23391, 1902, EOS]] + assert feed["attention_mask"].tolist() == [[1, 1, 1, 1]] + assert feed["input_ids"].dtype == np.int64 + + +def test_leaves_input_alone_when_eos_present(monkeypatch): + feed = _run(monkeypatch, [2, 23391, 1902, EOS]) + assert feed["input_ids"].tolist() == [[2, 23391, 1902, EOS]] + + +def test_full_length_input_ends_with_eos(monkeypatch): + feed = _run(monkeypatch, [2] + [7] * 511) + assert feed["input_ids"].shape == (1, 512) + assert feed["input_ids"][0, -1] == EOS + + +def test_loader_asks_the_tokenizer_for_eos(monkeypatch): + calls = {} + + class _AutoTokenizer: + @staticmethod + def from_pretrained(model_id, **kw): + calls["kw"] = kw + return _Tok([2, EOS]) + + fake_ort = types.SimpleNamespace( + SessionOptions=lambda: types.SimpleNamespace(graph_optimization_level=None), + GraphOptimizationLevel=types.SimpleNamespace(ORT_ENABLE_ALL=99), + InferenceSession=lambda path, opts: _Session(), + ) + monkeypatch.setitem(sys.modules, "onnxruntime", fake_ort) + monkeypatch.setitem(sys.modules, "transformers", types.SimpleNamespace(AutoTokenizer=_AutoTokenizer)) + monkeypatch.setitem( + sys.modules, "huggingface_hub", types.SimpleNamespace(hf_hub_download=lambda **kw: "/fake/model.onnx") + ) + monkeypatch.setattr(embedding, "_hf_cache_has_model", lambda: True) + monkeypatch.setattr(embedding, "_session", None) + monkeypatch.setattr(embedding, "_tokenizer", None) + + embedding._ensure_loaded() + assert calls["kw"].get("add_eos_token") is True