vtcode_core/llm/providers/ollama/
parser.rs1use serde_json::Value as JsonValue;
4
5use super::pull::OllamaPullEvent;
6
7pub 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}