clawft_kernel/workload_runtime/infer/
ollama.rs1use std::sync::{Arc, Mutex};
7
8use serde_json::{Value, json};
9
10use super::probe::ServerClient;
11
12#[derive(Debug, Clone, PartialEq, Eq)]
14pub enum LoadState {
15 Loading,
17 Done,
19 Failed(String),
21}
22
23pub type SharedLoad = Arc<Mutex<LoadState>>;
25
26fn keep_alive_value(k: &str) -> Value {
27 k.parse::<i64>()
28 .map(Value::from)
29 .unwrap_or_else(|_| json!(k))
30}
31
32pub async fn load(client: &ServerClient, tag: &str, keep_alive: &str) -> Result<(), String> {
35 let body = json!({"model": tag, "prompt": "", "stream": false,
36 "keep_alive": keep_alive_value(keep_alive)});
37 match client.post("/api/generate", &body).await? {
38 (200, _) => Ok(()),
39 (c, b) => Err(format!(
40 "ollama load returned HTTP {c}: {}",
41 b.chars().take(200).collect::<String>()
42 )),
43 }
44}
45
46pub async fn resident(client: &ServerClient, tag: &str) -> bool {
49 let Ok((200, body)) = client.get("/api/ps").await else {
50 return false;
51 };
52 let Ok(v) = serde_json::from_str::<Value>(&body) else {
53 return false;
54 };
55 let same = |n: &str| n == tag || n.strip_suffix(":latest") == Some(tag) || tag.strip_suffix(":latest") == Some(n);
56 v.get("models").and_then(Value::as_array).is_some_and(|ms| {
57 ms.iter()
58 .any(|m| ["name", "model"].iter().any(|k| m.get(k).and_then(Value::as_str).is_some_and(same)))
59 })
60}
61
62pub async fn unload(client: &ServerClient, tag: &str) -> Result<(), String> {
64 let body = json!({"model": tag, "keep_alive": 0});
65 match client.post("/api/generate", &body).await? {
66 (200, _) => Ok(()),
67 (c, _) => Err(format!("ollama unload returned HTTP {c}")),
68 }
69}