ollama-sdk 0.2.2

An idiomatic, unofficial Rust client for the Ollama API with support for streaming, tool calling, and custom transports.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

use bytes::Bytes;
use futures::{stream, StreamExt};

use ollama_sdk::stream::chat_stream_parser::ChatStreamParser;
use ollama_sdk::types::chat::ChatStreamEvent;
use ollama_sdk::Result;

// Helper function to create a stream from a vector of byte chunks
fn create_byte_stream(
    chunks: Vec<String>,
) -> impl futures::Stream<Item = Result<Bytes>> + Send + 'static {
    stream::iter(chunks.into_iter().map(|s| Ok(Bytes::from(s))))
}

#[tokio::test]
async fn test_parse_single_message_event() {
    let json_line =
        r#"{"model":"gpt-2","message":{"role":"assistant","content":"hello"},"done":false}"#
            .to_string();
    let byte_stream = create_byte_stream(vec![format!("{}\n", json_line)]);
    let mut parser = ChatStreamParser::new(byte_stream);

    let event = parser.next().await.unwrap().unwrap();
    match event {
        ChatStreamEvent::Message(response) => assert_eq!(response.message.content, "hello"),
        _ => panic!("Expected Message event, got {:?}", event),
    }
    assert!(parser.next().await.is_none());
}

#[tokio::test]
async fn test_parse_single_partial_event() {
    let json_line =
        r#"{"message":{"role":"assistant","content":"hello"},"done":false}"#.to_string();
    let byte_stream = create_byte_stream(vec![format!("{}\n", json_line)]);
    let mut parser = ChatStreamParser::new(byte_stream);

    let event = parser.next().await.unwrap().unwrap();
    match event {
        ChatStreamEvent::Partial { partial, .. } => assert_eq!(partial, json_line),
        _ => panic!("Expected Partial event, got {:?}", event),
    }
    assert!(parser.next().await.is_none());
}

#[tokio::test]
async fn test_parse_single_error_event() {
    let json_line = r#"{"error":"Some test error"}"#.to_string();
    let byte_stream = create_byte_stream(vec![format!("{}\n", json_line)]);
    let mut parser = ChatStreamParser::new(byte_stream);

    let event = parser.next().await.unwrap().unwrap();
    match event {
        ChatStreamEvent::Error(err) => assert_eq!(err, "Some test error"),
        _ => panic!("Expected Error event, got {:?}", event),
    }
    assert!(parser.next().await.is_none());
}

#[tokio::test]
async fn test_parse_multiple_message_events() {
    let json_line1 =
        r#"{"model":"gpt-2","message":{"role":"assistant","content":"hello"},"done":false}"#
            .to_string();
    let json_line2 =
        r#"{"model":"gpt-2","message":{"role":"assistant","content":" world"},"done":false}"#
            .to_string();
    let byte_stream = create_byte_stream(vec![
        format!("{}\n", json_line1),
        format!("{}\n", json_line2),
    ]);
    let mut parser = ChatStreamParser::new(byte_stream);

    let ev1 = parser.next().await.unwrap().unwrap();
    let ev2 = parser.next().await.unwrap().unwrap();

    match (ev1, ev2) {
        (ChatStreamEvent::Message(response1), ChatStreamEvent::Message(response2)) => {
            assert_eq!(response1.message.content, "hello");
            assert_eq!(response2.message.content, " world");
        }
        _ => panic!("Expected two Message events"),
    }
    assert!(parser.next().await.is_none());
}

#[tokio::test]
async fn test_handle_mix_of_partial_and_message_events() {
    let non_json_line = "This is a plain text message.".to_string();
    let json_line =
        r#"{"model":"gpt-2","message":{"role":"assistant","content":"JSON part"},"done":false}"#
            .to_string();

    let byte_stream = create_byte_stream(vec![
        format!("{}\n", non_json_line),
        format!("{}\n", json_line),
    ]);
    let mut parser = ChatStreamParser::new(byte_stream);

    let ev1 = parser.next().await.unwrap().unwrap();
    let ev2 = parser.next().await.unwrap().unwrap();

    match (ev1, ev2) {
        (ChatStreamEvent::Partial { partial, .. }, ChatStreamEvent::Message(response2)) => {
            assert_eq!(partial, non_json_line);
            assert_eq!(response2.message.content, "JSON part");
        }
        _ => panic!("Expected one Partial and one Message events"),
    }
    assert!(parser.next().await.is_none());
}

#[tokio::test]
async fn test_empty_stream() {
    let byte_stream = create_byte_stream(vec![]);
    let mut parser = ChatStreamParser::new(byte_stream);
    assert!(parser.next().await.is_none());
}

#[tokio::test]
async fn test_stream_with_empty_lines() {
    let json_line =
        r#"{"model":"gpt-2","message":{"role":"assistant","content":"test"},"done":false}"#
            .to_string();
    let byte_stream = create_byte_stream(vec![
        "\n".to_string(),
        format!("{}\n", json_line),
        "\n\n".to_string(),
    ]);
    let mut parser = ChatStreamParser::new(byte_stream);

    let ev = parser.next().await.unwrap().unwrap();
    match ev {
        ChatStreamEvent::Message(response) => assert_eq!(response.message.content, "test"),
        _ => panic!("Expected Message event"),
    }
    assert!(parser.next().await.is_none());
}