Add vector embedding support and hybrid search to ToolStore by aponcedeleonch · Pull Request #3808 · stacklok/toolhive

aponcedeleonch · 2026-02-12T17:19:01Z

Closes: #3732

Extend the SQLite-backed ToolStore with optional embedding-based semantic search. When an EmbeddingClient is provided, Search runs FTS5 and cosine similarity in parallel via errgroup, merges and deduplicates results by keeping the lower distance score, and caps output at 4 results.

Key additions:

EmbeddingClient interface in internal/types (Embed, EmbedBatch, Dimension, Close)
FakeEmbeddingClient using SHA-256 seeded RNG with L2-normalized vectors
Pure Go cosine similarity/distance in internal/similarity package
Binary encode/decode helpers for embedding BLOB storage
normalizeBM25 rescaled to [0, 2) range to align with cosine distance
Comprehensive unit tests, concurrency tests, and benchmarks

Large PR Justification

Multiple related changes that would break if separated. This PR introduces semantic search and a wait to combine results from full text search and semantic

github-actions

Large PR Detected

This PR exceeds 1000 lines of changes and requires justification before it can be reviewed.

How to unblock this PR:

Add a section to your PR description with the following format:

## Large PR Justification

[Explain why this PR must be large, such as:]
- Generated code that cannot be split
- Large refactoring that must be atomic
- Multiple related changes that would break if separated
- Migration or data transformation

Alternative:

Consider splitting this PR into smaller, focused changes (< 1000 lines each) for easier review and reduced risk.

See our Contributing Guidelines for more details.

This review will be automatically dismissed once you add the justification section.

codecov · 2026-02-12T17:31:39Z

Codecov Report

❌ Patch coverage is 86.30952% with 23 lines in your changes missing coverage. Please review.
✅ Project coverage is 66.91%. Comparing base (f1772c6) to head (09908e0).

Files with missing lines	Patch %	Lines
...cp/optimizer/internal/sqlite_store/sqlite_store.go	84.42%	9 Missing and 10 partials ⚠️
pkg/vmcp/optimizer/fake_embedding.go	86.20%	3 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3808      +/-   ##
==========================================
+ Coverage   66.84%   66.91%   +0.07%     
==========================================
  Files         439      441       +2     
  Lines       43509    43660     +151     
==========================================
+ Hits        29083    29216     +133     
- Misses      12175    12182       +7     
- Partials     2251     2262      +11

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

jerm-dro

The diff is quite large due to the conflicts with the base branch. The structure looks good, but I will take a closer look when the base branch is merged.

pkg/vmcp/optimizer/internal/similarity/cosine_bench_test.go

Large PR justification has been provided. Thank you!

github-actions · 2026-02-13T09:01:37Z

✅ Large PR justification has been provided. The size review has been dismissed and this PR can now proceed with normal review.

pkg/vmcp/optimizer/internal/similarity/cosine_bench_test.go

pkg/vmcp/optimizer/internal/sqlite_store/BENCHMARKS.md

Extend the SQLite-backed ToolStore with optional embedding-based semantic search. When an EmbeddingClient is provided, Search runs FTS5 and cosine similarity in parallel via errgroup, merges and deduplicates results by keeping the lower distance score, and caps output at 4 results. Key additions: - EmbeddingClient interface in internal/types (Embed, EmbedBatch, Dimension, Close) - FakeEmbeddingClient using SHA-256 seeded RNG with L2-normalized vectors - Pure Go cosine similarity/distance in internal/similarity package - Binary encode/decode helpers for embedding BLOB storage - normalizeBM25 rescaled to [0, 2) range to align with cosine distance - Comprehensive unit tests, concurrency tests, and benchmarks Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

pkg/vmcp/optimizer/internal/similarity/cosine_bench_test.go

jerm-dro · 2026-02-13T22:03:53Z

pkg/vmcp/optimizer/internal/sqlite_store/sqlite_store.go

 		  AND t.name IN (SELECT value FROM json_each(?))
-		ORDER BY rank`
+		ORDER BY rank
+		LIMIT ?`


Why limit? What happens if you have a hybrid search and a row's rank is low and below the limit, but the cosine similarity is high?

jerm-dro · 2026-02-13T22:05:51Z

pkg/vmcp/optimizer/internal/sqlite_store/sqlite_store.go

+	queryStr := `SELECT name, description, embedding
+		FROM llm_capabilities
+		WHERE embedding IS NOT NULL
+		  AND name IN (SELECT value FROM json_each(?))`
+
+	rows, err := s.db.QueryContext(ctx, queryStr, string(allowedJSON))


Blocker, question: Do we really need to query the DB twice?

Why can't we:

Load the allowed tools, their rank, and embeddings.

Compute the cosine similarity for each returned row

Compute the merged score

Return the top K results based on score.

jerm-dro · 2026-02-13T22:08:44Z

pkg/vmcp/optimizer/internal/types/types.go

+	// Dimension returns the dimensionality of the embeddings produced.
+	Dimension() int


suggestion: delete - you can just gather this from the returned embedding?

jerm-dro · 2026-02-13T22:09:05Z

pkg/vmcp/optimizer/embedding.go

+// EmbeddingClient generates vector embeddings from text.
+// It is defined in the internal/types package and aliased here so that
+// external consumers continue to use optimizer.EmbeddingClient.
+type EmbeddingClient = types.EmbeddingClient


Question, blocking: why does this need to be re-exported?

Suggestion: You may be able to avoid re-exporting things by making the entrypoint to this package, NewSQLiteToolStore, accept the config object. When that function is called, it can freely reference things in the internal packages to construct objects.

jerm-dro · 2026-02-13T22:10:12Z

pkg/vmcp/optimizer/fake_embedding.go

+// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package optimizer


suggestion: move this embedding related code to the similarity package it seems very related, and the actual CosineSimilarity function is an implementation detail.

aponcedeleonch requested review from JAORMX, amirejaz, jerm-dro, jhrozek and yrobla as code owners February 12, 2026 17:19

github-actions bot added the size/XL Extra large PR: 1000+ lines changed label Feb 12, 2026

github-actions bot previously requested changes Feb 12, 2026

View reviewed changes

aponcedeleonch marked this pull request as draft February 12, 2026 17:20

aponcedeleonch force-pushed the issue-3732/semantic-search-store branch from c06effa to 5d93afe Compare February 12, 2026 17:21

github-actions bot added size/XL Extra large PR: 1000+ lines changed and removed size/XL Extra large PR: 1000+ lines changed labels Feb 12, 2026

aponcedeleonch force-pushed the issue-3731/sqlite-fts5-store branch from 631f1d6 to 85a4c84 Compare February 12, 2026 20:00

jerm-dro reviewed Feb 12, 2026

View reviewed changes

pkg/vmcp/optimizer/internal/similarity/cosine_bench_test.go Show resolved Hide resolved

Base automatically changed from issue-3731/sqlite-fts5-store to main February 13, 2026 06:33

aponcedeleonch force-pushed the issue-3732/semantic-search-store branch from 5d93afe to 0fa9e53 Compare February 13, 2026 09:01

aponcedeleonch marked this pull request as ready for review February 13, 2026 09:01

github-actions bot added size/XL Extra large PR: 1000+ lines changed and removed size/XL Extra large PR: 1000+ lines changed labels Feb 13, 2026

aponcedeleonch requested a review from jerm-dro February 13, 2026 09:01

aponcedeleonch force-pushed the issue-3732/semantic-search-store branch from 0fa9e53 to d17bbb0 Compare February 13, 2026 09:21

github-actions bot added size/XL Extra large PR: 1000+ lines changed and removed size/XL Extra large PR: 1000+ lines changed labels Feb 13, 2026

aponcedeleonch force-pushed the issue-3732/semantic-search-store branch from d17bbb0 to 53657a3 Compare February 13, 2026 09:24

github-actions bot added size/XL Extra large PR: 1000+ lines changed and removed size/XL Extra large PR: 1000+ lines changed labels Feb 13, 2026

aponcedeleonch force-pushed the issue-3732/semantic-search-store branch from 53657a3 to 39ad0e2 Compare February 13, 2026 09:27

github-actions bot removed the size/XL Extra large PR: 1000+ lines changed label Feb 13, 2026

github-actions bot added the size/XL Extra large PR: 1000+ lines changed label Feb 13, 2026

yrobla reviewed Feb 13, 2026

View reviewed changes

pkg/vmcp/optimizer/internal/similarity/cosine_bench_test.go Show resolved Hide resolved

yrobla reviewed Feb 13, 2026

View reviewed changes

pkg/vmcp/optimizer/internal/sqlite_store/BENCHMARKS.md Show resolved Hide resolved

aponcedeleonch and others added 2 commits February 13, 2026 17:56

Added note about benchmarking

09908e0

aponcedeleonch force-pushed the issue-3732/semantic-search-store branch from 39ad0e2 to 09908e0 Compare February 13, 2026 15:59

github-actions bot added size/XL Extra large PR: 1000+ lines changed and removed size/XL Extra large PR: 1000+ lines changed labels Feb 13, 2026

jerm-dro reviewed Feb 13, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add vector embedding support and hybrid search to ToolStore#3808

Add vector embedding support and hybrid search to ToolStore#3808
aponcedeleonch wants to merge 2 commits intomainfrom
issue-3732/semantic-search-store

aponcedeleonch commented Feb 12, 2026 •

edited

Loading

Uh oh!

github-actions bot left a comment

Uh oh!

codecov bot commented Feb 12, 2026 •

edited

Loading

Uh oh!

jerm-dro left a comment

Uh oh!

Uh oh!

github-actions bot commented Feb 13, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jerm-dro Feb 13, 2026

Uh oh!

jerm-dro Feb 13, 2026

Uh oh!

jerm-dro Feb 13, 2026

Uh oh!

jerm-dro Feb 13, 2026

Uh oh!

jerm-dro Feb 13, 2026

Uh oh!

jerm-dro Feb 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		// Dimension returns the dimensionality of the embeddings produced.
		Dimension() int

Conversation

aponcedeleonch commented Feb 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Large PR Justification

Uh oh!

github-actions bot left a comment

Choose a reason for hiding this comment

Large PR Detected

How to unblock this PR:

Alternative:

Uh oh!

codecov bot commented Feb 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

jerm-dro left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions bot commented Feb 13, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jerm-dro Feb 13, 2026

Choose a reason for hiding this comment

Uh oh!

jerm-dro Feb 13, 2026

Choose a reason for hiding this comment

Uh oh!

jerm-dro Feb 13, 2026

Choose a reason for hiding this comment

Uh oh!

jerm-dro Feb 13, 2026

Choose a reason for hiding this comment

Uh oh!

jerm-dro Feb 13, 2026

Choose a reason for hiding this comment

Uh oh!

jerm-dro Feb 13, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

aponcedeleonch commented Feb 12, 2026 •

edited

Loading

codecov bot commented Feb 12, 2026 •

edited

Loading