Skip to main content

vtcode_core/llm/providers/ollama/
parser.rs

1/// Parse Ollama pull event JSON responses into structured events.
2/// Adapted from OpenAI Codex's codex-ollama/src/parser.rs
3use serde_json::Value as JsonValue;
4
5use super::pull::OllamaPullEvent;
6
7/// Convert a single JSON object representing a pull update into one or more events.
8pub fn pull_events_from_value(value: &JsonValue) -> Vec<OllamaPullEvent> {
9    let mut events = Vec::new();
10
11    if let Some(status) = value.get("status").and_then(|s| s.as_str()) {
12        events.push(OllamaPullEvent::Status(status.to_string()));
13        if status == "success" {
14            events.push(OllamaPullEvent::Success);
15        }
16    }
17
18    let digest = value
19        .get("digest")
20        .and_then(|d| d.as_str())
21        .unwrap_or("")
22        .to_string();
23    let total = value.get("total").and_then(JsonValue::as_u64);
24    let completed = value.get("completed").and_then(JsonValue::as_u64);
25
26    if total.is_some() || completed.is_some() {
27        events.push(OllamaPullEvent::ChunkProgress {
28            digest,
29            total,
30            completed,
31        });
32    }
33
34    events
35}
36
37#[cfg(test)]
38mod tests {
39    use super::*;
40
41    #[test]
42    fn test_pull_events_decoder_status_and_success() {
43        let v: JsonValue = serde_json::json!({"status":"verifying"});
44        let events = pull_events_from_value(&v);
45        assert_eq!(events.len(), 1);
46        match &events[0] {
47            OllamaPullEvent::Status(s) => assert_eq!(s, "verifying"),
48            _ => panic!("Expected Status event"),
49        }
50
51        let v2: JsonValue = serde_json::json!({"status":"success"});
52        let events2 = pull_events_from_value(&v2);
53        assert_eq!(events2.len(), 2);
54        match &events2[0] {
55            OllamaPullEvent::Status(s) => assert_eq!(s, "success"),
56            _ => panic!("Expected Status event"),
57        }
58        match &events2[1] {
59            OllamaPullEvent::Success => {}
60            _ => panic!("Expected Success event"),
61        }
62    }
63
64    #[test]
65    fn test_pull_events_decoder_progress() {
66        let v: JsonValue = serde_json::json!({"digest":"sha256:abc","total":100});
67        let events = pull_events_from_value(&v);
68        assert_eq!(events.len(), 1);
69        match &events[0] {
70            OllamaPullEvent::ChunkProgress {
71                digest,
72                total,
73                completed,
74            } => {
75                assert_eq!(digest, "sha256:abc");
76                assert_eq!(*total, Some(100));
77                assert_eq!(*completed, None);
78            }
79            _ => panic!("Expected ChunkProgress event"),
80        }
81
82        let v2: JsonValue = serde_json::json!({"digest":"sha256:def","completed":42});
83        let events2 = pull_events_from_value(&v2);
84        assert_eq!(events2.len(), 1);
85        match &events2[0] {
86            OllamaPullEvent::ChunkProgress {
87                digest,
88                total,
89                completed,
90            } => {
91                assert_eq!(digest, "sha256:def");
92                assert_eq!(*total, None);
93                assert_eq!(*completed, Some(42));
94            }
95            _ => panic!("Expected ChunkProgress event"),
96        }
97    }
98}