clnrm_core/services/
ollama.rs

1//! Ollama AI Service Plugin
2//!
3//! Provides integration with Ollama AI services for testing AI functionality.
4//! Supports model loading, text generation, and health monitoring.
5
6use crate::cleanroom::{HealthStatus, ServiceHandle, ServicePlugin};
7use crate::error::{CleanroomError, Result};
8use serde_json::json;
9use std::collections::HashMap;
10use std::sync::Arc;
11use tokio::sync::RwLock;
12use uuid::Uuid;
13
14/// Ollama AI service plugin configuration
15#[derive(Debug, Clone)]
16pub struct OllamaConfig {
17    /// Service endpoint URL
18    pub endpoint: String,
19    /// Default model to use
20    pub default_model: String,
21    /// Request timeout in seconds
22    pub timeout_seconds: u64,
23}
24
25/// Ollama AI service plugin
26#[derive(Debug)]
27pub struct OllamaPlugin {
28    name: String,
29    config: OllamaConfig,
30    client: Arc<RwLock<Option<reqwest::Client>>>,
31}
32
33impl OllamaPlugin {
34    /// Create a new Ollama plugin instance
35    pub fn new(name: &str, config: OllamaConfig) -> Self {
36        Self {
37            name: name.to_string(),
38            config,
39            client: Arc::new(RwLock::new(None)),
40        }
41    }
42
43    /// Initialize the HTTP client for Ollama API calls
44    async fn init_client(&self) -> Result<reqwest::Client> {
45        let client = reqwest::Client::builder()
46            .timeout(std::time::Duration::from_secs(self.config.timeout_seconds))
47            .build()
48            .map_err(|e| {
49                CleanroomError::internal_error(format!("Failed to create HTTP client: {}", e))
50            })?;
51
52        Ok(client)
53    }
54
55    /// Test connection to Ollama service
56    async fn test_connection(&self) -> Result<()> {
57        let client = self.init_client().await?;
58        let url = format!("{}/api/version", self.config.endpoint);
59
60        let response = client.get(&url).send().await.map_err(|e| {
61            CleanroomError::service_error(format!("Failed to connect to Ollama: {}", e))
62        })?;
63
64        if response.status().is_success() {
65            Ok(())
66        } else {
67            Err(CleanroomError::service_error(
68                "Ollama service not responding",
69            ))
70        }
71    }
72
73    /// Generate text using Ollama AI model
74    pub async fn generate_text(&self, model: &str, prompt: &str) -> Result<OllamaResponse> {
75        let mut client_guard = self.client.write().await;
76        if client_guard.is_none() {
77            *client_guard = Some(self.init_client().await?);
78        }
79        let client = client_guard
80            .as_ref()
81            .ok_or_else(|| CleanroomError::internal_error("HTTP client not initialized"))?;
82
83        let url = format!("{}/api/generate", self.config.endpoint);
84        let payload = json!({
85            "model": model,
86            "prompt": prompt,
87            "stream": false
88        });
89
90        let response = client
91            .post(&url)
92            .header("Content-Type", "application/json")
93            .json(&payload)
94            .send()
95            .await
96            .map_err(|e| {
97                CleanroomError::service_error(format!("Failed to generate text: {}", e))
98            })?;
99
100        if response.status().is_success() {
101            let ollama_response: OllamaResponse = response.json().await.map_err(|e| {
102                CleanroomError::service_error(format!("Failed to parse response: {}", e))
103            })?;
104
105            Ok(ollama_response)
106        } else {
107            let error_text = response
108                .text()
109                .await
110                .unwrap_or_else(|_| "Unknown error".to_string());
111
112            Err(CleanroomError::service_error(format!(
113                "Ollama API error: {}",
114                error_text
115            )))
116        }
117    }
118
119    /// Get available models from Ollama
120    pub async fn list_models(&self) -> Result<Vec<OllamaModel>> {
121        let mut client_guard = self.client.write().await;
122        if client_guard.is_none() {
123            *client_guard = Some(self.init_client().await?);
124        }
125        let client = client_guard
126            .as_ref()
127            .ok_or_else(|| CleanroomError::internal_error("HTTP client not initialized"))?;
128
129        let url = format!("{}/api/tags", self.config.endpoint);
130
131        let response =
132            client.get(&url).send().await.map_err(|e| {
133                CleanroomError::service_error(format!("Failed to list models: {}", e))
134            })?;
135
136        if response.status().is_success() {
137            let model_list: OllamaModelList = response.json().await.map_err(|e| {
138                CleanroomError::service_error(format!("Failed to parse model list: {}", e))
139            })?;
140
141            Ok(model_list.models)
142        } else {
143            Err(CleanroomError::service_error(
144                "Failed to retrieve model list",
145            ))
146        }
147    }
148}
149
150/// Response from Ollama text generation
151#[derive(Debug, serde::Deserialize)]
152pub struct OllamaResponse {
153    /// Generated text response
154    pub response: String,
155    /// Model used for generation
156    pub model: String,
157    /// Creation timestamp
158    pub created_at: String,
159    /// Whether generation is complete
160    pub done: bool,
161    /// Reason for completion
162    pub done_reason: String,
163    /// Total duration in nanoseconds
164    pub total_duration: u64,
165    /// Load duration in nanoseconds
166    pub load_duration: u64,
167    /// Number of tokens in prompt
168    pub prompt_eval_count: u32,
169    /// Duration for prompt evaluation
170    pub prompt_eval_duration: u64,
171    /// Number of tokens generated
172    pub eval_count: u32,
173    /// Duration for token generation
174    pub eval_duration: u64,
175}
176
177/// Ollama model information
178#[derive(Debug, serde::Deserialize)]
179pub struct OllamaModel {
180    /// Model name
181    pub name: String,
182    /// Model identifier
183    pub model: String,
184    /// Last modified timestamp
185    pub modified_at: String,
186    /// Model size in bytes
187    pub size: u64,
188    /// Model digest
189    pub digest: String,
190    /// Model details
191    pub details: OllamaModelDetails,
192}
193
194/// Ollama model details
195#[derive(Debug, serde::Deserialize)]
196pub struct OllamaModelDetails {
197    /// Parent model
198    pub parent_model: String,
199    /// Model format
200    pub format: String,
201    /// Model family
202    pub family: String,
203    /// Model families
204    pub families: Vec<String>,
205    /// Parameter size
206    pub parameter_size: String,
207    /// Quantization level
208    pub quantization_level: String,
209}
210
211/// List of available models
212#[derive(Debug, serde::Deserialize)]
213pub struct OllamaModelList {
214    /// List of available models
215    pub models: Vec<OllamaModel>,
216}
217
218impl ServicePlugin for OllamaPlugin {
219    fn name(&self) -> &str {
220        &self.name
221    }
222
223    fn start(&self) -> Result<ServiceHandle> {
224        // Use tokio::task::block_in_place for async operations
225        tokio::task::block_in_place(|| {
226            tokio::runtime::Handle::current().block_on(async {
227                // Test connection to Ollama service
228                let health_check = async {
229                    // Simple health check - try to connect
230                    match self.test_connection().await {
231                        Ok(_) => HealthStatus::Healthy,
232                        Err(_) => HealthStatus::Unhealthy,
233                    }
234                };
235
236                let health = health_check.await;
237
238                let mut metadata = HashMap::new();
239                metadata.insert("endpoint".to_string(), self.config.endpoint.clone());
240                metadata.insert(
241                    "default_model".to_string(),
242                    self.config.default_model.clone(),
243                );
244                metadata.insert(
245                    "timeout_seconds".to_string(),
246                    self.config.timeout_seconds.to_string(),
247                );
248                metadata.insert("health_status".to_string(), format!("{:?}", health));
249
250                Ok(ServiceHandle {
251                    id: Uuid::new_v4().to_string(),
252                    service_name: self.name.clone(),
253                    metadata,
254                })
255            })
256        })
257    }
258
259    fn stop(&self, _handle: ServiceHandle) -> Result<()> {
260        // HTTP-based service, no cleanup needed beyond dropping the client
261        Ok(())
262    }
263
264    fn health_check(&self, handle: &ServiceHandle) -> HealthStatus {
265        if let Some(health_status) = handle.metadata.get("health_status") {
266            match health_status.as_str() {
267                "Healthy" => HealthStatus::Healthy,
268                "Unhealthy" => HealthStatus::Unhealthy,
269                _ => HealthStatus::Unknown,
270            }
271        } else {
272            HealthStatus::Unknown
273        }
274    }
275}
clnrm_core/services/ollama.rs

clnrm_core/services/
ollama.rs
