Add the llama client types and the chat request builder
Implemented-By: Laguna S 2.1 (OpenCode)
This commit is contained in:
@@ -2,4 +2,5 @@
|
||||
|
||||
pub mod config;
|
||||
pub mod http;
|
||||
pub mod llama;
|
||||
pub mod sse;
|
||||
|
||||
@@ -0,0 +1,145 @@
|
||||
//! The llama inference server's chat-completions types and the request builder.
|
||||
|
||||
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
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,177 @@
|
||||
//! Building the chat request body the inference server expects.
|
||||
|
||||
use serde::Serialize;
|
||||
|
||||
use super::{ChatMessage, ToolSchema};
|
||||
|
||||
/// The request body as JSON text. The same inputs always give the same bytes.
|
||||
pub fn build_body(
|
||||
cfg: &crate::config::Config,
|
||||
req: &super::ChatRequest,
|
||||
) -> Result<String, serde_json::Error> {
|
||||
let messages = serialize_messages(&req.messages)?;
|
||||
let tools = req.tools.iter().map(serialized_tool).collect::<Vec<_>>();
|
||||
let body = RequestBody {
|
||||
model: cfg.infer.model.clone(),
|
||||
messages,
|
||||
tools,
|
||||
id_slot: req.slot,
|
||||
cache_prompt: true,
|
||||
stream: true,
|
||||
return_progress: true,
|
||||
timings_per_token: true,
|
||||
reasoning_control: true,
|
||||
max_tokens: cfg.limits.max_tokens,
|
||||
temperature: cfg.sampling.temperature,
|
||||
top_p: cfg.sampling.top_p,
|
||||
top_k: cfg.sampling.top_k,
|
||||
chat_template_kwargs: ChatTemplateKwargs {
|
||||
enable_thinking: req.thinking,
|
||||
},
|
||||
};
|
||||
serde_json::to_string(&body)
|
||||
}
|
||||
|
||||
fn serialize_messages(
|
||||
messages: &[ChatMessage],
|
||||
) -> Result<Vec<serde_json::Value>, serde_json::Error> {
|
||||
let mut out = Vec::with_capacity(messages.len());
|
||||
for message in messages {
|
||||
out.push(serialize_message(message)?);
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
fn serialize_message(message: &ChatMessage) -> Result<serde_json::Value, serde_json::Error> {
|
||||
match message {
|
||||
ChatMessage::System { content } => serde_json::to_value(SerializedSystem {
|
||||
role: "system",
|
||||
content: content.clone(),
|
||||
}),
|
||||
ChatMessage::User { content } => serde_json::to_value(SerializedUser {
|
||||
role: "user",
|
||||
content: content.clone(),
|
||||
}),
|
||||
ChatMessage::Tool {
|
||||
tool_call_id,
|
||||
content,
|
||||
} => serde_json::to_value(SerializedToolMessage {
|
||||
role: "tool",
|
||||
tool_call_id: tool_call_id.clone(),
|
||||
content: content.clone(),
|
||||
}),
|
||||
ChatMessage::Assistant {
|
||||
content,
|
||||
reasoning_content,
|
||||
tool_calls,
|
||||
} => serde_json::to_value(SerializedAssistant {
|
||||
role: "assistant",
|
||||
content: content.clone().unwrap_or_default(),
|
||||
reasoning_content: reasoning_content.clone(),
|
||||
tool_calls: tool_calls.iter().map(serialized_call).collect(),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn serialized_call(call: &proto::ToolCall) -> SerializedToolCall {
|
||||
SerializedToolCall {
|
||||
id: call.id.clone(),
|
||||
kind: "function",
|
||||
function: CallFunction {
|
||||
name: call.name.clone(),
|
||||
arguments: call.arguments.clone(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn serialized_tool(tool: &ToolSchema) -> SerializedTool {
|
||||
SerializedTool {
|
||||
kind: "function",
|
||||
function: ToolFunction {
|
||||
name: tool.name.clone(),
|
||||
description: tool.description.clone(),
|
||||
parameters: tool.parameters.clone(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct RequestBody {
|
||||
model: String,
|
||||
messages: Vec<serde_json::Value>,
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
tools: Vec<SerializedTool>,
|
||||
id_slot: u32,
|
||||
cache_prompt: bool,
|
||||
stream: bool,
|
||||
return_progress: bool,
|
||||
timings_per_token: bool,
|
||||
reasoning_control: bool,
|
||||
max_tokens: u64,
|
||||
temperature: f64,
|
||||
top_p: f64,
|
||||
top_k: u32,
|
||||
chat_template_kwargs: ChatTemplateKwargs,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ChatTemplateKwargs {
|
||||
enable_thinking: bool,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct SerializedSystem {
|
||||
role: &'static str,
|
||||
content: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct SerializedUser {
|
||||
role: &'static str,
|
||||
content: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct SerializedToolMessage {
|
||||
role: &'static str,
|
||||
tool_call_id: String,
|
||||
content: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct SerializedAssistant {
|
||||
role: &'static str,
|
||||
content: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
reasoning_content: Option<String>,
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
tool_calls: Vec<SerializedToolCall>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct SerializedToolCall {
|
||||
id: String,
|
||||
#[serde(rename = "type")]
|
||||
kind: &'static str,
|
||||
function: CallFunction,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct CallFunction {
|
||||
name: String,
|
||||
arguments: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct SerializedTool {
|
||||
#[serde(rename = "type")]
|
||||
kind: &'static str,
|
||||
function: ToolFunction,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ToolFunction {
|
||||
name: String,
|
||||
description: String,
|
||||
parameters: serde_json::Value,
|
||||
}
|
||||
Reference in New Issue
Block a user