1use serde::{Deserialize, Serialize};
4use serde_json::{Map, Value, json};
5use std::collections::BTreeMap;
6use std::fmt;
7use std::io::{Read, Write};
8use std::net::{TcpStream, ToSocketAddrs};
9use std::time::Duration;
10use traverse_contracts::ExecutionTarget;
11use traverse_registry::{
12 ApplicationModelDependency, ModelAvailabilityProbe, ModelCandidate, ModelCandidateAvailability,
13 ModelCandidateRejectionCode, ModelResolutionEvidence, ModelResolutionPhase,
14 ModelResolutionRequest, resolve_model_dependency,
15};
16
17const OLLAMA_PROVIDER: &str = "ollama";
18const GENERATE_INTERFACE: &str = "traverse.inference.generate";
19const DEFAULT_TIMEOUT_MS: u64 = 30_000;
20const DEFAULT_MAX_RESPONSE_BYTES: u64 = 8 * 1024 * 1024;
24
25#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
26pub struct OllamaProviderConfig {
27 pub base_url: String,
28 #[serde(default)]
29 pub request_timeout_ms: Option<u64>,
30 #[serde(default)]
33 pub max_response_bytes: Option<u64>,
34}
35
36impl OllamaProviderConfig {
37 #[must_use]
38 pub fn timeout(&self) -> Duration {
39 Duration::from_millis(self.request_timeout_ms.unwrap_or(DEFAULT_TIMEOUT_MS))
40 }
41
42 #[must_use]
43 pub fn max_response_bytes(&self) -> u64 {
44 self.max_response_bytes
45 .unwrap_or(DEFAULT_MAX_RESPONSE_BYTES)
46 }
47}
48
49#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
50pub struct OllamaInferenceRequest {
51 pub model: String,
52 pub prompt: String,
53 #[serde(default)]
54 pub system_prompt: Option<String>,
55 #[serde(default)]
56 pub options: Value,
57}
58
59#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
60pub struct OllamaInferenceOutput {
61 pub interface_id: String,
62 pub provider: String,
63 pub provider_implementation_id: String,
64 pub model: String,
65 pub response: String,
66 pub done: bool,
67 pub evidence: OllamaInferenceEvidence,
68}
69
70#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
71pub struct OllamaInferenceEvidence {
72 pub placement_target: String,
73 pub selected_provider: String,
74 pub selected_model: String,
75}
76
77#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
78pub struct GovernedModelExecutionRequest {
79 pub interface_id: String,
80 pub prompt: String,
81 #[serde(default)]
82 pub system_prompt: Option<String>,
83 #[serde(default)]
84 pub options: Value,
85 pub requested_placement: ExecutionTarget,
86 #[serde(default)]
87 pub provider_configs: BTreeMap<String, OllamaProviderConfig>,
88}
89
90#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
91pub struct GovernedModelExecutionOutcome {
92 pub output: OllamaInferenceOutput,
93 pub model_resolution: ModelResolutionEvidence,
94}
95
96#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
97#[serde(rename_all = "snake_case")]
98pub enum GovernedModelExecutionErrorCode {
99 InterfaceNotDeclared,
100 ModelDependencyUnsatisfied,
101 ProviderExecutionFailed,
102}
103
104impl GovernedModelExecutionErrorCode {
105 #[must_use]
106 pub const fn as_str(self) -> &'static str {
107 match self {
108 Self::InterfaceNotDeclared => "model_interface_not_declared",
109 Self::ModelDependencyUnsatisfied => "model_dependency_unsatisfied",
110 Self::ProviderExecutionFailed => "model_provider_failure",
111 }
112 }
113}
114
115#[derive(Debug, Clone, PartialEq, Eq)]
116pub struct GovernedModelExecutionError {
117 pub code: GovernedModelExecutionErrorCode,
118 pub message: String,
119 pub model_resolution: Option<Box<ModelResolutionEvidence>>,
120}
121
122impl GovernedModelExecutionError {
123 #[must_use]
124 pub fn new(code: GovernedModelExecutionErrorCode, message: impl Into<String>) -> Self {
125 Self {
126 code,
127 message: message.into(),
128 model_resolution: None,
129 }
130 }
131
132 #[must_use]
133 pub fn with_model_resolution(mut self, evidence: ModelResolutionEvidence) -> Self {
134 self.model_resolution = Some(Box::new(evidence));
135 self
136 }
137}
138
139impl fmt::Display for GovernedModelExecutionError {
140 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
141 write!(f, "{}: {}", self.code.as_str(), self.message)
142 }
143}
144
145impl std::error::Error for GovernedModelExecutionError {}
146
147#[derive(Debug, Clone, PartialEq, Eq)]
148pub struct OllamaInferenceProvider {
149 config: OllamaProviderConfig,
150}
151
152impl OllamaInferenceProvider {
153 pub fn new(config: OllamaProviderConfig) -> Result<Self, OllamaInferenceError> {
159 parse_base_url(&config.base_url)?;
160 Ok(Self { config })
161 }
162
163 fn from_validated_config(config: OllamaProviderConfig) -> Self {
164 Self { config }
165 }
166
167 #[must_use]
168 pub fn provider_implementation_id(&self) -> &'static str {
169 "ollama.local.generate"
170 }
171
172 pub fn check_model_available(&self, model: &str) -> Result<(), OllamaInferenceError> {
179 if model.trim().is_empty() {
180 return Err(OllamaInferenceError::new(
181 OllamaInferenceErrorCode::InvalidConfig,
182 "model_identifier is required",
183 ));
184 }
185
186 let response = self.request("GET", "/api/tags", None)?;
187 let models = response
188 .get("models")
189 .and_then(Value::as_array)
190 .ok_or_else(|| {
191 OllamaInferenceError::new(
192 OllamaInferenceErrorCode::InvalidResponse,
193 "Ollama tags response must contain models array",
194 )
195 })?;
196
197 if models.iter().any(|entry| model_entry_matches(entry, model)) {
198 return Ok(());
199 }
200
201 Err(OllamaInferenceError::new(
202 OllamaInferenceErrorCode::ModelUnavailable,
203 format!("Ollama model {model} is not installed"),
204 ))
205 }
206
207 pub fn generate(
214 &self,
215 request: &OllamaInferenceRequest,
216 ) -> Result<OllamaInferenceOutput, OllamaInferenceError> {
217 validate_generate_request(request)?;
218 self.check_model_available(&request.model)?;
219
220 let mut body = Map::new();
221 body.insert("model".to_string(), json!(request.model));
222 body.insert("prompt".to_string(), json!(request.prompt));
223 body.insert("stream".to_string(), json!(false));
224 if let Some(system_prompt) = &request.system_prompt {
225 body.insert("system".to_string(), json!(system_prompt));
226 }
227 if !request.options.is_null() {
228 body.insert("options".to_string(), request.options.clone());
229 }
230
231 let response = self.request("POST", "/api/generate", Some(Value::Object(body)))?;
232 parse_generate_response(self.provider_implementation_id(), &request.model, &response)
233 }
234
235 fn request(
236 &self,
237 method: &str,
238 path: &str,
239 body: Option<Value>,
240 ) -> Result<Value, OllamaInferenceError> {
241 let endpoint = parse_base_url(&self.config.base_url)?;
242 let body_text = body.map_or_else(String::new, |value| value.to_string());
243 let request_path = endpoint.path_for(path);
244 let response_text = send_http_json(
245 &endpoint.host,
246 endpoint.port,
247 &request_path,
248 method,
249 &body_text,
250 self.config.timeout(),
251 self.config.max_response_bytes(),
252 )?;
253 parse_http_json_response(&response_text)
254 }
255}
256
257#[derive(Debug, Clone, PartialEq, Eq, Default)]
258pub struct OllamaModelAvailabilityProbe {
259 configs_by_implementation: BTreeMap<String, OllamaProviderConfig>,
260}
261
262impl OllamaModelAvailabilityProbe {
263 #[must_use]
264 pub fn new(config: OllamaProviderConfig) -> Self {
265 let mut configs_by_implementation = BTreeMap::new();
266 configs_by_implementation.insert("ollama.local.generate".to_string(), config);
267 Self {
268 configs_by_implementation,
269 }
270 }
271
272 #[must_use]
273 pub fn with_provider_config(
274 mut self,
275 provider_implementation_id: impl Into<String>,
276 config: OllamaProviderConfig,
277 ) -> Self {
278 self.configs_by_implementation
279 .insert(provider_implementation_id.into(), config);
280 self
281 }
282}
283
284impl ModelAvailabilityProbe for OllamaModelAvailabilityProbe {
285 fn check_candidate(
286 &self,
287 _dependency: &ApplicationModelDependency,
288 candidate: &ModelCandidate,
289 ) -> ModelCandidateAvailability {
290 let Some(config) = self
291 .configs_by_implementation
292 .get(&candidate.provider_implementation_id)
293 else {
294 return ModelCandidateAvailability::rejected(
295 ModelCandidateRejectionCode::ModelCandidateConfigInvalid,
296 "missing provider config for model candidate",
297 );
298 };
299 let provider = match OllamaInferenceProvider::new(config.clone()) {
300 Ok(provider) => provider,
301 Err(error) => return model_candidate_availability_error(&error),
302 };
303 match provider.check_model_available(&candidate.model_identifier) {
304 Ok(()) => ModelCandidateAvailability::ready(),
305 Err(error) => model_candidate_availability_error(&error),
306 }
307 }
308}
309
310#[must_use]
311pub fn resolve_ollama_model_dependency(
312 dependency: &ApplicationModelDependency,
313 request: &ModelResolutionRequest,
314 probe: &OllamaModelAvailabilityProbe,
315) -> ModelResolutionEvidence {
316 resolve_model_dependency(dependency, request, probe)
317}
318
319pub fn execute_governed_ollama_model_dependency(
330 dependency: &ApplicationModelDependency,
331 request: &GovernedModelExecutionRequest,
332) -> Result<GovernedModelExecutionOutcome, GovernedModelExecutionError> {
333 if dependency.interface_id != request.interface_id {
334 return Err(GovernedModelExecutionError::new(
335 GovernedModelExecutionErrorCode::InterfaceNotDeclared,
336 "requested inference interface is not declared by this app dependency",
337 ));
338 }
339
340 let probe = request.provider_configs.iter().fold(
341 OllamaModelAvailabilityProbe::default(),
342 |probe, (implementation_id, config)| {
343 probe.with_provider_config(implementation_id.clone(), config.clone())
344 },
345 );
346 let resolution_request = ModelResolutionRequest {
347 phase: ModelResolutionPhase::Execution,
348 requested_interface_id: request.interface_id.clone(),
349 requested_placement: request.requested_placement.clone(),
350 };
351 let evidence = resolve_ollama_model_dependency(dependency, &resolution_request, &probe);
352 let Some(selected) = evidence.selected.as_ref() else {
353 return Err(GovernedModelExecutionError::new(
354 GovernedModelExecutionErrorCode::ModelDependencyUnsatisfied,
355 "no app-declared model candidate satisfied execution-time resolution",
356 )
357 .with_model_resolution(evidence));
358 };
359 let provider = OllamaInferenceProvider::from_validated_config(
360 request.provider_configs[&selected.provider_implementation_id].clone(),
361 );
362 let output = provider
363 .generate(&OllamaInferenceRequest {
364 model: selected.model_identifier.clone(),
365 prompt: request.prompt.clone(),
366 system_prompt: request.system_prompt.clone(),
367 options: request.options.clone(),
368 })
369 .map_err(|error| {
370 GovernedModelExecutionError::new(
371 GovernedModelExecutionErrorCode::ProviderExecutionFailed,
372 error.to_string(),
373 )
374 .with_model_resolution(evidence.clone())
375 })?;
376
377 Ok(GovernedModelExecutionOutcome {
378 output,
379 model_resolution: evidence,
380 })
381}
382
383#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
384#[serde(rename_all = "snake_case")]
385pub enum OllamaInferenceErrorCode {
386 InvalidConfig,
387 ModelProviderUnavailable,
388 ModelUnavailable,
389 ProviderFailure,
390 InvalidResponse,
391}
392
393impl OllamaInferenceErrorCode {
394 #[must_use]
395 pub const fn as_str(self) -> &'static str {
396 match self {
397 Self::InvalidConfig => "model_candidate_config_invalid",
398 Self::ModelProviderUnavailable => "model_provider_unavailable",
399 Self::ModelUnavailable => "model_candidate_unavailable",
400 Self::ProviderFailure => "model_provider_failure",
401 Self::InvalidResponse => "model_provider_invalid_response",
402 }
403 }
404}
405
406#[derive(Debug, Clone, PartialEq, Eq)]
407pub struct OllamaInferenceError {
408 pub code: OllamaInferenceErrorCode,
409 pub message: String,
410}
411
412impl OllamaInferenceError {
413 #[must_use]
414 pub fn new(code: OllamaInferenceErrorCode, message: impl Into<String>) -> Self {
415 Self {
416 code,
417 message: message.into(),
418 }
419 }
420
421 #[must_use]
422 pub fn machine_code(&self) -> &'static str {
423 self.code.as_str()
424 }
425}
426
427impl fmt::Display for OllamaInferenceError {
428 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
429 write!(f, "{}: {}", self.machine_code(), self.message)
430 }
431}
432
433impl std::error::Error for OllamaInferenceError {}
434
435impl From<std::io::Error> for OllamaInferenceError {
436 fn from(error: std::io::Error) -> Self {
437 provider_unavailable(format!("Ollama I/O failed: {error}"))
438 }
439}
440
441#[derive(Debug, Clone, PartialEq, Eq)]
442struct HttpEndpoint {
443 host: String,
444 port: u16,
445 base_path: String,
446}
447
448impl HttpEndpoint {
449 fn path_for(&self, path: &str) -> String {
450 let base = self.base_path.trim_end_matches('/');
451 let suffix = path.trim_start_matches('/');
452 if base.is_empty() {
453 format!("/{suffix}")
454 } else {
455 format!("{base}/{suffix}")
456 }
457 }
458}
459
460fn validate_generate_request(request: &OllamaInferenceRequest) -> Result<(), OllamaInferenceError> {
461 if request.model.trim().is_empty() {
462 return Err(OllamaInferenceError::new(
463 OllamaInferenceErrorCode::InvalidConfig,
464 "model_identifier is required",
465 ));
466 }
467 if request.prompt.trim().is_empty() {
468 return Err(OllamaInferenceError::new(
469 OllamaInferenceErrorCode::InvalidConfig,
470 "prompt is required for traverse.inference.generate",
471 ));
472 }
473 if !request.options.is_null() && !request.options.is_object() {
474 return Err(OllamaInferenceError::new(
475 OllamaInferenceErrorCode::InvalidConfig,
476 "options must be a JSON object when provided",
477 ));
478 }
479 Ok(())
480}
481
482fn parse_generate_response(
483 provider_implementation_id: &str,
484 requested_model: &str,
485 response: &Value,
486) -> Result<OllamaInferenceOutput, OllamaInferenceError> {
487 let text = response
488 .get("response")
489 .and_then(Value::as_str)
490 .ok_or_else(|| {
491 OllamaInferenceError::new(
492 OllamaInferenceErrorCode::InvalidResponse,
493 "Ollama generate response must contain response text",
494 )
495 })?;
496 let done = response
497 .get("done")
498 .and_then(Value::as_bool)
499 .unwrap_or(false);
500 let model = response
501 .get("model")
502 .and_then(Value::as_str)
503 .unwrap_or(requested_model);
504
505 Ok(OllamaInferenceOutput {
506 interface_id: GENERATE_INTERFACE.to_string(),
507 provider: OLLAMA_PROVIDER.to_string(),
508 provider_implementation_id: provider_implementation_id.to_string(),
509 model: model.to_string(),
510 response: text.to_string(),
511 done,
512 evidence: OllamaInferenceEvidence {
513 placement_target: "local".to_string(),
514 selected_provider: OLLAMA_PROVIDER.to_string(),
515 selected_model: model.to_string(),
516 },
517 })
518}
519
520fn model_entry_matches(entry: &Value, model: &str) -> bool {
521 entry
522 .get("name")
523 .or_else(|| entry.get("model"))
524 .and_then(Value::as_str)
525 .is_some_and(|name| name == model)
526}
527
528fn parse_base_url(base_url: &str) -> Result<HttpEndpoint, OllamaInferenceError> {
529 let trimmed = base_url.trim();
530 let remainder = trimmed.strip_prefix("http://").ok_or_else(|| {
531 OllamaInferenceError::new(
532 OllamaInferenceErrorCode::InvalidConfig,
533 "ollama_base_url must use http:// for local Ollama",
534 )
535 })?;
536 let remainder = remainder.trim_end_matches('/');
537 let (authority, path) = remainder.split_once('/').unwrap_or((remainder, ""));
538 if authority.is_empty() {
539 return Err(OllamaInferenceError::new(
540 OllamaInferenceErrorCode::InvalidConfig,
541 "ollama_base_url must include host",
542 ));
543 }
544
545 let (host, port) = parse_authority(authority)?;
546 Ok(HttpEndpoint {
547 host,
548 port,
549 base_path: format!("/{path}").trim_end_matches('/').to_string(),
550 })
551}
552
553fn parse_authority(authority: &str) -> Result<(String, u16), OllamaInferenceError> {
554 let (host, port) = if let Some((host, port_text)) = authority.rsplit_once(':') {
555 let port = port_text.parse::<u16>().map_err(|error| {
556 OllamaInferenceError::new(
557 OllamaInferenceErrorCode::InvalidConfig,
558 format!("ollama_base_url port is invalid: {error}"),
559 )
560 })?;
561 (host.to_string(), port)
562 } else {
563 (authority.to_string(), 11_434)
564 };
565
566 if host.is_empty() {
567 return Err(OllamaInferenceError::new(
568 OllamaInferenceErrorCode::InvalidConfig,
569 "ollama_base_url host is required",
570 ));
571 }
572
573 Ok((host, port))
574}
575
576fn send_http_json(
577 host: &str,
578 port: u16,
579 path: &str,
580 method: &str,
581 body: &str,
582 timeout: Duration,
583 max_response_bytes: u64,
584) -> Result<String, OllamaInferenceError> {
585 let address = format!("{host}:{port}");
586 let socket_address = address
587 .to_socket_addrs()?
588 .next()
589 .ok_or_else(|| provider_unavailable("Ollama endpoint did not resolve"))?;
590 let mut stream = TcpStream::connect_timeout(&socket_address, timeout)?;
591 stream.set_read_timeout(Some(timeout))?;
595 stream.set_write_timeout(Some(timeout))?;
596 let request = format!(
597 "{method} {path} HTTP/1.1\r\nHost: {host}:{port}\r\nContent-Type: application/json\r\nAccept: application/json\r\nConnection: close\r\nContent-Length: {}\r\n\r\n{body}",
598 body.len()
599 );
600 stream.write_all(request.as_bytes())?;
601
602 read_capped_response(&mut stream, max_response_bytes)
603}
604
605fn read_capped_response(
609 stream: &mut TcpStream,
610 max_response_bytes: u64,
611) -> Result<String, OllamaInferenceError> {
612 let mut buffer = Vec::new();
613 let limit = max_response_bytes.saturating_add(1);
614 stream.take(limit).read_to_end(&mut buffer)?;
615 if buffer.len() as u64 > max_response_bytes {
616 return Err(OllamaInferenceError::new(
617 OllamaInferenceErrorCode::InvalidResponse,
618 format!("Ollama response exceeded the {max_response_bytes}-byte limit"),
619 ));
620 }
621 String::from_utf8(buffer).map_err(|error| {
622 OllamaInferenceError::new(
623 OllamaInferenceErrorCode::InvalidResponse,
624 format!("Ollama response was not valid UTF-8: {error}"),
625 )
626 })
627}
628
629fn parse_http_json_response(response: &str) -> Result<Value, OllamaInferenceError> {
630 let (head, body) = response.split_once("\r\n\r\n").ok_or_else(|| {
631 OllamaInferenceError::new(
632 OllamaInferenceErrorCode::InvalidResponse,
633 "Ollama HTTP response is missing header separator",
634 )
635 })?;
636 let status_line = head.lines().next().ok_or_else(|| {
637 OllamaInferenceError::new(
638 OllamaInferenceErrorCode::InvalidResponse,
639 "Ollama HTTP response is missing status line",
640 )
641 })?;
642 let status = parse_status_code(status_line)?;
643 if !(200..300).contains(&status) {
644 return Err(OllamaInferenceError::new(
645 OllamaInferenceErrorCode::ProviderFailure,
646 format!("Ollama returned HTTP {status}"),
647 ));
648 }
649
650 serde_json::from_str(body).map_err(|error| {
651 OllamaInferenceError::new(
652 OllamaInferenceErrorCode::InvalidResponse,
653 format!("Ollama response body is not valid JSON: {error}"),
654 )
655 })
656}
657
658fn parse_status_code(status_line: &str) -> Result<u16, OllamaInferenceError> {
659 status_line
660 .split_whitespace()
661 .nth(1)
662 .ok_or_else(|| {
663 OllamaInferenceError::new(
664 OllamaInferenceErrorCode::InvalidResponse,
665 "Ollama HTTP status line is malformed",
666 )
667 })?
668 .parse::<u16>()
669 .map_err(|error| {
670 OllamaInferenceError::new(
671 OllamaInferenceErrorCode::InvalidResponse,
672 format!("Ollama HTTP status code is invalid: {error}"),
673 )
674 })
675}
676
677fn provider_unavailable(message: impl Into<String>) -> OllamaInferenceError {
678 OllamaInferenceError::new(OllamaInferenceErrorCode::ModelProviderUnavailable, message)
679}
680
681fn model_candidate_availability_error(error: &OllamaInferenceError) -> ModelCandidateAvailability {
682 match error.code {
683 OllamaInferenceErrorCode::InvalidConfig => ModelCandidateAvailability::rejected(
684 ModelCandidateRejectionCode::ModelCandidateConfigInvalid,
685 error.message.clone(),
686 ),
687 OllamaInferenceErrorCode::ModelProviderUnavailable => ModelCandidateAvailability::rejected(
688 ModelCandidateRejectionCode::ModelProviderUnavailable,
689 error.message.clone(),
690 ),
691 OllamaInferenceErrorCode::ModelUnavailable => ModelCandidateAvailability::rejected(
692 ModelCandidateRejectionCode::ModelCandidateUnavailable,
693 error.message.clone(),
694 ),
695 OllamaInferenceErrorCode::ProviderFailure | OllamaInferenceErrorCode::InvalidResponse => {
696 ModelCandidateAvailability::rejected(
697 ModelCandidateRejectionCode::ModelProviderUnavailable,
698 error.message.clone(),
699 )
700 }
701 }
702}