Adds pull method and example
This commit is contained in:
38
src/lib.rs
38
src/lib.rs
@@ -13,6 +13,7 @@ use crate::{
|
||||
chat::{ChatRequest, ChatResponse},
|
||||
generate::{GenerateRequest, GenerateResponse},
|
||||
ps::RunningModel,
|
||||
pull::{PullRequest, PullResponse},
|
||||
tags::Model,
|
||||
},
|
||||
};
|
||||
@@ -154,4 +155,41 @@ impl OllamaClient {
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn pull(
|
||||
&self,
|
||||
request: PullRequest,
|
||||
) -> impl Stream<Item = OllamaResult<PullResponse>> {
|
||||
let request_address = format!("{}/api/pull", self.server_address);
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
Box::pin(stream! {
|
||||
let response = client
|
||||
.post(request_address)
|
||||
.json(&request)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| OllamaError::from(e))?; // Adjust based on your error type
|
||||
|
||||
let bytes_stream = response.bytes_stream();
|
||||
|
||||
let body_reader = StreamReader::new(
|
||||
bytes_stream.map(|res| res.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))),
|
||||
);
|
||||
|
||||
let mut lines_stream = FramedRead::new(body_reader, LinesCodec::new());
|
||||
|
||||
while let Some(line_result) = lines_stream.next().await {
|
||||
match line_result {
|
||||
Ok(line_content) => {
|
||||
println!("{line_content}");
|
||||
if let Ok(parsed) = serde_json::from_str::<PullResponse>(&line_content) {
|
||||
yield Ok(parsed);
|
||||
}
|
||||
}
|
||||
Err(e) => yield Err(OllamaError::from(e)),
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,4 +2,5 @@ pub mod chat;
|
||||
pub mod common;
|
||||
pub mod generate;
|
||||
pub mod ps;
|
||||
pub mod pull;
|
||||
pub mod tags;
|
||||
|
||||
45
src/types/pull.rs
Normal file
45
src/types/pull.rs
Normal file
@@ -0,0 +1,45 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct PullRequest {
|
||||
pub model: String,
|
||||
pub insecure: Option<bool>,
|
||||
pub stream: Option<bool>,
|
||||
}
|
||||
|
||||
impl PullRequest {
|
||||
pub fn builder<M: Into<String>>(model: M) -> PullRequestBuilder {
|
||||
PullRequestBuilder {
|
||||
pull_request: PullRequest {
|
||||
model: model.into(),
|
||||
insecure: None,
|
||||
stream: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct PullRequestBuilder {
|
||||
pull_request: PullRequest,
|
||||
}
|
||||
|
||||
impl PullRequestBuilder {
|
||||
pub fn stream(mut self, stream: bool) -> Self {
|
||||
self.pull_request.stream = Some(stream);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn insecure(mut self, insecure: bool) -> Self {
|
||||
self.pull_request.insecure = Some(insecure);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn build(self) -> PullRequest {
|
||||
self.pull_request
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
pub struct PullResponse {
|
||||
pub status: String,
|
||||
}
|
||||
Reference in New Issue
Block a user