synaptic_nomic/
lib.rs

1use async_trait::async_trait;
2use serde_json::json;
3use synaptic_core::{Embeddings, SynapticError};
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub enum NomicModel {
7    NomicEmbedTextV1_5,
8    NomicEmbedTextV1,
9    Custom(String),
10}
11
12impl NomicModel {
13    pub fn as_str(&self) -> &str {
14        match self {
15            NomicModel::NomicEmbedTextV1_5 => "nomic-embed-text-v1.5",
16            NomicModel::NomicEmbedTextV1 => "nomic-embed-text-v1",
17            NomicModel::Custom(s) => s.as_str(),
18        }
19    }
20}
21
22impl std::fmt::Display for NomicModel {
23    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
24        write!(f, "{}", self.as_str())
25    }
26}
27
28/// Task type for Nomic embeddings (affects how the model encodes text).
29#[derive(Debug, Clone, PartialEq, Eq)]
30pub enum NomicTaskType {
31    SearchDocument,
32    SearchQuery,
33    Classification,
34    Clustering,
35}
36
37impl NomicTaskType {
38    pub fn as_str(&self) -> &str {
39        match self {
40            NomicTaskType::SearchDocument => "search_document",
41            NomicTaskType::SearchQuery => "search_query",
42            NomicTaskType::Classification => "classification",
43            NomicTaskType::Clustering => "clustering",
44        }
45    }
46}
47
48#[derive(Debug, Clone)]
49pub struct NomicConfig {
50    pub api_key: String,
51    pub model: String,
52    pub base_url: String,
53}
54
55impl NomicConfig {
56    pub fn new(api_key: impl Into<String>) -> Self {
57        Self {
58            api_key: api_key.into(),
59            model: NomicModel::NomicEmbedTextV1_5.to_string(),
60            base_url: "https://api-atlas.nomic.ai/v1".to_string(),
61        }
62    }
63
64    pub fn with_model(mut self, model: NomicModel) -> Self {
65        self.model = model.to_string();
66        self
67    }
68}
69
70pub struct NomicEmbeddings {
71    config: NomicConfig,
72    client: reqwest::Client,
73}
74
75impl NomicEmbeddings {
76    pub fn new(config: NomicConfig) -> Self {
77        Self {
78            config,
79            client: reqwest::Client::new(),
80        }
81    }
82
83    async fn embed_with_task(
84        &self,
85        texts: &[&str],
86        task_type: NomicTaskType,
87    ) -> Result<Vec<Vec<f32>>, SynapticError> {
88        let body = json!({
89            "model": self.config.model,
90            "texts": texts,
91            "task_type": task_type.as_str(),
92        });
93        let resp = self
94            .client
95            .post(format!("{}/embedding/text", self.config.base_url))
96            .header("Authorization", format!("Bearer {}", self.config.api_key))
97            .header("Content-Type", "application/json")
98            .json(&body)
99            .send()
100            .await
101            .map_err(|e| SynapticError::Embedding(format!("Nomic request: {e}")))?;
102        let status = resp.status().as_u16();
103        let json: serde_json::Value = resp
104            .json()
105            .await
106            .map_err(|e| SynapticError::Embedding(format!("Nomic parse: {e}")))?;
107        if status != 200 {
108            return Err(SynapticError::Embedding(format!(
109                "Nomic API error ({}): {}",
110                status, json
111            )));
112        }
113        let embeddings = json
114            .get("embeddings")
115            .and_then(|e| e.as_array())
116            .ok_or_else(|| SynapticError::Embedding("missing 'embeddings' field".to_string()))?;
117        let result = embeddings
118            .iter()
119            .map(|row| {
120                row.as_array()
121                    .unwrap_or(&vec![])
122                    .iter()
123                    .map(|v| v.as_f64().unwrap_or(0.0) as f32)
124                    .collect()
125            })
126            .collect();
127        Ok(result)
128    }
129}
130
131#[async_trait]
132impl Embeddings for NomicEmbeddings {
133    async fn embed_documents(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, SynapticError> {
134        self.embed_with_task(texts, NomicTaskType::SearchDocument)
135            .await
136    }
137
138    async fn embed_query(&self, text: &str) -> Result<Vec<f32>, SynapticError> {
139        let mut results = self
140            .embed_with_task(&[text], NomicTaskType::SearchQuery)
141            .await?;
142        results
143            .pop()
144            .ok_or_else(|| SynapticError::Embedding("empty response".to_string()))
145    }
146}
synaptic_nomic/lib.rs

synaptic_nomic/
lib.rs
