39 lines
993 B
Rust
39 lines
993 B
Rust
use std::{error::Error, io::Write};
|
|
|
|
use clap::Parser;
|
|
use futures_util::StreamExt;
|
|
use ollama_rs::{OllamaClient, types::pull::PullRequest};
|
|
|
|
#[derive(Parser, Debug)]
|
|
#[command(author, version, about, long_about = None)]
|
|
struct Args {
|
|
/// The model to pull
|
|
#[arg(short, long)]
|
|
model: String,
|
|
|
|
/// The Ollama server address
|
|
#[arg(
|
|
short,
|
|
long,
|
|
env = "OLLAMA_SERVER",
|
|
default_value = "http://localhost:11434"
|
|
)]
|
|
server_address: String,
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<(), Box<dyn Error>> {
|
|
let _ = dotenvy::dotenv();
|
|
let args = Args::parse();
|
|
let ollama_client = OllamaClient::new(args.server_address);
|
|
|
|
let request = PullRequest::builder(&args.model).stream(true).build();
|
|
let mut stream = ollama_client.pull(request);
|
|
while let Some(response) = stream.next().await {
|
|
let response = response?;
|
|
println!("{}", response.status);
|
|
std::io::stdout().flush()?;
|
|
}
|
|
Ok(())
|
|
}
|