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-mpnet-base-v2` at its native 768 dimensions, matching
8// `mcp_store.db`'s `semantic_endpoints` schema exactly (mcpify's shared
9// pipeline hard-codes `FLOAT[768]`) and the same model
10// `targets::typescript`'s `@xenova/transformers`-based embedding-service.ts
11// uses, so `-l rust` and `-l typescript` output are search-behavior
12// equivalent for the same spec.
13
14use std::sync::{Mutex, OnceLock};
15
16use fastembed::{EmbeddingModel, TextEmbedding, TextInitOptions};
17
18fn model() -> &'static Mutex<TextEmbedding> {
19 static MODEL: OnceLock<Mutex<TextEmbedding>> = OnceLock::new();
20 MODEL.get_or_init(|| {
21 // Downloads on first use and caches locally afterward (no network
22 // needed once cached) — mirrors `@xenova/transformers`' own
23 // caching UX. `.expect()` here matches the TS target's own
24 // failure mode: an unrecoverable startup error either way if the
25 // model can't be fetched/loaded.
26 Mutex::new(
27 TextEmbedding::try_new(TextInitOptions::new(EmbeddingModel::AllMpnetBaseV2))
28 .expect("failed to load the all-mpnet-base-v2 embedding model"),
29 )
30 })
31}
32
33/// Computes a 768-dim embedding vector for `text`, mean-pooled and
34/// normalized (fastembed's default behavior for this model, replicating
35/// the sentence-transformers reference implementation).
36pub fn embed(text: &str) -> anyhow::Result<Vec<f32>> {
37 let model = model();
38 let mut model = model.lock().unwrap();
39 let mut embeddings = model.embed(vec![text], None)?;
40 embeddings
41 .pop()
42 .ok_or_else(|| anyhow::anyhow!("embedding model returned no output for the given text"))
43}