148 lines
3.6 KiB
Rust
148 lines
3.6 KiB
Rust
//! The llama inference server's chat-completions types and the request builder.
|
|
|
|
pub mod assemble;
|
|
pub mod info;
|
|
pub mod request;
|
|
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub enum ChatMessage {
|
|
System {
|
|
content: String,
|
|
},
|
|
User {
|
|
content: String,
|
|
},
|
|
Assistant {
|
|
content: Option<String>,
|
|
reasoning_content: Option<String>,
|
|
tool_calls: Vec<proto::ToolCall>,
|
|
},
|
|
Tool {
|
|
tool_call_id: String,
|
|
content: String,
|
|
},
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub struct ToolSchema {
|
|
pub name: String,
|
|
pub description: String,
|
|
pub parameters: serde_json::Value,
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub struct ChatRequest {
|
|
pub slot: u32,
|
|
pub messages: Vec<ChatMessage>,
|
|
pub tools: Vec<ToolSchema>,
|
|
pub thinking: bool,
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub enum ChatEvent {
|
|
Queued {
|
|
ahead: usize,
|
|
},
|
|
Waiting {
|
|
slot_busy: bool,
|
|
},
|
|
Progress {
|
|
total: u64,
|
|
cache: u64,
|
|
processed: u64,
|
|
},
|
|
Reasoning(String),
|
|
Content(String),
|
|
ToolCallDelta {
|
|
index: u32,
|
|
id: Option<String>,
|
|
name: Option<String>,
|
|
arguments: String,
|
|
},
|
|
ThinkingCapped {
|
|
tokens: u64,
|
|
},
|
|
Retrying {
|
|
attempt: u32,
|
|
after_ms: u64,
|
|
error: String,
|
|
},
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum FinishReason {
|
|
Stop,
|
|
ToolCalls,
|
|
Length,
|
|
}
|
|
|
|
/// The three counts loopd uses. The server sends more fields beside them; they are ignored.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Deserialize)]
|
|
pub struct Timings {
|
|
pub cache_n: u64,
|
|
pub prompt_n: u64,
|
|
pub predicted_n: u64,
|
|
}
|
|
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub struct Completion {
|
|
pub id: String,
|
|
pub content: Option<String>,
|
|
pub reasoning_content: Option<String>,
|
|
pub tool_calls: Vec<proto::ToolCall>,
|
|
pub finish_reason: FinishReason,
|
|
pub timings: Timings,
|
|
pub reasoning_tokens: u64,
|
|
pub thinking_capped: bool,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub enum InferError {
|
|
Busy,
|
|
Connect(std::io::Error),
|
|
WaitTimeout,
|
|
LoadTimeout,
|
|
Stalled,
|
|
StreamClosedEarly,
|
|
ThinkingOverrun,
|
|
Http { status: u16, body: String },
|
|
Protocol(String),
|
|
}
|
|
|
|
impl std::fmt::Display for InferError {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
|
|
match self {
|
|
InferError::Busy => write!(f, "the slot is busy"),
|
|
InferError::Connect(err) => {
|
|
write!(f, "could not connect to the inference server: {err}")
|
|
}
|
|
InferError::WaitTimeout => write!(f, "timed out waiting for the slot to free up"),
|
|
InferError::LoadTimeout => write!(f, "timed out waiting for the model to load"),
|
|
InferError::Stalled => write!(f, "the stream stalled"),
|
|
InferError::StreamClosedEarly => {
|
|
write!(f, "the stream closed before the response was complete")
|
|
}
|
|
InferError::ThinkingOverrun => write!(f, "thinking ran over its token budget"),
|
|
InferError::Http { status, body } => {
|
|
write!(f, "the server responded with {status}: {body}")
|
|
}
|
|
InferError::Protocol(what) => write!(f, "the server spoke the protocol wrong: {what}"),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl std::error::Error for InferError {}
|
|
|
|
pub struct Client {
|
|
pub(crate) cfg: crate::config::Config,
|
|
}
|
|
|
|
impl Client {
|
|
pub fn new(cfg: crate::config::Config) -> Self {
|
|
Self { cfg }
|
|
}
|
|
pub fn config(&self) -> &crate::config::Config {
|
|
&self.cfg
|
|
}
|
|
}
|