Skip to main content

github_mcp/services/
embedding_service.rs

1// GitHub v3 REST API MCP server — generated by mcpify. Do not hand-edit.
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 — see the plan's embeddings decision.
7// `all-MiniLM-L6-v2` at its native 384 dimensions, matching
8// `mcp_store.db`'s `semantic_endpoints` schema exactly (`FLOAT[384]`).
9// Switched from `all-mpnet-base-v2` (768-dim) so the packaged crate's
10// embedded stores fit under crates.io's 10MiB upload limit even at max
11// zstd compression — see the "switch to all-MiniLM-L6-v2" fix.
12//
13// This intentionally diverges from `targets::typescript`'s
14// `@xenova/transformers`-based embedding-service.ts, which still uses
15// `all-mpnet-base-v2`/768 dims — `-l rust` and `-l typescript` output are
16// no longer search-behavior equivalent for the same spec as a result of
17// this crate-size-driven change.
18
19use std::sync::{Mutex, OnceLock};
20
21use fastembed::{EmbeddingModel, TextEmbedding, TextInitOptions};
22
23fn model() -> &'static Mutex<TextEmbedding> {
24    static MODEL: OnceLock<Mutex<TextEmbedding>> = OnceLock::new();
25    MODEL.get_or_init(|| {
26        // Downloads on first use and caches locally afterward (no network
27        // needed once cached) — mirrors `@xenova/transformers`' own
28        // caching UX. `.expect()` here matches the TS target's own
29        // failure mode: an unrecoverable startup error either way if the
30        // model can't be fetched/loaded.
31        Mutex::new(
32            TextEmbedding::try_new(TextInitOptions::new(EmbeddingModel::AllMiniLML6V2))
33                .expect("failed to load the all-MiniLM-L6-v2 embedding model"),
34        )
35    })
36}
37
38/// Computes a 384-dim embedding vector for `text`, mean-pooled and
39/// normalized (fastembed's default behavior for this model, replicating
40/// the sentence-transformers reference implementation).
41pub fn embed(text: &str) -> anyhow::Result<Vec<f32>> {
42    let model = model();
43    let mut model = model.lock().unwrap();
44    let mut embeddings = model.embed(vec![text], None)?;
45    embeddings
46        .pop()
47        .ok_or_else(|| anyhow::anyhow!("embedding model returned no output for the given text"))
48}