Skip to main content

clawft_kernel/workload_runtime/infer/
ollama.rs

1//! Ollama is driven through its API: it owns its process and decides when
2//! weights are resident, so the adapter asks it to load or unload a model
3//! and reports what `/api/ps` says. It never starts, stops or signals the
4//! Ollama server itself.
5
6use std::sync::{Arc, Mutex};
7
8use serde_json::{Value, json};
9
10use super::probe::ServerClient;
11
12/// Progress of the background load request.
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub enum LoadState {
15    /// The request is in flight.
16    Loading,
17    /// The request returned.
18    Done,
19    /// The request failed.
20    Failed(String),
21}
22
23/// Shared handle on a [`LoadState`].
24pub 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
32/// Ask Ollama to bring `tag` into memory (an empty prompt generates
33/// nothing). Blocks until it answers, so run it on a task.
34pub 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
46/// Whether `tag` is in memory now (`/api/ps`); false when Ollama does not
47/// answer, so a load is attempted and reports the failure itself.
48pub 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
62/// Ask Ollama to drop `tag` from memory now (`keep_alive: 0`).
63pub 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}