ML-API/src/models.rs

81 lines
2.7 KiB
Rust

use reqwest::Client;
use tracing::{debug, instrument, warn, info};
use std::fmt;
use serde::{Deserialize, Serialize};
///
#[derive(Debug)]
pub struct ApiSessionConfig {
// integration config
pub target_url : String,
pub model_name : String,
pub request_timeout : usize,
pub client : Client,
}
impl fmt::Display for ApiSessionConfig {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
let param_name_width = 25;
let param_name_val = (self.model_name.as_bytes().iter().count() + 4).max(self.target_url.as_bytes().iter().count() + 4);
let param_name_measure = 10;
// header
writeln!(f, "+{:-<param_name_width$}-+-{:-<param_name_val$}+{:-<param_name_measure$}-+",
"", "", "")?;
writeln!(f, "| {: <param_name_width$}| {: <param_name_val$}| {: <param_name_measure$}|",
"Param", "Value", "Measure")?;
writeln!(f, "+{:-<param_name_width$}-+-{:-<param_name_val$}+{:-<param_name_measure$}-+",
"", "", "")?;
// data
writeln!(f, "| {: <param_name_width$}| {: <param_name_val$}| {: <param_name_measure$}|",
"Target model url", self.target_url, "-")?;
writeln!(f, "| {: <param_name_width$}| {: <param_name_val$}| {: <param_name_measure$}|",
"Tagert model name", self.model_name, "-")?;
writeln!(f, "| {: <param_name_width$}| {: <param_name_val$}| {: <param_name_measure$}|",
"API response wait-time", self.request_timeout, "secs")?;
// footer
writeln!(f, "+{:-<param_name_width$}-+-{:-<param_name_val$}+{:-<param_name_measure$}-+",
"", "", "")?;
Ok(())
}
}
impl ApiSessionConfig {
#[instrument("app-config")]
pub fn from_env() -> anyhow::Result<Self> {
let config = Self {
target_url : std::env::var("ML_TARGET_URL")?,
model_name : std::env::var("ML_MODEL_NAME")?,
request_timeout : std::env::var("ML_REQUEST_TIMEOUT")?.parse().unwrap_or_else(|_| {
warn!("invalid port was given, setting up default one ...");
15
}),
client : Client::new(),
};
debug!("{:?}", config);
info!("app is configurated. config :\n{}", config);
Ok(config)
}
}
#[derive(Debug, Serialize)]
pub struct OutputModel {
model : String,
prompt : String,
stream : bool,
}
impl OutputModel {
pub fn build(model_name: &str, prompt: String) -> Self {
Self {
model : model_name.to_owned(),
prompt : prompt,
stream : false,
}
}
}
#[derive(Debug, Serialize, Deserialize)]
pub struct ModelResponse {
pub response: String,
}