Skip to main content

sqlserver_mcp_catalog/services/
embedding_service.rs

1// SQL Server 2025 - master/msdb/sandbox combined catalog MCP server.
2//
3// Single source of truth for embedding computation. bin/populate_embeddings.rs
4// (indexing time) and tools/search_tool.rs (live query time) both call
5// `embed()` from here, so the two are structurally guaranteed to share the
6// same model and vector space.
7//
8// `all-MiniLM-L6-v2` at its native 384 dimensions -- switched from mcpify's
9// originally-generated `all-mpnet-base-v2` (768-dim) purely to fit
10// crates.io's 10 MiB package-size limit: with 4 SQL Server versions each
11// embedding ~700-800 operations, 768-dim float32 vectors alone pushed the
12// packaged crate to ~10.0 MiB compressed (measured via `cargo package`),
13// right at the limit. Embedding vectors are close to incompressible
14// (unlike the relational/schema data, which compresses ~20x), so halving
15// the dimension was the single most effective lever -- it brought the
16// package to a comfortable ~7 MiB. The tradeoff is somewhat coarser
17// `search` tool relevance than the larger model would give; acceptable
18// here since queries are short technical phrases ("rename a table",
19// "current session's SQL text"), not long natural-language passages.
20//
21// IMPORTANT: mcpify's own generator hard-codes the `semantic_endpoints`
22// vec0 table's column as `FLOAT[768]` (`mcpify/src/db/schema.rs`, not
23// something this project's source controls) -- every `mcp_store*.db` file
24// mcpify produces is born with a 768-dim column regardless of what this
25// module computes. Docs/sqlserver-eda-openapi-pipeline/scripts/
26// regenerate_mcp_server.sh's final step recreates that table at `FLOAT[384]`
27// in each store after every `mcpify sync`, before `populate_embeddings`
28// runs -- skipping that step means `populate_embeddings` will fail with a
29// sqlite-vec dimension-mismatch error (384-dim vectors don't fit a
30// 768-dim column), not silently write wrong data.
31
32use std::sync::{Mutex, OnceLock};
33
34use fastembed::{EmbeddingModel, TextEmbedding, TextInitOptions};
35
36fn model() -> &'static Mutex<TextEmbedding> {
37    static MODEL: OnceLock<Mutex<TextEmbedding>> = OnceLock::new();
38    MODEL.get_or_init(|| {
39        // Downloads on first use and caches locally afterward (no network
40        // needed once cached). `.expect()` matches this project's other
41        // unrecoverable-startup-error handling: nothing useful can happen
42        // if the model can't be fetched/loaded.
43        Mutex::new(
44            TextEmbedding::try_new(TextInitOptions::new(EmbeddingModel::AllMiniLML6V2))
45                .expect("failed to load the all-MiniLM-L6-v2 embedding model"),
46        )
47    })
48}
49
50/// Computes a 384-dim embedding vector for `text`, mean-pooled and
51/// normalized (fastembed's default behavior for this model, replicating
52/// the sentence-transformers reference implementation).
53pub fn embed(text: &str) -> anyhow::Result<Vec<f32>> {
54    let model = model();
55    let mut model = model.lock().unwrap();
56    let mut embeddings = model.embed(vec![text], None)?;
57    embeddings
58        .pop()
59        .ok_or_else(|| anyhow::anyhow!("embedding model returned no output for the given text"))
60}