Reuse tool embeddings across sessions

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

@@ -4,13 +4,24 @@
 // Package similarity provides vector distance functions for semantic search.
 package similarity
 
-import "math"
+import (
+	"fmt"
+	"math"
+)
 
 // CosineSimilarity computes the cosine similarity between two vectors.
 // Returns a value in [-1, 1] where 1 means identical direction,
 // 0 means orthogonal, and -1 means opposite direction.
-// Both vectors must have the same length.
-func CosineSimilarity(a, b []float32) float64 {
+//
+// Vectors of different lengths return an error. The guard lives here rather
+// than at the call sites because the loop below indexes both slices
+// positionally: a shorter b would panic and a longer b would silently ignore
+// its tail, and a caller that forgets its own check gets one or the other.
+func CosineSimilarity(a, b []float32) (float64, error) {
+	if len(a) != len(b) {
+		return 0, fmt.Errorf("vectors have different dimensions: %d and %d", len(a), len(b))
+	}
+
 	var dot, normA, normB float64
 	for i := range a {
 		ai := float64(a[i])
@@ -22,14 +33,19 @@ func CosineSimilarity(a, b []float32) float64 {
 
 	denom := math.Sqrt(normA) * math.Sqrt(normB)
 	if denom == 0 {
-		return 0
+		return 0, nil
 	}
-	return dot / denom
+	return dot / denom, nil
 }
 
 // CosineDistance computes the cosine distance between two vectors.
 // Returns a value in [0, 2] where 0 means identical direction and 2 means
 // opposite direction. Lower values indicate more similar vectors.
-func CosineDistance(a, b []float32) float64 {
-	return 1 - CosineSimilarity(a, b)
+// Vectors of different lengths return an error (see CosineSimilarity).
+func CosineDistance(a, b []float32) (float64, error) {
+	sim, err := CosineSimilarity(a, b)
+	if err != nil {
+		return 0, err
+	}
+	return 1 - sim, nil
 }

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

@@ -23,7 +23,7 @@ func BenchmarkCosineDistance_384(b *testing.B) {
 	b.ResetTimer()
 	b.ReportAllocs()
 	for b.Loop() {
-		CosineDistance(a, v)
+		_, _ = CosineDistance(a, v)
 	}
 }
 
@@ -33,6 +33,6 @@ func BenchmarkCosineDistance_768(b *testing.B) {
 	b.ResetTimer()
 	b.ReportAllocs()
 	for b.Loop() {
-		CosineDistance(a, v)
+		_, _ = CosineDistance(a, v)
 	}
 }

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

@@ -28,7 +28,37 @@ func TestCosineSimilarity(t *testing.T) {
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
-			require.InDelta(t, tc.want, CosineSimilarity(tc.a, tc.b), 1e-7)
+			got, err := CosineSimilarity(tc.a, tc.b)
+			require.NoError(t, err)
+			require.InDelta(t, tc.want, got, 1e-7)
+		})
+	}
+}
+
+// TestCosineSimilarity_DimensionMismatch asserts mismatched widths are refused
+// rather than computed. Without the guard, a shorter b panics on indexing and
+// a longer b silently ignores its tail — a wrong answer, not an error.
+func TestCosineSimilarity_DimensionMismatch(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name string
+		a, b []float32
+	}{
+		{name: "b shorter would panic", a: []float32{1, 2, 3}, b: []float32{1, 2}},
+		{name: "b longer would be silently truncated", a: []float32{1, 2}, b: []float32{1, 2, 3}},
+		{name: "empty against non-empty", a: nil, b: []float32{1}},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			_, err := CosineSimilarity(tc.a, tc.b)
+			require.Error(t, err)
+
+			_, err = CosineDistance(tc.a, tc.b)
+			require.Error(t, err)
 		})
 	}
 }
@@ -49,7 +79,9 @@ func TestCosineDistance(t *testing.T) {
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
-			require.InDelta(t, tc.want, CosineDistance(tc.a, tc.b), 1e-7)
+			got, err := CosineDistance(tc.a, tc.b)
+			require.NoError(t, err)
+			require.InDelta(t, tc.want, got, 1e-7)
 		})
 	}
 }

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

@@ -178,6 +178,13 @@ func (c *openAIClient) embedChunk(ctx context.Context, texts []string) ([][]floa
 	return embeddings, nil
 }
 
+// ModelID returns the configured model name. The OpenAI /embeddings API
+// selects the model per request from this same field, so unlike TEI there is
+// no server-side state that could drift from it.
+func (c *openAIClient) ModelID(context.Context) (string, error) {
+	return c.model, nil
+}
+
 // Close is a no-op for the OpenAI client.
 func (*openAIClient) Close() error {
 	return nil

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

@@ -59,6 +59,18 @@ func Test_newOpenAIClient(t *testing.T) {
 	})
 }
 
+func TestOpenAIClient_ModelID(t *testing.T) {
+	t.Parallel()
+
+	client, err := newOpenAIClient("http://embeddings:8080/v1", "text-embedding-3-small", "key", nil, 0)
+	require.NoError(t, err)
+
+	id, err := client.ModelID(context.Background())
+	require.NoError(t, err)
+	require.Equal(t, "text-embedding-3-small", id,
+		"the OpenAI client's model is fixed by configuration and sent per request")
+}
+
 func TestOpenAIClient_Embed(t *testing.T) {
 	t.Parallel()
 

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

@@ -68,24 +68,41 @@ func newTEIClient(baseURL string, timeout time.Duration) (*teiClient, error) {
 
 // teiInfoResponse is a subset of the TEI /info endpoint response.
 type teiInfoResponse struct {
-	MaxClientBatchSize int `json:"max_client_batch_size"`
+	ModelID            string `json:"model_id"`
+	MaxClientBatchSize int    `json:"max_client_batch_size"`
 }
 
-// fetchMaxBatchSize queries the TEI /info endpoint and returns the max client batch size.
-func fetchMaxBatchSize(baseURL string, httpClient *http.Client) (int, error) {
-	resp, err := httpClient.Get(baseURL + infoPath) // #nosec G107 -- URL is built from the configured TEI base URL
+// fetchInfo queries the TEI /info endpoint.
+func fetchInfo(ctx context.Context, baseURL string, httpClient *http.Client) (teiInfoResponse, error) {
+	var info teiInfoResponse
+
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+infoPath, nil)
+	if err != nil {
+		return info, fmt.Errorf("failed to create TEI /info request: %w", err)
+	}
+
+	resp, err := httpClient.Do(req) // #nosec G704 -- URL is built from the configured TEI base URL
 	if err != nil {
-		return 0, fmt.Errorf("TEI /info request failed: %w", err)
+		return info, fmt.Errorf("TEI /info request failed: %w", err)
 	}
 	defer func() { _ = resp.Body.Close() }()
 
 	if resp.StatusCode != http.StatusOK {
-		return 0, fmt.Errorf("TEI /info returned status %d", resp.StatusCode)
+		return info, fmt.Errorf("TEI /info returned status %d", resp.StatusCode)
 	}
 
-	var info teiInfoResponse
 	if err := json.NewDecoder(resp.Body).Decode(&info); err != nil {
-		return 0, fmt.Errorf("failed to decode TEI /info response: %w", err)
+		return info, fmt.Errorf("failed to decode TEI /info response: %w", err)
+	}
+
+	return info, nil
+}
+
+// fetchMaxBatchSize queries the TEI /info endpoint and returns the max client batch size.
+func fetchMaxBatchSize(baseURL string, httpClient *http.Client) (int, error) {
+	info, err := fetchInfo(context.Background(), baseURL, httpClient)
+	if err != nil {
+		return 0, err
 	}
 
 	if info.MaxClientBatchSize <= 0 {
@@ -95,6 +112,22 @@ func fetchMaxBatchSize(baseURL string, httpClient *http.Client) (int, error) {
 	return info.MaxClientBatchSize, nil
 }
 
+// ModelID returns the id of the model the TEI server is currently running,
+// read from /info on every call. The model is a property of the running
+// container, not of this client's configuration, so it is deliberately not
+// cached: the point is letting a caller observe a redeploy that swapped the
+// model behind an unchanged URL.
+func (c *teiClient) ModelID(ctx context.Context) (string, error) {
+	info, err := fetchInfo(ctx, c.baseURL, c.httpClient)
+	if err != nil {
+		return "", err
+	}
+	if info.ModelID == "" {
+		return "", fmt.Errorf("TEI /info reported no model_id")
+	}
+	return info.ModelID, nil
+}
+
 // embedRequest is the JSON body sent to the TEI /embed endpoint.
 type embedRequest struct {
 	Inputs []string `json:"inputs"`

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

@@ -9,6 +9,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -69,6 +70,89 @@ func Test_newTEIClient(t *testing.T) {
 	})
 }
 
+func TestTEIClient_ModelID(t *testing.T) {
+	t.Parallel()
+
+	t.Run("returns the model id from /info", func(t *testing.T) {
+		t.Parallel()
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			require.Equal(t, infoPath, r.URL.Path)
+			w.Header().Set("Content-Type", "application/json")
+			_, _ = w.Write([]byte(`{"model_id": "BAAI/bge-small-en-v1.5", "max_client_batch_size": 16}`))
+		}))
+		defer srv.Close()
+
+		client, err := newTEIClient(srv.URL, 0)
+		require.NoError(t, err)
+
+		id, err := client.ModelID(context.Background())
+		require.NoError(t, err)
+		require.Equal(t, "BAAI/bge-small-en-v1.5", id)
+	})
+
+	t.Run("reads per call so a swap is observable", func(t *testing.T) {
+		t.Parallel()
+		var calls atomic.Int64
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			w.Header().Set("Content-Type", "application/json")
+			if calls.Add(1) > 2 { // the constructor's own /info read is call 1
+				_, _ = w.Write([]byte(`{"model_id": "model-b", "max_client_batch_size": 16}`))
+				return
+			}
+			_, _ = w.Write([]byte(`{"model_id": "model-a", "max_client_batch_size": 16}`))
+		}))
+		defer srv.Close()
+
+		client, err := newTEIClient(srv.URL, 0)
+		require.NoError(t, err)
+
+		first, err := client.ModelID(context.Background())
+		require.NoError(t, err)
+		second, err := client.ModelID(context.Background())
+		require.NoError(t, err)
+
+		require.Equal(t, "model-a", first)
+		require.Equal(t, "model-b", second,
+			"the id must be read live per call, not cached at construction")
+	})
+
+	t.Run("missing model_id is an error", func(t *testing.T) {
+		t.Parallel()
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			w.Header().Set("Content-Type", "application/json")
+			_, _ = w.Write([]byte(`{"max_client_batch_size": 16}`))
+		}))
+		defer srv.Close()
+
+		client, err := newTEIClient(srv.URL, 0)
+		require.NoError(t, err)
+
+		_, err = client.ModelID(context.Background())
+		require.ErrorContains(t, err, "no model_id")
+	})
+
+	t.Run("non-200 is an error", func(t *testing.T) {
+		t.Parallel()
+		var constructed atomic.Bool
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			if !constructed.Load() {
+				w.Header().Set("Content-Type", "application/json")
+				_, _ = w.Write([]byte(`{"model_id": "model-a", "max_client_batch_size": 16}`))
+				return
+			}
+			w.WriteHeader(http.StatusServiceUnavailable)
+		}))
+		defer srv.Close()
+
+		client, err := newTEIClient(srv.URL, 0)
+		require.NoError(t, err)
+		constructed.Store(true)
+
+		_, err = client.ModelID(context.Background())
+		require.ErrorContains(t, err, "status 503")
+	})
+}
+
 func TestTEIClient_Embed(t *testing.T) {
 	t.Parallel()
 

pkg/vmcp/optimizer/internal/toolstore/schema.sql

@@ -1,13 +1,27 @@
 -- SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
 -- SPDX-License-Identifier: Apache-2.0
 
--- Capabilities table stores tool/resource/prompt metadata
+-- Capabilities table stores tool/resource/prompt metadata.
+--
+-- content_hash identifies the exact input the stored embedding was produced
+-- from, covering both the embedded text and the embedding backend (see
+-- embeddingCacheKey), so UpsertTools can reuse a vector when the hash still
+-- matches. NULL means no embedding (FTS5-only mode).
+--
+-- The database is recreated in memory on every process start, so this column
+-- needs no migration. That changes if the store ever becomes file-backed.
 CREATE TABLE IF NOT EXISTS llm_capabilities (
     name TEXT PRIMARY KEY,
     description TEXT NOT NULL DEFAULT '',
-    embedding BLOB
+    embedding BLOB,
+    content_hash TEXT
 );
 
+-- The reuse lookup selects by content_hash across the whole table, not by the
+-- name primary key.
+CREATE INDEX IF NOT EXISTS llm_capabilities_content_hash_idx
+    ON llm_capabilities (content_hash);
+
 -- FTS5 virtual table for full-text search with BM25 ranking.
 -- tokenize='porter' uses the Porter stemming algorithm so that morphological
 -- variants of a word (e.g. "running", "runs", "ran") match the root form "run".