Adds tags and ps methods

This commit is contained in:
2025-12-23 22:04:38 +00:00
parent 4b96a0b9dc
commit 8a0c10c6fc
14 changed files with 2042 additions and 3 deletions

1
.gitignore vendored
View File

@@ -1 +1,2 @@
/target
.env

1715
Cargo.lock generated Normal file

File diff suppressed because it is too large Load Diff

View File

@@ -4,3 +4,14 @@ version = "0.1.0"
edition = "2024"
[dependencies]
reqwest = { version = "0.12.28", features = ["json", "stream"] }
serde = { version = "1.0.228", features = ["derive"] }
serde_json = "1.0.146"
tokio-util = "0.7.17"
tracing = "0.1.44"
futures-util = "0.3.31"
[dev-dependencies]
dotenvy = "0.15.7"
tokio = { version = "1.48.0", features = ["full"] }
tracing-subscriber = "0.3.22"

16
examples/generate.rs Normal file
View File

@@ -0,0 +1,16 @@
use std::{env, error::Error};
use ollama_rs::{OllamaClient, types::generate::GenerateRequest};
#[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 request = GenerateRequest::builder("dolphin3:8b")
.prompt("Why is the sky blue?")
.build();
let response = ollama_client.generate(request).await?;
println!("{:?}", response);
Ok(())
}

16
examples/ps.rs Normal file
View File

@@ -0,0 +1,16 @@
use std::{env, error::Error};
use ollama_rs::OllamaClient;
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
tracing_subscriber::fmt().init();
let _ = dotenvy::dotenv();
let server_address = env::var("OLLAMA_SERVER")?;
let ollama_client = OllamaClient::new(server_address);
let models = ollama_client.list_runnning_models().await?;
for model in models {
println!("{:?}", model);
}
Ok(())
}

16
examples/tags.rs Normal file
View File

@@ -0,0 +1,16 @@
use std::{env, error::Error};
use ollama_rs::OllamaClient;
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
tracing_subscriber::fmt().init();
let _ = dotenvy::dotenv();
let server_address = env::var("OLLAMA_SERVER")?;
let ollama_client = OllamaClient::new(server_address);
let models = ollama_client.tags().await?;
for model in models {
println!("{:?}", model);
}
Ok(())
}

32
src/error.rs Normal file
View File

@@ -0,0 +1,32 @@
use std::{error::Error, fmt::Display};
pub type OllamaResult<T> = Result<T, OllamaError>;
#[derive(Debug)]
pub enum OllamaError {
NetworkError(reqwest::Error),
ResponseParseError(serde_json::Error),
}
impl Error for OllamaError {}
impl Display for OllamaError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
OllamaError::NetworkError(e) => writeln!(f, "Network error: {}", e),
OllamaError::ResponseParseError(e) => writeln!(f, "ResponseParseError error: {}", e),
}
}
}
impl From<reqwest::Error> for OllamaError {
fn from(error: reqwest::Error) -> Self {
Self::NetworkError(error)
}
}
impl From<serde_json::Error> for OllamaError {
fn from(error: serde_json::Error) -> Self {
Self::ResponseParseError(error)
}
}

102
src/lib.rs Normal file
View File

@@ -0,0 +1,102 @@
use futures_util::{StreamExt};
use serde_json::Value;
use tokio_util::io::StreamReader;
use tracing::info;
use crate::{
error::OllamaResult,
types::{
generate::{GenerateRequest, GenerateResponse},
ps::RunningModel,
tags::Model,
},
};
pub mod error;
pub mod types;
pub struct OllamaClient {
server_address: String,
}
impl OllamaClient {
pub fn new<S: AsRef<str>>(server_address: S) -> Self {
Self {
server_address: server_address.as_ref().to_string(),
}
}
/// Fetch a list of models and their details
pub async fn tags(&self) -> OllamaResult<Vec<Model>> {
let request_address = format!("{}/api/tags", self.server_address);
info!("List models: {}", request_address);
let mut response: Value = reqwest::get(request_address)
.await?
.error_for_status()?
.json()
.await?;
let Some(response) = response.as_object_mut() else {
return Ok(vec![]);
};
let Some(models) = response.remove("models") else {
return Ok(vec![]);
};
let models = serde_json::from_value(models)?;
Ok(models)
}
/// Retrieve a list of models that are currently running
pub async fn list_runnning_models(&self) -> OllamaResult<Vec<RunningModel>> {
let request_address = format!("{}/api/ps", self.server_address);
info!("List models: {}", request_address);
let mut response: Value = reqwest::get(request_address)
.await?
.error_for_status()?
.json()
.await?;
let Some(response) = response.as_object_mut() else {
return Ok(vec![]);
};
let Some(models) = response.remove("models") else {
return Ok(vec![]);
};
let models = serde_json::from_value(models)?;
Ok(models)
}
/// Generates a response for the provided prompt
pub async fn generate(&self, request: GenerateRequest) -> OllamaResult<()> {
let request_address = format!("{}/api/generate", self.server_address);
let client = reqwest::Client::new();
let response = client
.post(request_address)
.json(&request)
.send()
.await?
.error_for_status()?;
let stream = response.bytes_stream().;
// let reader = BufReader::new(stream);
let reader = StreamReader(stream);
while reader
while let Some(item) = stream.next().await {
let item = item?;
println!("Chunk: {:?}", item?);
}
// let stream_reader = tokio_util::io::StreamReader::new(stream);
// let reder = BufReader::new(stream);
// let full_response = response.text().await?;
// let parts = full_response
// .lines()
// .map(|line| serde_json::from_str::<GenerateResponse>(line).unwrap())
// .collect::<Vec<_>>();
// println!("{:#?}", parts);
Ok(())
}
}

View File

@@ -1,3 +0,0 @@
fn main() {
println!("Hello, world!");
}

10
src/types/common.rs Normal file
View File

@@ -0,0 +1,10 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Serialize, Deserialize)]
pub struct ModelDetails {
pub format: String,
pub family: String,
pub families: Vec<String>,
pub parameter_size: String,
pub quantization_level: String,
}

91
src/types/generate.rs Normal file
View File

@@ -0,0 +1,91 @@
use serde::{Deserialize, Serialize};
#[derive(Debug, Serialize, Deserialize)]
pub struct GenerateRequest {
/// Model name
pub model: String,
#[serde(skip_serializing_if = "Option::is_none")]
/// Text for the model to generate a response from
pub prompt: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
/// Used for fill-in-the-middle models, text that appears after the user prompt and before the
/// model response
pub suffix: Option<String>,
/// System prompt for the model to generate a response from
#[serde(skip_serializing_if = "Option::is_none")]
pub system: Option<String>,
}
impl GenerateRequest {
pub fn builder<M: Into<String>>(model: M) -> GenerateRequestBuilder {
GenerateRequestBuilder::new(model)
}
}
pub struct GenerateRequestBuilder {
generate_request: GenerateRequest,
}
impl GenerateRequestBuilder {
fn new<M: Into<String>>(model: M) -> Self {
Self {
generate_request: GenerateRequest {
model: model.into(),
prompt: None,
suffix: None,
system: None,
},
}
}
pub fn prompt<P: Into<String>>(mut self, prompt: P) -> Self {
self.generate_request.prompt = Some(prompt.into());
self
}
pub fn build(self) -> GenerateRequest {
self.generate_request
}
}
#[derive(Debug, Serialize, Deserialize)]
pub struct GenerateResponse {
/// Model name
pub model: String,
/// ISO 8601 timestamp of response creation
pub created_at: String,
/// The model's generated text response
pub response: String,
/// The model's generated thinking output
pub thinking: Option<String>,
/// Indicates whether generation has finished
pub done: bool,
/// Reason the generation stopped
pub done_reason: Option<String>,
/// Time spent generating the response in nanoseconds
pub total_duration: Option<usize>,
/// Time spent loading the model in nanoseconds
pub load_duration: Option<usize>,
/// Number of input tokens in the prompt
pub prompt_eval_count: Option<usize>,
/// Time spent evaluating the prompt in nanoseconds
pub prompt_eval_duration: Option<usize>,
/// Number of output tokens generated in the response
pub eval_count: Option<usize>,
/// Time spent generating tokens in nanoseconds
pub eval_duration: Option<usize>,
}

4
src/types/mod.rs Normal file
View File

@@ -0,0 +1,4 @@
pub mod common;
pub mod generate;
pub mod ps;
pub mod tags;

15
src/types/ps.rs Normal file
View File

@@ -0,0 +1,15 @@
use serde::{Deserialize, Serialize};
use crate::types::common::ModelDetails;
#[derive(Debug, Serialize, Deserialize)]
pub struct RunningModel {
pub name: String,
pub model: String,
pub size: usize,
pub digest: String,
pub details: ModelDetails,
pub expires_at: String,
pub size_vram: usize,
pub context_length: u32,
}

13
src/types/tags.rs Normal file
View File

@@ -0,0 +1,13 @@
use serde::{Deserialize, Serialize};
use crate::types::common::ModelDetails;
#[derive(Debug, Serialize, Deserialize)]
pub struct Model {
pub name: String,
pub model: String,
pub modified_at: String,
pub size: usize,
pub digest: String,
pub details: ModelDetails,
}