spire-ai 0.1.3

AI-native SDK for SpireDB — RAG, code search, agents
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

//! Ollama embedding provider.

use async_trait::async_trait;
use serde::{Deserialize, Serialize};

use crate::embedding::Embedder;
use crate::error::{Error, Result};

/// Ollama embedding provider (offline, local).
pub struct OllamaEmbedder {
    client: reqwest::Client,
    url: String,
    model: String,
    dims: usize,
}

impl OllamaEmbedder {
    /// Create with default model (`qwen3-embedding`, 768 dimensions).
    pub fn new(url: impl Into<String>) -> Self {
        Self::with_model(url, "qwen3-embedding")
    }

    /// Create with a specific model.
    pub fn with_model(url: impl Into<String>, model: impl Into<String>) -> Self {
        let model = model.into();
        let dims = known_dimensions(&model);
        Self {
            client: reqwest::Client::new(),
            url: url.into(),
            model,
            dims,
        }
    }
}

#[async_trait]
impl Embedder for OllamaEmbedder {
    async fn embed(&self, text: &str) -> Result<Vec<f32>> {
        let request = OllamaEmbedRequest {
            model: &self.model,
            input: vec![text],
        };

        let response: OllamaEmbedResponse = self
            .client
            .post(format!("{}/api/embed", self.url))
            .json(&request)
            .send()
            .await
            .map_err(|e| Error::Embedding(format!("ollama request failed: {e}")))?
            .json()
            .await
            .map_err(|e| Error::Embedding(format!("ollama response parse failed: {e}")))?;

        response
            .embeddings
            .into_iter()
            .next()
            .ok_or_else(|| Error::Embedding("ollama returned no embeddings".into()))
    }

    async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
        let input: Vec<&str> = texts.iter().map(|s| s.as_str()).collect();
        let request = OllamaEmbedRequest {
            model: &self.model,
            input,
        };

        let response: OllamaEmbedResponse = self
            .client
            .post(format!("{}/api/embed", self.url))
            .json(&request)
            .send()
            .await
            .map_err(|e| Error::Embedding(format!("ollama batch request failed: {e}")))?
            .json()
            .await
            .map_err(|e| Error::Embedding(format!("ollama batch response parse failed: {e}")))?;

        if response.embeddings.len() != texts.len() {
            return Err(Error::Embedding(format!(
                "expected {} embeddings, got {}",
                texts.len(),
                response.embeddings.len()
            )));
        }

        Ok(response.embeddings)
    }

    fn dimensions(&self) -> usize {
        self.dims
    }

    fn model_name(&self) -> &str {
        &self.model
    }
}

#[derive(Serialize)]
struct OllamaEmbedRequest<'a> {
    model: &'a str,
    input: Vec<&'a str>,
}

#[derive(Deserialize)]
struct OllamaEmbedResponse {
    embeddings: Vec<Vec<f32>>,
}

fn known_dimensions(model: &str) -> usize {
    match model {
        "qwen3-embedding" => 768,
        "mxbai-embed-large" => 1024,
        "all-minilm" => 384,
        "snowflake-arctic-embed" => 1024,
        _ => 768,
    }
}