ollama_vision/
captioner.rs

1use crate::parser;
2use crate::types::{CaptionOptions, OllamaVisionConfig};
3use reqwest::Client;
4use serde_json::json;
5use std::path::Path;
6
7const DEFAULT_CAPTION_PROMPT: &str = r#"Describe this image in 1-2 sentences. Focus on the main subject, art style, composition, lighting, and mood. Be specific and concise. Do not start with "This image shows" or "The image depicts". Just describe what you see directly."#;
8
9/// Generate a caption for an image using an Ollama vision model.
10///
11/// Returns a cleaned caption string with `<think>` blocks stripped.
12///
13/// # Errors
14///
15/// Returns an error if:
16/// - The image file cannot be read
17/// - The Ollama endpoint is unreachable
18/// - The model returns an error or empty caption
19pub async fn caption_image(
20    client: &Client,
21    config: &OllamaVisionConfig,
22    image_path: &Path,
23    options: &CaptionOptions,
24) -> Result<String, CaptionError> {
25    let image_b64 = read_image_base64(image_path)?;
26
27    let prompt = options.prompt.as_deref().unwrap_or(DEFAULT_CAPTION_PROMPT);
28
29    let body = json!({
30        "model": config.model,
31        "prompt": prompt,
32        "images": [image_b64],
33        "stream": false,
34        "options": config.options,
35    });
36
37    let url = format!("{}/api/generate", config.endpoint);
38    let resp = client
39        .post(&url)
40        .timeout(config.timeout)
41        .json(&body)
42        .send()
43        .await
44        .map_err(|e| CaptionError::Connection(config.endpoint.clone(), e.to_string()))?;
45
46    if !resp.status().is_success() {
47        let status = resp.status().as_u16();
48        let text = resp.text().await.unwrap_or_default();
49        return Err(CaptionError::OllamaError(status, text));
50    }
51
52    let json: serde_json::Value = resp
53        .json()
54        .await
55        .map_err(|e| CaptionError::InvalidResponse(e.to_string()))?;
56
57    let raw = json.get("response").and_then(|v| v.as_str()).unwrap_or("");
58
59    let caption = parser::strip_think_tags(raw).trim().to_string();
60
61    if caption.is_empty() {
62        return Err(CaptionError::EmptyCaption);
63    }
64
65    Ok(truncate_caption(&caption, options.max_caption_length))
66}
67
68/// Caption an image from raw base64-encoded bytes (no file I/O).
69pub async fn caption_image_base64(
70    client: &Client,
71    config: &OllamaVisionConfig,
72    image_b64: &str,
73    options: &CaptionOptions,
74) -> Result<String, CaptionError> {
75    let prompt = options.prompt.as_deref().unwrap_or(DEFAULT_CAPTION_PROMPT);
76
77    let body = json!({
78        "model": config.model,
79        "prompt": prompt,
80        "images": [image_b64],
81        "stream": false,
82        "options": config.options,
83    });
84
85    let url = format!("{}/api/generate", config.endpoint);
86    let resp = client
87        .post(&url)
88        .timeout(config.timeout)
89        .json(&body)
90        .send()
91        .await
92        .map_err(|e| CaptionError::Connection(config.endpoint.clone(), e.to_string()))?;
93
94    if !resp.status().is_success() {
95        let status = resp.status().as_u16();
96        let text = resp.text().await.unwrap_or_default();
97        return Err(CaptionError::OllamaError(status, text));
98    }
99
100    let json: serde_json::Value = resp
101        .json()
102        .await
103        .map_err(|e| CaptionError::InvalidResponse(e.to_string()))?;
104
105    let raw = json.get("response").and_then(|v| v.as_str()).unwrap_or("");
106
107    let caption = parser::strip_think_tags(raw).trim().to_string();
108
109    if caption.is_empty() {
110        return Err(CaptionError::EmptyCaption);
111    }
112
113    Ok(truncate_caption(&caption, options.max_caption_length))
114}
115
116/// Truncate caption at word boundary if it exceeds the limit.
117fn truncate_caption(caption: &str, max_len: usize) -> String {
118    if caption.chars().count() <= max_len {
119        return caption.to_string();
120    }
121
122    let mut truncated = String::new();
123    for (idx, ch) in caption.chars().enumerate() {
124        if idx >= max_len {
125            break;
126        }
127        truncated.push(ch);
128    }
129
130    match truncated.rfind(' ') {
131        Some(pos) => truncated[..pos].to_string(),
132        None => truncated,
133    }
134}
135
136/// Errors that can occur during image captioning.
137#[derive(Debug, thiserror::Error)]
138pub enum CaptionError {
139    #[error("Cannot connect to Ollama at {0}: {1}")]
140    Connection(String, String),
141
142    #[error("Ollama returned HTTP {0}: {1}")]
143    OllamaError(u16, String),
144
145    #[error("Invalid response from Ollama: {0}")]
146    InvalidResponse(String),
147
148    #[error("Failed to read image: {0}")]
149    ImageRead(String),
150
151    #[error("Ollama returned empty caption")]
152    EmptyCaption,
153}
154
155impl CaptionError {
156    /// Stable string discriminant for programmatic matching.
157    pub fn kind(&self) -> &'static str {
158        match self {
159            Self::Connection(..) => "connection",
160            Self::OllamaError(..) => "ollama_error",
161            Self::InvalidResponse(_) => "invalid_response",
162            Self::ImageRead(_) => "image_read",
163            Self::EmptyCaption => "empty_caption",
164        }
165    }
166}
167
168fn read_image_base64(path: &Path) -> Result<String, CaptionError> {
169    let bytes = std::fs::read(path)
170        .map_err(|e| CaptionError::ImageRead(format!("{}: {}", path.display(), e)))?;
171    Ok(base64::Engine::encode(
172        &base64::engine::general_purpose::STANDARD,
173        &bytes,
174    ))
175}
176
177#[cfg(test)]
178mod tests {
179    use super::truncate_caption;
180
181    #[test]
182    fn truncate_caption_preserves_utf8_boundaries() {
183        assert_eq!(truncate_caption("hello world", 7), "hello");
184        assert_eq!(truncate_caption("こんにちは 世界", 5), "こんにちは");
185        assert_eq!(truncate_caption("emoji 😀 test", 8), "emoji 😀");
186    }
187}
ollama_vision/captioner.rs

ollama_vision/
captioner.rs
