Use clap for command-line arguments in pull example
This commit is contained in:
@@ -1,21 +1,37 @@
|
||||
use std::{env, error::Error, io::Write};
|
||||
use std::{error::Error, io::Write};
|
||||
|
||||
use clap::Parser;
|
||||
use futures_util::StreamExt;
|
||||
use ollama_rs::{OllamaClient, types::pull::PullRequest};
|
||||
|
||||
const MODEL: &str = "functiongemma";
|
||||
#[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 server_address = env::var("OLLAMA_SERVER")?;
|
||||
let ollama_client = OllamaClient::new(server_address);
|
||||
let args = Args::parse();
|
||||
let ollama_client = OllamaClient::new(args.server_address);
|
||||
|
||||
let request = PullRequest::builder(MODEL).stream(true).build();
|
||||
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);
|
||||
println!("{}", response.status);
|
||||
std::io::stdout().flush()?;
|
||||
}
|
||||
Ok(())
|
||||
|
||||
Reference in New Issue
Block a user