Skip to main content

ollama_vision/
types.rs

1use serde::{Deserialize, Serialize};
2use std::time::Duration;
3
4/// Configuration for the Ollama vision client.
5#[derive(Debug, Clone)]
6pub struct OllamaVisionConfig {
7    /// Ollama API endpoint (e.g., "http://localhost:11434")
8    pub endpoint: String,
9    /// Vision model name (e.g., "llava", "llava-llama3", "minicpm-v")
10    pub model: String,
11    /// Request timeout (default: 120s)
12    pub timeout: Duration,
13    /// TCP connection timeout (default: 10s)
14    pub connect_timeout: Duration,
15    /// Generation options sent to Ollama
16    pub options: GenerateOptions,
17}
18
19/// Ollama generation options controlling output quality.
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct GenerateOptions {
22    /// Maximum tokens to generate
23    pub num_predict: u32,
24    /// Penalize repeated tokens (default: 1.2)
25    pub repeat_penalty: f32,
26    /// Window for repeat penalty (default: 128)
27    pub repeat_last_n: u32,
28    /// Temperature (default: None, uses Ollama default)
29    #[serde(skip_serializing_if = "Option::is_none")]
30    pub temperature: Option<f32>,
31    /// Top-p sampling (default: None, uses Ollama default)
32    #[serde(skip_serializing_if = "Option::is_none")]
33    pub top_p: Option<f32>,
34}
35
36impl Default for GenerateOptions {
37    fn default() -> Self {
38        Self {
39            num_predict: 512,
40            repeat_penalty: 1.2,
41            repeat_last_n: 128,
42            temperature: None,
43            top_p: None,
44        }
45    }
46}
47
48impl Default for OllamaVisionConfig {
49    fn default() -> Self {
50        Self {
51            endpoint: "http://localhost:11434".to_string(),
52            model: "llava".to_string(),
53            timeout: Duration::from_secs(120),
54            connect_timeout: Duration::from_secs(10),
55            options: GenerateOptions::default(),
56        }
57    }
58}
59
60impl OllamaVisionConfig {
61    /// Create a new config with the given model name.
62    pub fn with_model(model: impl Into<String>) -> Self {
63        Self {
64            model: model.into(),
65            ..Default::default()
66        }
67    }
68
69    /// Set the connection timeout.
70    pub fn connect_timeout(mut self, timeout: Duration) -> Self {
71        self.connect_timeout = timeout;
72        self
73    }
74
75    /// Set the Ollama endpoint.
76    pub fn endpoint(mut self, endpoint: impl Into<String>) -> Self {
77        self.endpoint = endpoint.into();
78        self
79    }
80
81    /// Set the request timeout.
82    pub fn timeout(mut self, timeout: Duration) -> Self {
83        self.timeout = timeout;
84        self
85    }
86
87    /// Set the generation options.
88    pub fn options(mut self, options: GenerateOptions) -> Self {
89        self.options = options;
90        self
91    }
92}
93
94/// Tag configuration for controlling tag extraction behavior.
95#[derive(Debug, Clone)]
96pub struct TagOptions {
97    /// Custom system prompt (overrides default)
98    pub prompt: Option<String>,
99    /// Request JSON format from Ollama (default: true)
100    pub request_json_format: bool,
101    /// Maximum number of tags to return. Tags beyond this limit are dropped.
102    /// Default: 30.
103    pub max_tags: usize,
104    /// Maximum length of a single tag in characters. Tags longer than this
105    /// are truncated. Default: 50.
106    pub max_tag_length: usize,
107    /// Maximum number of retries on transient failures. Default: 2.
108    pub max_retries: u32,
109}
110
111impl Default for TagOptions {
112    fn default() -> Self {
113        Self {
114            prompt: None,
115            request_json_format: true,
116            max_tags: 30,
117            max_tag_length: 50,
118            max_retries: 2,
119        }
120    }
121}
122
123/// Caption configuration for controlling caption generation.
124#[derive(Debug, Clone)]
125pub struct CaptionOptions {
126    /// Custom prompt (overrides default)
127    pub prompt: Option<String>,
128    /// Maximum caption length in characters. Captions beyond this are truncated
129    /// at the last word boundary. Default: 500.
130    pub max_caption_length: usize,
131    /// Maximum number of retries on transient failures. Default: 2.
132    pub max_retries: u32,
133}
134
135impl Default for CaptionOptions {
136    fn default() -> Self {
137        Self {
138            prompt: None,
139            max_caption_length: 500,
140            max_retries: 2,
141        }
142    }
143}