Skip to content

Make cosine the only similarity, and add get_cosine_similarity (speedkick) - #1593

Closed
edwinyyyu wants to merge 8 commits into
MemMachine:speedkickfrom
edwinyyyu:feat/cosine-only-vector-store-speedkick
Closed

edwinyyyu wants to merge 8 commits into
MemMachine:speedkickfrom
edwinyyyu:feat/cosine-only-vector-store-speedkick

Conversation

@edwinyyyu

Copy link
Copy Markdown
Contributor

Stacked on #1591 (which is itself stacked on #1588) — only the last commit, Make cosine the only similarity, and read it back per record, belongs to this PR.

Follows the default branch, where this removal has already been made.

Purpose of the change

Every embedder MemMachine ships produces vectors meant to be compared by cosine similarity. OpenAIEmbedder hard-coded COSINE; AmazonBedrockEmbedder defaulted to it; SentenceTransformerEmbedder only reported whatever the model declared. SimilarityMetric let the rest of the stack ask which of four metrics was in play, and the answer was cosine everywhere — while every layer that touched a score paid for the other three: direction flags threaded through scoring, thresholds applied in whichever direction the metric implied, and a per-backend table in each store mapping the enum onto native metric names.

Description

The enum is gone. Scores are cosine similarities in [-1, 1], higher is better, and the names say so:

Before After
QueryMatch.score QueryMatch.cosine_similarity
SearchMatch.score SearchMatch.cosine_similarity
query(score_threshold=…) query(min_cosine_similarity=…)
compute_similarity(a, b, metric) compute_cosine_similarity(a, b)
VectorStoreCollectionConfig.similarity_metric removed
Embedder.similarity_metric removed
similarity_gt(metric) removed (no callers)
VectorSearchEngineFactory = (int, SimilarityMetric) -> … (int) -> …

min_cosine_similarity no longer needs a direction to be meaningful — that is the point of the rename, not just tidier spelling.

And get_cosine_similarity arrives. With one metric there is finally something worth reading back per record, so this adds VectorStoreCollection.get_cosine_similarity and VectorSearchEngine.get_cosine_similarities. Every backend can serve it, by one of two routes:

  • Keyed vector access — hnswlib and usearch gather stored vectors and score locally; sqlite-vec pushes it into SQL (vec_distance_cosine per rowid, one point query each because vec0 plans rowid IN (…) as a full scan); Qdrant and Milvus retrieve by id, which is exact where an id-filtered ANN query is not guaranteed to return every matching point.
  • Comprehensive filtered search — turbovec cannot hand a vector back at all (it stores only TurboQuant codes), so it leans on its filtered search, which stops early only once it already holds limit matches; asking for limit = len(keys) therefore returns every key present.

Similarities may be computed from quantized stored vectors, so the contract says plainly that they may not match a similarity computed from a fresh embedding — which covers both turbovec's estimate and hnswlib's normalize-on-insert.

Two consequences beyond the vector store

default's packages/core has no vector graph store, so these two have no counterpart there and are worth reviewing on their own terms.

  1. The vector graph stores carried a metric per stored embedding, as a companion similarity_metric_for_<name> property written beside every vector, in both Neo4j and NebulaGraph. That property is gone, and Node.embeddings / Edge.embeddings are dict[str, list[float]] rather than dict[str, tuple[list[float], SimilarityMetric]]. Existing graph data keeps the orphaned companion properties; nothing reads them.

  2. NebulaGraph's similarity search is now always exact KNN. Its vector indexes offer only L2 and IP, and its cosine() does not support APPROXIMATE — so with cosine as the only similarity, no index it can build serves a query. use_ann could never be true and _create_vector_index_if_not_exists could only ever take its skip path, so both are removed rather than left as unreachable code. This is a real behavioral narrowing for NebulaGraph deployments that were relying on an L2 or IP index. The now-unused ANN tuning parameters (ann_index_type, ivf_nprobe, hnsw_ef_search, force_exact_similarity_search) are deliberately left in place for a separate change, since they are user-visible configuration.

Tests that existed only to pin multi-metric behavior are removed: euclidean/dot search sections in the engine suites, euclidean ordering and threshold-direction tests in the store suites, the _similarity_metric_to_nebula mapping table test, and the euclidean threshold-inversion regressions in the event-backend wiring suite (that inversion is no longer reachable).

Verification

  • pytest packages/server/server_tests: 1916 passed, 3 skipped. The one failure, test_get_version, is a git-describe version-string artifact of the local checkout and fails identically on untouched speedkick (0.3.9.post2.dev5+g8dd1c0178); see style: reformat an assert ruff 0.15.14 formats differently (speedkick) #1592's sibling discussion.
  • ty check --project packages/server: 17 diagnostics, the same 17 as the base branch.
  • ruff check / ruff format --check: clean.

🤖 Generated with Claude Code

https://claude.ai/code/session_01NKmF9xNph9QH3ozNw3ZnJL

@edwinyyyu
edwinyyyu force-pushed the feat/cosine-only-vector-store-speedkick branch 8 times, most recently from d229e74 to 1e2361a Compare September 9, 2026 18:50
edwinyyyu and others added 8 commits September 9, 2026 11:52
turbovec keeps TurboQuant-compressed vectors in RAM, so an index is a
fraction of the size of the f32 engines' and a full scan has a much
smaller scaling constant than sqlite-vec's on-disk one. Search is
approximate as a result: scores land near the exact value rather than on
it, and the tests assert ranking and membership instead of magnitudes.

Two consequences of storing only compressed vectors are worth naming.
`get_vectors` raises `NotImplementedError`, since the originals are not
recoverable -- which makes this engine usable by EventMemory but not by
semantic memory, whose feature updates read stored embeddings back. And
removal is exact and cheap: turbovec drops the id from its map rather
than tombstoning, so a deleted key cannot resurface in results.

`save` publishes through the shared atomic-write helper, like the other
engines, so an interrupted save leaves the previously published index
intact rather than a truncated file the store would treat as a hard
load error. `load` clears any temp file a previous save left behind.

Signed-off-by: Edwin Yu <[email protected]>
Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
1.0.0 is upstream's first stable release, and what it commits to is the
on-disk format: v7 is the only container turbovec reads or writes, and a
file written by this release stays readable by later ones. The extra
floors there rather than at 0.7.0.

v7 is also what makes `sync` possible, and that changes how this engine
saves. `save` no longer routes through the shared `atomic_index_write`
helper. turbovec publishes the index itself: a checkpoint appends what
changed since the last one rather than restating the whole index, and
commits it durably -- a crash at any byte leaves the previous commit
intact, and unlike the helper's rename the publication survives a power
failure. Wrapping that would restate the index on every checkpoint and
publish it through the weaker of the two protocols. `load` drops
`clear_stale_index_temp` with it, since the engine no longer writes the
`<path>.tmp` sibling it cleared.

The tests follow. Save leaving no temp file is now an assertion about
the whole directory, since turbovec names its own temp; the stale-temp
test goes with the protocol it tested; and a new test pins what a later
checkpoint carries, reading a bulk removal and an append back through a
fresh engine.

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Signed-off-by: Edwin Yu <[email protected]>
`SearchMatch` documents a cosine score as a cosine similarity in [-1, 1],
and `SQLiteVectorStoreCollection` compares `score_threshold` against it
directly. turbovec's score is a quantized inner product, so it lands a
hair outside that range: at 8 dimensions and the default bit width, 40 of
50 self-matches score above 1.0 (max 1.0066), and at bit width 2 they
exceed it at every dimension measured (1.0038 at 768, 1.0025 at 1536).
Publishing 1.0066 as a cosine similarity is a broken promise, however
small, so the hair is clamped where the match is built.

A dot product is unbounded by contract and passes through untouched.

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Signed-off-by: Edwin Yu <[email protected]>
turbovec indexes a dimensionality that is a positive multiple of 8, so the
engine refused a width every other engine in the tree accepts -- 300, the
width of a spaCy or word2vec vector, among them -- with turbovec's own
error rather than a store-level one.

The width is now rounded up and vectors are written into the leading
columns of a zeroed buffer. Padding is exact for both metrics this engine
serves: a zero coordinate adds nothing to an inner product and nothing to
an L2 norm, so the padded index answers as the unpadded one would. The
same write is what rejects a wrong-width vector, which previously reached
turbovec as a shape it would report in its own terms.

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Signed-off-by: Edwin Yu <[email protected]>
`prepare()` warms turbovec's per-index caches and its lazy id-to-slot map
so the first search, contains or remove after a write does not pay a
one-time cost. Against turbovec 0.8.0 that cost was the whole SIMD-blocked
code layout, rebuilt per add and proportional to the index rather than the
batch -- 199 ms at 100k vectors, which the first unlucky search paid if
this call did not. Upstream made that re-convergence incremental in 1.0.0,
and the call now costs 12-20 us and buys nothing measurable: at 100k a
search after a one-vector add reads 0.52 ms with it and 0.51 ms without,
and in a cold process the first search reads 0.10 ms against 0.09 ms.

The remaining warm it offers -- the id-to-slot map after a load, worth
0.5 ms off the first remove at 100k -- is not on this path, and the map
materializes safely under concurrent readers regardless.

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Signed-off-by: Edwin Yu <[email protected]>
`VectorStoreCollection.get` had one production caller: the semantic
storage read a feature's stored embedding back so it could write that
same embedding again alongside fresh properties, because `upsert` demands
a whole record. `return_vector=True` was passed at exactly that one call
site, and `VectorSearchEngine.get_vectors` existed to serve it.

Replace it with `set_properties`, which says what the caller wanted:
replace a record's properties and leave its vector where it is. On SQLite
and sqlite-vec that is a plain UPDATE on the records table that never
touches the index. On Qdrant it is an `overwrite_payload` selected by
filter rather than by id -- an id list raises on an id the collection
does not hold, and would reach a point another logical collection owns in
the same native one. Milvus has no partial update, so the read-back
survives there, inside the one backend that needs it.

With no caller left, `get`, `return_vector` and `get_vectors` all go.
turbovec could not implement `get_vectors` at all -- it stores only
TurboQuant-compressed vectors -- so this also retires the one contract
method a conforming engine was allowed to raise on.

Two consequences worth naming:

- A feature's embedding is a function of its `value` alone, on add and on
  update alike, so keeping the stored vector when `value` did not change
  is the behavior that was already there. What changes is that it is no
  longer expressed as a read-modify-write through the index.

- `get` was the only read that could see a record whose vector the index
  lost to a reverted publication. `query` goes through the index, so such
  a record is now invisible through the contract; the row survives, and
  the durability test reads the row directly to say so.

Tests read records back through `query` instead, via a `_fetch_records`
helper per suite.

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01NKmF9xNph9QH3ozNw3ZnJL
Every embedder MemMachine ships produces vectors meant to be compared by
cosine similarity: OpenAI hard-coded COSINE, Bedrock defaulted to it, and
the SentenceTransformer embedder only reported what the model declared.
`SimilarityMetric` let the rest of the stack ask which of four metrics was
in play, and the answer was cosine everywhere -- while every layer that
touched a score paid for the other three: direction flags threaded through
scoring, thresholds applied in whichever direction the metric implied, and
per-backend tables mapping the enum onto native metric names.

Drop the enum. Scores become cosine similarities in [-1, 1], higher is
better, and the names say so: `QueryMatch.score` and `SearchMatch.score`
become `cosine_similarity`, and `query(score_threshold=...)` becomes
`query(min_cosine_similarity=...)`, which no longer needs a direction to
be meaningful. `compute_similarity(a, b, metric)` becomes
`compute_cosine_similarity(a, b)`, and `similarity_gt` -- a comparator
picked by metric, with no callers -- goes with it.

With one metric there is also something worth reading back per record, so
this adds `VectorStoreCollection.get_cosine_similarity` and
`VectorSearchEngine.get_cosine_similarities`. Engines with keyed vector
access (hnswlib, usearch) gather and score locally; sqlite-vec pushes it
into SQL; Qdrant and Milvus retrieve by id, which is exact where an
id-filtered ANN query would not be; turbovec, which cannot hand a vector
back, leans on its filtered search, comprehensive when asked for as many
matches as there are keys.

Two consequences beyond the vector store:

- The vector graph stores carried a metric per stored embedding, as a
  companion property written beside every vector. That property is gone,
  and `Node.embeddings` / `Edge.embeddings` are plain vectors.

- NebulaGraph indexes only L2 and IP, and its `cosine()` does not support
  APPROXIMATE, so with cosine the only similarity no vector index it can
  build serves a query. Its ANN search branch and vector index creation
  could no longer run, and are removed; search there is always exact KNN.
  The now-unused ANN tuning parameters are left for a separate change.

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01NKmF9xNph9QH3ozNw3ZnJL
`test_search_similar_nodes` stored the same two vectors under `embedding1`
and `embedding2` and told them apart by metric -- cosine on one, euclidean
on the other -- so searching each name gave a different nearest node. With
one metric the two became byte-identical, and the assertion that
`embedding2` ranks Node2 first started failing: under cosine the query
[1, 0] is exactly aligned with Node1's [1000, 0].

Giving them identical vectors and correcting the expectation would have
made `embedding2` prove nothing `embedding1` does not. Swapping Node1's
and Node2's `embedding2` vectors restores what the metric used to supply:
the two names disagree about which node is nearest, so searching by name
has to select the right vector to get the right answer. The assertions
stand as written.

Verified against a real Neo4j container, and the whole `common/`
integration suite passes locally: 292 passed, 97 skipped.

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01NKmF9xNph9QH3ozNw3ZnJL
@edwinyyyu

Copy link
Copy Markdown
Contributor Author

Superseded by #1598, which combines the cosine-only change with the read-back replacement so the implementation moves rather than being deleted and re-added.

@edwinyyyu edwinyyyu closed this Sep 9, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant