1use std::{
2 net::IpAddr,
3 path::Path,
4 time::{Duration, Instant},
5};
6
7use anyhow::{Context, Result, anyhow, bail};
8#[cfg(feature = "native-nomic")]
9use fastembed::{EmbeddingModel, TextEmbedding, TextInitOptions};
10use serde::{Deserialize, Serialize};
11use serde_json::json;
12
13use crate::{
14 cache,
15 models::{EmbeddingStats, FunctionRecord, ReportConfig},
16 normalize::content_hash,
17};
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum ProviderKind {
21 OpenAi,
22 Ollama,
23 Nomic,
24 Lexical,
25 None,
26}
27
28impl ProviderKind {
29 pub fn as_str(self) -> &'static str {
30 match self {
31 Self::OpenAi => "openai",
32 Self::Ollama => "ollama",
33 Self::Nomic => "nomic",
34 Self::Lexical => "lexical",
35 Self::None => "none",
36 }
37 }
38}
39
40pub fn embeddings_for(
41 functions: &[FunctionRecord],
42 config: &ReportConfig,
43 cache_root: &Path,
44) -> Result<(Vec<Option<Vec<f32>>>, EmbeddingStats)> {
45 let started = Instant::now();
46 match config.provider {
47 ProviderKind::None => Ok((vec![None; functions.len()], elapsed_stats(started))),
48 ProviderKind::Lexical => {
49 let vectors: Vec<_> = functions
50 .iter()
51 .map(|function| Some(lexical_embedding(&function.normalized)))
52 .collect();
53 let mut stats = elapsed_stats(started);
54 stats.cache_misses = functions.len();
55 stats.dimensions = vectors
56 .iter()
57 .find_map(|vector| vector.as_ref().map(Vec::len));
58 Ok((vectors, stats))
59 }
60 ProviderKind::OpenAi => {
61 if !config.allow_source_upload {
62 bail!(
63 "openai provider would send source-derived text; rerun with --allow-source-upload to opt in"
64 );
65 }
66 let model = config.model.as_deref().unwrap_or("text-embedding-3-small");
67 let provider =
68 OpenAiProvider::new(model, Duration::from_secs(config.ollama_timeout_secs))?;
69 let mut out = Vec::with_capacity(functions.len());
70 let mut stats = elapsed_stats(started);
71 for function in functions {
72 let text = embedding_text(function);
73 let key = content_hash(&format!("openai:{model}:text={}", content_hash(&text)));
74 if let Some(vector) = cache::load_embedding(cache_root, &key)? {
75 stats.cache_hits += 1;
76 stats.dimensions.get_or_insert(vector.len());
77 out.push(Some(vector));
78 continue;
79 }
80 let vector = provider.embed(&text)?;
81 stats.cache_misses += 1;
82 stats.dimensions.get_or_insert(vector.len());
83 cache::save_embedding(cache_root, &key, &vector)?;
84 out.push(Some(vector));
85 }
86 stats.elapsed_ms = elapsed_ms(started);
87 Ok((out, stats))
88 }
89 ProviderKind::Ollama => {
90 let model = config
91 .model
92 .as_deref()
93 .context("--model is required when using --provider ollama")?;
94 let provider = OllamaProvider::new(config, model)?;
95 provider.embed_functions(functions, cache_root, started)
96 }
97 ProviderKind::Nomic => {
98 #[cfg(not(feature = "native-nomic"))]
99 {
100 bail!(
101 "native Nomic support is not included in this build; install with `cargo install funcvec --features native-nomic` or use `--provider lexical`"
102 );
103 }
104 #[cfg(feature = "native-nomic")]
105 {
106 let provider = NativeNomicProvider::new(config)?;
107 provider.embed_functions(functions, cache_root, started)
108 }
109 }
110 }
111}
112
113fn embedding_text(function: &FunctionRecord) -> String {
114 format!(
115 "name: {}\nlines: {}-{}\ncode:\n{}",
116 function.name, function.start_line, function.end_line, function.normalized
117 )
118}
119
120fn lexical_embedding(text: &str) -> Vec<f32> {
121 const DIMS: usize = 96;
122 let mut vector = vec![0.0; DIMS];
123 for token in text.split_whitespace() {
124 let hash = content_hash(token);
125 let bucket = usize::from_str_radix(&hash[..8], 16).unwrap_or(0) % DIMS;
126 vector[bucket] += 1.0;
127 }
128 normalize(&mut vector);
129 vector
130}
131
132fn normalize(vector: &mut [f32]) {
133 let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
134 if norm > 0.0 {
135 for value in vector {
136 *value /= norm;
137 }
138 }
139}
140
141fn elapsed_stats(started: Instant) -> EmbeddingStats {
142 EmbeddingStats {
143 elapsed_ms: elapsed_ms(started),
144 ..EmbeddingStats::default()
145 }
146}
147
148fn elapsed_ms(started: Instant) -> u64 {
149 started.elapsed().as_millis().try_into().unwrap_or(u64::MAX)
150}
151
152#[cfg(feature = "native-nomic")]
153#[derive(Debug, Clone)]
154struct NomicModel {
155 alias: &'static str,
156 model: EmbeddingModel,
157}
158
159#[cfg(feature = "native-nomic")]
160struct NativeNomicProvider {
161 model: NomicModel,
162 model_cache_dir: std::path::PathBuf,
163 native_threads: Option<usize>,
164}
165
166#[cfg(feature = "native-nomic")]
167impl NativeNomicProvider {
168 fn new(config: &ReportConfig) -> Result<Self> {
169 let model = parse_nomic_model(config.model.as_deref())?;
170 let model_cache_dir = model_cache_dir(config)?;
171 Ok(Self {
172 model,
173 model_cache_dir,
174 native_threads: config.native_threads,
175 })
176 }
177
178 fn embed_functions(
179 &self,
180 functions: &[FunctionRecord],
181 cache_root: &Path,
182 started: Instant,
183 ) -> Result<(Vec<Option<Vec<f32>>>, EmbeddingStats)> {
184 let mut stats = elapsed_stats(started);
185 let mut out = vec![None; functions.len()];
186 let mut pending_indices = Vec::new();
187 let mut pending_inputs = Vec::new();
188 let mut pending_keys = Vec::new();
189
190 for (idx, function) in functions.iter().enumerate() {
191 let text = nomic_embedding_text(function);
192 let key = self.cache_key(function, &text);
193 if let Some(vector) = cache::load_embedding(cache_root, &key)? {
194 stats.cache_hits += 1;
195 stats.dimensions.get_or_insert(vector.len());
196 out[idx] = Some(vector);
197 } else {
198 pending_indices.push(idx);
199 pending_inputs.push(text);
200 pending_keys.push(key);
201 }
202 }
203
204 if !pending_inputs.is_empty() {
205 let mut options = TextInitOptions::new(self.model.model.clone())
206 .with_cache_dir(self.model_cache_dir.clone())
207 .with_show_download_progress(false);
208 if let Some(threads) = self.native_threads {
209 options = options.with_intra_threads(threads);
210 }
211
212 let mut model = TextEmbedding::try_new(options).with_context(|| {
213 format!(
214 "failed to initialize native Nomic model `{}`",
215 self.model.alias
216 )
217 })?;
218 let vectors = model.embed(&pending_inputs, None).with_context(|| {
219 format!(
220 "failed to embed {} functions with native Nomic model `{}`",
221 pending_inputs.len(),
222 self.model.alias
223 )
224 })?;
225 if vectors.len() != pending_inputs.len() {
226 bail!(
227 "native Nomic returned {} embeddings for {} inputs",
228 vectors.len(),
229 pending_inputs.len()
230 );
231 }
232
233 for (idx, (key, vector)) in pending_indices
234 .into_iter()
235 .zip(pending_keys.into_iter().zip(vectors))
236 {
237 let dimension = vector.len();
238 if let Some(expected) = stats.dimensions {
239 if expected != dimension {
240 bail!(
241 "native Nomic returned inconsistent embedding dimensions: expected {expected}, got {dimension}"
242 );
243 }
244 } else {
245 stats.dimensions = Some(dimension);
246 }
247 cache::save_embedding(cache_root, &key, &vector)?;
248 stats.cache_misses += 1;
249 out[idx] = Some(vector);
250 }
251 }
252
253 stats.elapsed_ms = elapsed_ms(started);
254 Ok((out, stats))
255 }
256
257 fn cache_key(&self, function: &FunctionRecord, embedding_text: &str) -> String {
258 content_hash(&nomic_cache_key_seed(
259 self.model.alias,
260 function,
261 embedding_text,
262 ))
263 }
264}
265
266#[cfg(feature = "native-nomic")]
267fn nomic_cache_key_seed(
268 model_alias: &str,
269 function: &FunctionRecord,
270 embedding_text: &str,
271) -> String {
272 format!(
273 "nomic-fastembed-v1:model={model_alias}:prefix=clustering:function={}:text={}",
274 function.content_hash,
275 content_hash(embedding_text)
276 )
277}
278
279#[cfg(feature = "native-nomic")]
280fn nomic_embedding_text(function: &FunctionRecord) -> String {
281 format!("clustering: {}", embedding_text(function))
282}
283
284#[cfg(feature = "native-nomic")]
285fn parse_nomic_model(model: Option<&str>) -> Result<NomicModel> {
286 match model.unwrap_or(default_nomic_model()) {
287 "nomic-v1" | "nomic-embed-text-v1" => Ok(NomicModel {
288 alias: "nomic-v1",
289 model: EmbeddingModel::NomicEmbedTextV1,
290 }),
291 "nomic-v1.5" | "nomic-embed-text-v1.5" | "nomic-embed-text" => Ok(NomicModel {
292 alias: "nomic-v1.5",
293 model: EmbeddingModel::NomicEmbedTextV15,
294 }),
295 value => bail!(
296 "unsupported native Nomic model `{value}`; supported models: nomic-v1, nomic-v1.5"
297 ),
298 }
299}
300
301pub fn default_nomic_model() -> &'static str {
302 "nomic-v1.5"
303}
304
305#[cfg(feature = "native-nomic")]
306fn model_cache_dir(config: &ReportConfig) -> Result<std::path::PathBuf> {
307 if let Some(path) = &config.model_cache_dir {
308 return Ok(path.clone());
309 }
310 if let Some(path) = std::env::var_os("FUNCVEC_MODEL_CACHE_DIR")
311 .or_else(|| std::env::var_os("RFV_MODEL_CACHE_DIR"))
312 {
313 return Ok(std::path::PathBuf::from(path));
314 }
315 let cache_dir = dirs::cache_dir()
316 .context("could not determine OS cache directory; pass --model-cache-dir")?;
317 Ok(cache_dir.join("funcvec").join("models"))
318}
319
320struct OllamaProvider {
321 client: reqwest::blocking::Client,
322 host: String,
323 model: String,
324 keep_alive: Option<String>,
325 dimensions: Option<usize>,
326 truncate: bool,
327}
328
329impl OllamaProvider {
330 fn new(config: &ReportConfig, model: &str) -> Result<Self> {
331 let host = normalize_ollama_host(&config.ollama_host)?;
332 if !config.allow_nonlocal_ollama_host && !is_loopback_url(&host)? {
333 bail!(
334 "refusing to send source-derived text to non-loopback Ollama host `{host}`; rerun with --allow-nonlocal-ollama-host to opt in"
335 );
336 }
337
338 let client = reqwest::blocking::Client::builder()
339 .timeout(Duration::from_secs(config.ollama_timeout_secs))
340 .build()?;
341 Ok(Self {
342 client,
343 host,
344 model: model.to_owned(),
345 keep_alive: config.ollama_keep_alive.clone(),
346 dimensions: config.ollama_dimensions,
347 truncate: config.ollama_truncate,
348 })
349 }
350
351 fn embed_functions(
352 &self,
353 functions: &[FunctionRecord],
354 cache_root: &Path,
355 started: Instant,
356 ) -> Result<(Vec<Option<Vec<f32>>>, EmbeddingStats)> {
357 let model_digest = self.model_digest().unwrap_or(None);
358 let mut stats = elapsed_stats(started);
359 stats.model_digest = model_digest.clone();
360
361 let mut out = vec![None; functions.len()];
362 let mut pending_indices = Vec::new();
363 let mut pending_inputs = Vec::new();
364 let mut pending_keys = Vec::new();
365
366 for (idx, function) in functions.iter().enumerate() {
367 let text = embedding_text(function);
368 let key = self.cache_key(function, &text, model_digest.as_deref());
369 if let Some(vector) = cache::load_embedding(cache_root, &key)? {
370 stats.cache_hits += 1;
371 stats.dimensions.get_or_insert(vector.len());
372 out[idx] = Some(vector);
373 } else {
374 pending_indices.push(idx);
375 pending_inputs.push(text);
376 pending_keys.push(key);
377 }
378 }
379
380 if !pending_inputs.is_empty() {
381 let vectors = self.embed_batch(&pending_inputs)?;
382 if vectors.len() != pending_inputs.len() {
383 bail!(
384 "ollama returned {} embeddings for {} inputs",
385 vectors.len(),
386 pending_inputs.len()
387 );
388 }
389 for (idx, (key, vector)) in pending_indices
390 .into_iter()
391 .zip(pending_keys.into_iter().zip(vectors))
392 {
393 let dimension = vector.len();
394 if let Some(expected) = stats.dimensions {
395 if expected != dimension {
396 bail!(
397 "ollama returned inconsistent embedding dimensions: expected {expected}, got {dimension}"
398 );
399 }
400 } else {
401 stats.dimensions = Some(dimension);
402 }
403 cache::save_embedding(cache_root, &key, &vector)?;
404 stats.cache_misses += 1;
405 out[idx] = Some(vector);
406 }
407 }
408
409 stats.elapsed_ms = elapsed_ms(started);
410 Ok((out, stats))
411 }
412
413 fn cache_key(
414 &self,
415 function: &FunctionRecord,
416 embedding_text: &str,
417 model_digest: Option<&str>,
418 ) -> String {
419 content_hash(&format!(
420 "ollama:api_embed_v1:host={}:model={}:digest={}:truncate={}:dimensions={:?}:function={}:text={}",
421 self.host,
422 self.model,
423 model_digest.unwrap_or("unknown"),
424 self.truncate,
425 self.dimensions,
426 function.content_hash,
427 content_hash(embedding_text)
428 ))
429 }
430
431 fn embed_batch(&self, inputs: &[String]) -> Result<Vec<Vec<f32>>> {
432 let request = OllamaEmbedRequest {
433 model: &self.model,
434 input: inputs,
435 truncate: self.truncate,
436 keep_alive: self.keep_alive.as_deref(),
437 dimensions: self.dimensions,
438 };
439 let url = format!("{}/api/embed", self.host.trim_end_matches('/'));
440 let response = self
441 .client
442 .post(url)
443 .json(&request)
444 .send()
445 .map_err(|err| ollama_transport_error(err, &self.host))?;
446 let status = response.status();
447 if !status.is_success() {
448 let body = response.text().unwrap_or_default();
449 if status.as_u16() == 404 {
450 bail!(
451 "ollama model `{}` is not available at {}; run `ollama pull {}` ({})",
452 self.model,
453 self.host,
454 self.model,
455 body.trim()
456 );
457 }
458 bail!(
459 "ollama embed request failed with HTTP {status} at {}: {}",
460 self.host,
461 body.trim()
462 );
463 }
464 let response: OllamaEmbedResponse = response.json()?;
465 Ok(response.embeddings)
466 }
467
468 fn model_digest(&self) -> Result<Option<String>> {
469 let url = format!("{}/api/tags", self.host.trim_end_matches('/'));
470 let response = self
471 .client
472 .get(url)
473 .send()
474 .map_err(|err| ollama_transport_error(err, &self.host))?;
475 if !response.status().is_success() {
476 return Ok(None);
477 }
478 let tags: OllamaTagsResponse = response.json()?;
479 Ok(tags
480 .models
481 .into_iter()
482 .find(|model| model.name == self.model || model.model == self.model)
483 .and_then(|model| model.digest))
484 }
485}
486
487fn normalize_ollama_host(host: &str) -> Result<String> {
488 let trimmed = host.trim();
489 if trimmed.is_empty() {
490 bail!("--ollama-host cannot be empty");
491 }
492 let host = if trimmed.contains("://") {
493 trimmed.to_owned()
494 } else {
495 format!("http://{trimmed}")
496 };
497 let parsed = url::Url::parse(&host).with_context(|| format!("invalid Ollama host `{host}`"))?;
498 if parsed.scheme() != "http" && parsed.scheme() != "https" {
499 bail!("Ollama host must use http or https: {host}");
500 }
501 Ok(host.trim_end_matches('/').to_owned())
502}
503
504fn is_loopback_url(host: &str) -> Result<bool> {
505 let parsed = url::Url::parse(host)?;
506 let Some(host) = parsed.host_str() else {
507 return Ok(false);
508 };
509 if host.eq_ignore_ascii_case("localhost") {
510 return Ok(true);
511 }
512 Ok(host.parse::<IpAddr>().is_ok_and(|addr| addr.is_loopback()))
513}
514
515fn ollama_transport_error(err: reqwest::Error, host: &str) -> anyhow::Error {
516 if err.is_connect() {
517 anyhow!("could not connect to Ollama at {host}; start it with `ollama serve`")
518 } else if err.is_timeout() {
519 anyhow!("timed out waiting for Ollama at {host}")
520 } else {
521 anyhow!("ollama request to {host} failed: {err}")
522 }
523}
524
525#[derive(Debug, Serialize)]
526struct OllamaEmbedRequest<'a> {
527 model: &'a str,
528 input: &'a [String],
529 truncate: bool,
530 #[serde(skip_serializing_if = "Option::is_none")]
531 keep_alive: Option<&'a str>,
532 #[serde(skip_serializing_if = "Option::is_none")]
533 dimensions: Option<usize>,
534}
535
536#[derive(Debug, Deserialize)]
537struct OllamaEmbedResponse {
538 embeddings: Vec<Vec<f32>>,
539}
540
541#[derive(Debug, Deserialize)]
542struct OllamaTagsResponse {
543 models: Vec<OllamaTagModel>,
544}
545
546#[derive(Debug, Deserialize)]
547struct OllamaTagModel {
548 name: String,
549 model: String,
550 digest: Option<String>,
551}
552
553struct OpenAiProvider {
554 client: reqwest::blocking::Client,
555 api_key: String,
556 model: String,
557}
558
559impl OpenAiProvider {
560 fn new(model: &str, timeout: Duration) -> Result<Self> {
561 let api_key = std::env::var("OPENAI_API_KEY")
562 .context("OPENAI_API_KEY is required when using --provider openai")?;
563 Ok(Self {
564 client: reqwest::blocking::Client::builder()
565 .timeout(timeout)
566 .build()?,
567 api_key,
568 model: model.to_owned(),
569 })
570 }
571
572 fn embed(&self, input: &str) -> Result<Vec<f32>> {
573 let response = self
574 .client
575 .post("https://api.openai.com/v1/embeddings")
576 .bearer_auth(&self.api_key)
577 .json(&json!({
578 "model": self.model,
579 "input": input,
580 }))
581 .send()?
582 .error_for_status()?
583 .json::<OpenAiEmbeddingResponse>()?;
584 response
585 .data
586 .into_iter()
587 .next()
588 .map(|item| item.embedding)
589 .context("OpenAI embedding response did not contain an embedding")
590 }
591}
592
593#[derive(Debug, Deserialize)]
594struct OpenAiEmbeddingResponse {
595 data: Vec<OpenAiEmbeddingItem>,
596}
597
598#[derive(Debug, Deserialize)]
599struct OpenAiEmbeddingItem {
600 embedding: Vec<f32>,
601}
602
603#[cfg(test)]
604mod tests {
605 use super::*;
606
607 #[test]
608 fn lexical_embeddings_are_deterministic() {
609 assert_eq!(lexical_embedding("a b c"), lexical_embedding("a b c"));
610 }
611
612 #[test]
613 fn rejects_non_loopback_ollama_hosts_by_default() {
614 let mut config = ReportConfig {
615 provider: ProviderKind::Ollama,
616 model: Some("nomic-embed-text".to_owned()),
617 ollama_host: "http://example.com:11434".to_owned(),
618 ..ReportConfig::default()
619 };
620 assert!(OllamaProvider::new(&config, "nomic-embed-text").is_err());
621
622 config.allow_nonlocal_ollama_host = true;
623 assert!(OllamaProvider::new(&config, "nomic-embed-text").is_ok());
624 }
625
626 #[test]
627 fn accepts_loopback_ollama_hosts() {
628 let config = ReportConfig {
629 provider: ProviderKind::Ollama,
630 model: Some("nomic-embed-text".to_owned()),
631 ollama_host: "127.0.0.1:11434".to_owned(),
632 ..ReportConfig::default()
633 };
634 assert!(OllamaProvider::new(&config, "nomic-embed-text").is_ok());
635 }
636
637 #[cfg(feature = "native-nomic")]
638 #[test]
639 fn parses_native_nomic_model_aliases() {
640 assert_eq!(parse_nomic_model(None).unwrap().alias, "nomic-v1.5");
641 assert_eq!(
642 parse_nomic_model(Some("nomic-embed-text-v1"))
643 .unwrap()
644 .alias,
645 "nomic-v1"
646 );
647 assert_eq!(
648 parse_nomic_model(Some("nomic-embed-text")).unwrap().alias,
649 "nomic-v1.5"
650 );
651 assert!(parse_nomic_model(Some("nomic-v1.5-q")).is_err());
652 }
653
654 #[cfg(feature = "native-nomic")]
655 #[test]
656 fn native_nomic_embedding_text_uses_clustering_prefix() {
657 let function = sample_function();
658 let text = nomic_embedding_text(&function);
659 assert!(text.starts_with("clustering: name: sample"));
660 }
661
662 #[cfg(feature = "native-nomic")]
663 #[test]
664 fn native_nomic_cache_seed_versions_embedding_behavior() {
665 let function = sample_function();
666 let text = nomic_embedding_text(&function);
667 let seed = nomic_cache_key_seed("nomic-v1.5", &function, &text);
668 assert!(seed.contains("nomic-fastembed-v1"));
669 assert!(seed.contains("model=nomic-v1.5"));
670 assert!(seed.contains("prefix=clustering"));
671 assert!(seed.contains("function=abc123"));
672 }
673
674 #[cfg(feature = "native-nomic")]
675 #[test]
676 fn native_nomic_smoke_test_is_explicitly_opted_in() {
677 if std::env::var_os("FUNCVEC_RUN_NATIVE_MODEL_TESTS")
678 .or_else(|| std::env::var_os("RFV_RUN_NATIVE_MODEL_TESTS"))
679 .is_none()
680 {
681 return;
682 }
683
684 let config = ReportConfig {
685 provider: ProviderKind::Nomic,
686 model: Some(default_nomic_model().to_owned()),
687 native_threads: Some(1),
688 ..ReportConfig::default()
689 };
690 let provider = NativeNomicProvider::new(&config).unwrap();
691 let function = sample_function();
692 let cache_root = tempfile::tempdir().unwrap();
693 let (embeddings, stats) = provider
694 .embed_functions(&[function], cache_root.path(), Instant::now())
695 .unwrap();
696 assert_eq!(embeddings.len(), 1);
697 assert_eq!(stats.dimensions, Some(768));
698 }
699
700 #[cfg(feature = "native-nomic")]
701 fn sample_function() -> FunctionRecord {
702 FunctionRecord {
703 id: "id".to_owned(),
704 name: "sample".to_owned(),
705 file: "src/lib.rs".into(),
706 start_line: 1,
707 end_line: 3,
708 source: "fn sample() -> i32 { 1 }".to_owned(),
709 normalized: "fn ID ( ) -> ID { NUM }".to_owned(),
710 token_count: 8,
711 line_count: 3,
712 content_hash: "abc123".to_owned(),
713 expected_group: None,
714 }
715 }
716}