ollama-rs 0.3.5

A Rust library for interacting with the Ollama API
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
120
121
122
123
124
125
126
127
128

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

use crate::{generation::parameters::KeepAlive, models::ModelOptions};

#[derive(Debug, Deserialize, PartialEq)]
#[serde(untagged)]
pub enum EmbeddingsInput {
    Single(String),
    Multiple(Vec<String>),
}

impl Default for EmbeddingsInput {
    fn default() -> Self {
        Self::Single(String::default())
    }
}

impl From<String> for EmbeddingsInput {
    fn from(s: String) -> Self {
        Self::Single(s)
    }
}

impl From<&str> for EmbeddingsInput {
    fn from(s: &str) -> Self {
        Self::Single(s.to_string())
    }
}

impl From<Vec<String>> for EmbeddingsInput {
    fn from(v: Vec<String>) -> Self {
        Self::Multiple(v)
    }
}

impl From<Vec<&str>> for EmbeddingsInput {
    fn from(v: Vec<&str>) -> Self {
        Self::Multiple(v.iter().map(|s| s.to_string()).collect())
    }
}

impl Serialize for EmbeddingsInput {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        match self {
            EmbeddingsInput::Single(s) => s.serialize(serializer),
            EmbeddingsInput::Multiple(v) => v.serialize(serializer),
        }
    }
}

/// An embeddings generation request to Ollama.
#[derive(Debug, Serialize, Deserialize, Default)]
pub struct GenerateEmbeddingsRequest {
    #[serde(rename = "model")]
    pub model_name: String,
    pub input: EmbeddingsInput,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub truncate: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub options: Option<ModelOptions>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub keep_alive: Option<KeepAlive>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub dimensions: Option<u32>,
}

impl GenerateEmbeddingsRequest {
    pub fn new(model_name: String, input: EmbeddingsInput) -> Self {
        Self {
            model_name,
            input,
            ..Default::default()
        }
    }

    pub fn options(mut self, options: ModelOptions) -> Self {
        self.options = Some(options);
        self
    }

    pub fn keep_alive(mut self, keep_alive: KeepAlive) -> Self {
        self.keep_alive = Some(keep_alive);
        self
    }

    pub fn truncate(mut self, truncate: bool) -> Self {
        self.truncate = Some(truncate);
        self
    }

    pub fn dimensions(mut self, dimensions: u32) -> Self {
        self.dimensions = Some(dimensions);
        self
    }
}

#[cfg(test)]
mod tests {
    use crate::generation::embeddings::request::{EmbeddingsInput, GenerateEmbeddingsRequest};

    #[test]
    fn serde_embedding_request_single() {
        let request = GenerateEmbeddingsRequest::new(
            "test".to_string(),
            EmbeddingsInput::Single("test".to_string()),
        );
        let json = serde_json::to_vec(&request).unwrap();
        let parsed_request: GenerateEmbeddingsRequest = serde_json::from_slice(&json).unwrap();
        assert_eq!(request.model_name, parsed_request.model_name);
        assert_eq!(request.input, parsed_request.input);
        assert_eq!(request.truncate, parsed_request.truncate);
        assert_eq!(request.keep_alive, parsed_request.keep_alive);
    }

    #[test]
    fn serde_embedding_request_multiple() {
        let request = GenerateEmbeddingsRequest::new(
            "test".to_string(),
            EmbeddingsInput::Multiple(vec!["test".to_string()]),
        );
        let json = serde_json::to_vec(&request).unwrap();
        let parsed_request: GenerateEmbeddingsRequest = serde_json::from_slice(&json).unwrap();
        assert_eq!(request.model_name, parsed_request.model_name);
        assert_eq!(request.input, parsed_request.input);
        assert_eq!(request.truncate, parsed_request.truncate);
        assert_eq!(request.keep_alive, parsed_request.keep_alive);
    }
}