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