ollama_vision/
tagger.rs

1use crate::parser::{self, ParseError};
2use crate::types::{OllamaVisionConfig, TagOptions};
3use reqwest::Client;
4use serde_json::json;
5use std::path::Path;
6
7const DEFAULT_TAG_PROMPT: &str = r#"You are an image tagging assistant. Analyze the provided image and return a JSON array of relevant tags. Each tag should be a single word or short phrase (2-3 words max) that describes a key visual element, style, subject, or mood in the image.
8
9Return ONLY a JSON array of strings. Example: ["portrait", "fantasy", "dark lighting", "woman", "medieval", "oil painting"]
10
11Return between 5 and 15 tags. Focus on:
12- Subject matter (person, animal, landscape, object)
13- Art style (photorealistic, anime, oil painting, digital art)
14- Mood/atmosphere (dark, bright, serene, dramatic)
15- Colors (warm tones, blue, monochrome)
16- Composition (close-up, wide shot, symmetrical)
17- Notable elements (fire, water, armor, flowers)"#;
18
19/// Tag an image using an Ollama vision model.
20///
21/// Returns a list of cleaned, lowercase tag strings extracted from the
22/// model's response using the 7-strategy parser.
23///
24/// # Errors
25///
26/// Returns an error if:
27/// - The image file cannot be read
28/// - The Ollama endpoint is unreachable
29/// - The model returns an error status
30/// - The response cannot be parsed into tags
31pub async fn tag_image(
32    client: &Client,
33    config: &OllamaVisionConfig,
34    image_path: &Path,
35    options: &TagOptions,
36) -> Result<Vec<String>, TagError> {
37    let image_b64 = read_image_base64(image_path)?;
38
39    let prompt = options.prompt.as_deref().unwrap_or(DEFAULT_TAG_PROMPT);
40
41    let mut body = json!({
42        "model": config.model,
43        "prompt": prompt,
44        "images": [image_b64],
45        "stream": false,
46        "options": config.options,
47    });
48
49    if options.request_json_format {
50        body["format"] = json!("json");
51    }
52
53    let url = format!("{}/api/generate", config.endpoint);
54    let resp = client
55        .post(&url)
56        .timeout(config.timeout)
57        .json(&body)
58        .send()
59        .await
60        .map_err(|e| TagError::Connection(config.endpoint.clone(), e.to_string()))?;
61
62    if !resp.status().is_success() {
63        let status = resp.status().as_u16();
64        let text = resp.text().await.unwrap_or_default();
65        return Err(TagError::OllamaError(status, text));
66    }
67
68    let json: serde_json::Value = resp
69        .json()
70        .await
71        .map_err(|e| TagError::InvalidResponse(e.to_string()))?;
72
73    let content = json
74        .get("response")
75        .and_then(|v| v.as_str())
76        .unwrap_or("[]");
77
78    let tags = parser::parse_tags(content).map_err(TagError::Parse)?;
79    Ok(truncate_tags(
80        tags,
81        options.max_tags,
82        options.max_tag_length,
83    ))
84}
85
86/// Tag an image from raw base64-encoded bytes (no file I/O).
87///
88/// Useful when you already have the image in memory.
89pub async fn tag_image_base64(
90    client: &Client,
91    config: &OllamaVisionConfig,
92    image_b64: &str,
93    options: &TagOptions,
94) -> Result<Vec<String>, TagError> {
95    let prompt = options.prompt.as_deref().unwrap_or(DEFAULT_TAG_PROMPT);
96
97    let mut body = json!({
98        "model": config.model,
99        "prompt": prompt,
100        "images": [image_b64],
101        "stream": false,
102        "options": config.options,
103    });
104
105    if options.request_json_format {
106        body["format"] = json!("json");
107    }
108
109    let url = format!("{}/api/generate", config.endpoint);
110    let resp = client
111        .post(&url)
112        .timeout(config.timeout)
113        .json(&body)
114        .send()
115        .await
116        .map_err(|e| TagError::Connection(config.endpoint.clone(), e.to_string()))?;
117
118    if !resp.status().is_success() {
119        let status = resp.status().as_u16();
120        let text = resp.text().await.unwrap_or_default();
121        return Err(TagError::OllamaError(status, text));
122    }
123
124    let json: serde_json::Value = resp
125        .json()
126        .await
127        .map_err(|e| TagError::InvalidResponse(e.to_string()))?;
128
129    let content = json
130        .get("response")
131        .and_then(|v| v.as_str())
132        .unwrap_or("[]");
133
134    let tags = parser::parse_tags(content).map_err(TagError::Parse)?;
135    Ok(truncate_tags(
136        tags,
137        options.max_tags,
138        options.max_tag_length,
139    ))
140}
141
142/// Apply tag count and length limits.
143fn truncate_tags(tags: Vec<String>, max_tags: usize, max_tag_length: usize) -> Vec<String> {
144    tags.into_iter()
145        .take(max_tags)
146        .map(|t| truncate_utf8(&t, max_tag_length))
147        .collect()
148}
149
150fn truncate_utf8(text: &str, max_chars: usize) -> String {
151    if text.chars().count() <= max_chars {
152        return text.to_string();
153    }
154
155    text.chars().take(max_chars).collect()
156}
157
158#[cfg(test)]
159mod tests {
160    use super::truncate_utf8;
161
162    #[test]
163    fn truncate_utf8_preserves_codepoint_boundaries() {
164        assert_eq!(truncate_utf8("naive", 3), "nai");
165        assert_eq!(truncate_utf8("cafe\u{301}", 4), "cafe");
166        assert_eq!(truncate_utf8("こんにちは", 3), "こんに");
167    }
168}
169
170/// Errors that can occur during image tagging.
171#[derive(Debug, thiserror::Error)]
172pub enum TagError {
173    #[error("Cannot connect to Ollama at {0}: {1}")]
174    Connection(String, String),
175
176    #[error("Ollama returned HTTP {0}: {1}")]
177    OllamaError(u16, String),
178
179    #[error("Invalid response from Ollama: {0}")]
180    InvalidResponse(String),
181
182    #[error("Failed to read image: {0}")]
183    ImageRead(String),
184
185    #[error("{0}")]
186    Parse(#[from] ParseError),
187}
188
189impl TagError {
190    /// Stable string discriminant for programmatic matching.
191    pub fn kind(&self) -> &'static str {
192        match self {
193            Self::Connection(..) => "connection",
194            Self::OllamaError(..) => "ollama_error",
195            Self::InvalidResponse(_) => "invalid_response",
196            Self::ImageRead(_) => "image_read",
197            Self::Parse(_) => "parse",
198        }
199    }
200}
201
202fn read_image_base64(path: &Path) -> Result<String, TagError> {
203    let bytes = std::fs::read(path)
204        .map_err(|e| TagError::ImageRead(format!("{}: {}", path.display(), e)))?;
205    Ok(base64::Engine::encode(
206        &base64::engine::general_purpose::STANDARD,
207        &bytes,
208    ))
209}
ollama_vision/tagger.rs

ollama_vision/
tagger.rs
