use serde::{Deserialize, Serialize};
use crate::v1::models::OllamaModel;
#[derive(Serialize, Debug, Clone)]
pub struct OllamaCompletionRequest {
pub model: OllamaModel,
pub prompt: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_tokens: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_p: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_k: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub presence_penalty: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub frequency_penalty: Option<f32>,
}
impl OllamaCompletionRequest {
pub fn to_stream_request(&self) -> OllamaCompletionStreamRequest {
OllamaCompletionStreamRequest {
model: self.model.clone(),
prompt: self.prompt.clone(),
max_tokens: self.max_tokens,
temperature: self.temperature,
top_p: self.top_p,
top_k: self.top_k,
stream: Some(true),
presence_penalty: self.presence_penalty,
frequency_penalty: self.frequency_penalty,
}
}
}
#[derive(Serialize, Debug, Clone)]
pub struct OllamaCompletionStreamRequest {
pub model: OllamaModel,
pub prompt: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_tokens: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_p: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_k: Option<u32>,
pub stream: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub presence_penalty: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub frequency_penalty: Option<f32>,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct OllamaCompletionResponse {
pub model: String,
pub created_at: String,
pub response: String,
pub done: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub total_duration: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub load_duration: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub sample_count: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub sample_duration: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub prompt_eval_count: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub prompt_eval_duration: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub eval_count: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub eval_duration: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub context: Option<Vec<u32>>,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct CompletionStreamChoice {
pub text: String,
pub index: u32,
pub finish_reason: String,
}
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Usage {
pub prompt_tokens: u32,
pub completion_tokens: u32,
pub total_tokens: u32,
}