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 std::pin::Pin;

use crate::types::Thinking;
use crate::Result;
use futures::Stream;
use serde::{Deserialize, Serialize};

use super::Role;

#[derive(Serialize, Default, Debug, Clone)]
pub struct ChatRequest {
    pub model: String,
    pub messages: Vec<ChatRequestMessage>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stream: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tools: Option<Vec<ToolSpec>>,
    #[serde(default)]
    pub think: Thinking,
}

#[derive(Deserialize, Serialize, Default, Debug, Clone)]
pub struct ChatRequestMessage {
    pub role: Role,
    pub content: String,
    #[serde(default)]
    pub tool_calls: Vec<FunctionalTool>,
}

#[derive(Serialize, Debug, Clone)]
#[serde(tag = "type")]
pub enum ToolSpec {
    #[serde(rename = "function")]
    Function(FunctionalTool),
}

#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct FunctionalTool {
    pub name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    pub parameters: serde_json::Value,
}

#[derive(Deserialize, Serialize, Default, Debug, Clone)]
pub struct ChatResponse {
    pub model: String,
    #[serde(default)]
    pub created_at: String,
    pub message: ChatResponseMessage,
    pub done: bool,
}

#[derive(Deserialize, Serialize, Default, Debug, Clone)]
pub struct ChatResponseMessage {
    pub role: Role,
    pub content: String,
    #[serde(default)]
    pub thinking: String,
    #[serde(default)]
    pub tool_calls: Vec<FunctionalTool>,
}

impl From<ChatResponseMessage> for ChatRequestMessage {
    fn from(value: ChatResponseMessage) -> Self {
        ChatRequestMessage {
            role: value.role,
            content: value.content,
            tool_calls: value.tool_calls,
        }
    }
}

#[derive(Serialize, Default, Debug, Clone)]
pub struct SimpleChatRequest {
    pub model: String,
    pub messages: Vec<ChatRequestMessage>,
    pub think: Thinking,
}

#[derive(Serialize, Default, Debug, Clone)]
pub struct StreamingChatRequest {
    pub model: String,
    pub messages: Vec<ChatRequestMessage>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tools: Option<Vec<ToolSpec>>,
    pub think: Thinking,
}

impl From<SimpleChatRequest> for ChatRequest {
    fn from(value: SimpleChatRequest) -> Self {
        ChatRequest {
            model: value.model,
            messages: value.messages,
            stream: Some(false),
            think: value.think,
            tools: None,
        }
    }
}

impl From<StreamingChatRequest> for ChatRequest {
    fn from(value: StreamingChatRequest) -> Self {
        ChatRequest {
            model: value.model,
            messages: value.messages,
            stream: Some(true),
            think: value.think,
            tools: value.tools,
        }
    }
}

#[derive(Deserialize, Serialize, Debug)]
pub enum ChatStreamEvent {
    Message(ChatResponse),
    Error(String),
    Partial {
        partial: String,
        error: Option<String>,
    },
}

// ChatStream type
pub struct ChatStream {
    pub inner: Pin<Box<dyn Stream<Item = Result<ChatStreamEvent>> + Send>>,
}

impl Stream for ChatStream {
    type Item = Result<ChatStreamEvent>;

    fn poll_next(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        self.inner.as_mut().poll_next(cx)
    }
}