use serde::{Deserialize, Serialize};
use std::error::Error;
use reqwest;
#[derive(Serialize, Deserialize, Debug)]
pub struct GenerateResponse {
pub model: String,
pub created_at: String,
pub done: bool,
pub done_reason: String,
pub context: Vec<u64>,
pub total_duration: u64,
pub load_duration: u64,
pub prompt_eval_duration: u64,
pub eval_count: u32,
pub eval_duration: u64,
pub response: String,
}
#[derive(Deserialize, Debug)]
pub struct HealthResponse {
pub status: String,
pub timestamp: String,
}
pub struct ApiMyLlama {
server_ip: String,
server_port: u16,
}
impl ApiMyLlama {
pub fn new(server_ip: String, server_port: u16) -> Self {
ApiMyLlama { server_ip, server_port }
}
pub async fn generate(
&self,
api_key: &str,
prompt: &str,
model: &str,
stream: bool,
) -> Result<GenerateResponse, Box<dyn Error>> {
let url = format!("http://{}:{}/generate", self.server_ip, self.server_port);
let payload = serde_json::json!({
"apikey": api_key,
"prompt": prompt,
"model": model,
"stream": stream,
});
let client = reqwest::Client::new();
let response = client
.post(&url)
.json(&payload)
.send()
.await?;
let generate_response: GenerateResponse = response.json().await?;
Ok(generate_response)
}
pub async fn get_health(&self, api_key: &str) -> Result<HealthResponse, Box<dyn Error>> {
let url = format!("http://{}:{}/health", self.server_ip, self.server_port);
let client = reqwest::Client::new();
let response = client
.get(&url)
.query(&[("apikey", api_key)])
.send()
.await?;
let health_response: HealthResponse = response.json().await?;
Ok(health_response)
}
}