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

use std::pin::Pin;

#[cfg(feature = "tracing")]
use tracing::instrument;

use async_trait::async_trait;
use bytes::Bytes;
use futures::{Stream, StreamExt};
use reqwest::{Client, Url};

use crate::transport::Transport;
use crate::types::chat::ChatRequest;
use crate::types::generate::GenerateRequest;
use crate::{Error, Result};

pub struct ReqwestTransport {
    client: Client,
    base_url: Url,
    api_key: Option<String>,
}

impl ReqwestTransport {
    pub fn new(base_url: Url, api_key: Option<String>) -> Result<Self> {
        let client = Client::builder()
            .build()
            .map_err(|e| Error::Client(e.to_string()))?;
        Ok(Self {
            client,
            base_url,
            api_key,
        })
    }
}

#[async_trait]
impl Transport for ReqwestTransport {
    #[cfg_attr(feature = "tracing", instrument(skip(self, request)))]
    async fn send_generate_request(
        &self,
        request: GenerateRequest,
    ) -> Result<Pin<Box<dyn Stream<Item = Result<Bytes>> + Send>>> {
        let url = self
            .base_url
            .join("/api/generate")
            .map_err(|e| Error::Client(e.to_string()))?;

        let mut request_builder = self.client.post(url).json(&request);

        if let Some(api_key) = &self.api_key {
            request_builder = request_builder.bearer_auth(api_key);
        }

        let response = request_builder.send().await.map_err(Error::Transport)?;

        response.error_for_status_ref().map_err(Error::Transport)?;

        let stream = response
            .bytes_stream()
            .map(|item| item.map_err(Error::Transport))
            .boxed();

        Ok(stream)
    }

    #[cfg_attr(feature = "tracing", instrument(skip(self, request)))]
    async fn send_chat_request(
        &self,
        request: ChatRequest,
    ) -> Result<Pin<Box<dyn futures::Stream<Item = Result<Bytes>> + Send>>> {
        let mut url = self
            .base_url
            .join("/api/chat")
            .map_err(|e| Error::Client(e.to_string()))?;
        if request.stream.unwrap_or(false) {
            url.query_pairs_mut().append_pair("stream", "true");
        }

        let mut request_builder = self.client.post(url).json(&request);

        if let Some(api_key) = &self.api_key {
            request_builder = request_builder.bearer_auth(api_key);
        }

        let response = request_builder.send().await.map_err(Error::Transport)?;

        response.error_for_status_ref().map_err(Error::Transport)?;

        let stream = response
            .bytes_stream()
            .map(|item| item.map_err(Error::Transport))
            .boxed();

        Ok(stream)
    }
}