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,
|
||||
}
|
||||
@@ -0,0 +1,179 @@
|
||||
//! Tests for the request body. Do not edit.
|
||||
//!
|
||||
//! Bodies are compared as JSON values, so key order is free. What is not free: which keys
|
||||
//! appear, and that the same input always gives the same bytes.
|
||||
|
||||
mod support;
|
||||
|
||||
use loopd::llama::request::build_body;
|
||||
use loopd::llama::{ChatMessage, ChatRequest, ToolSchema};
|
||||
use serde_json::{Value, json};
|
||||
use std::path::Path;
|
||||
|
||||
fn body(req: &ChatRequest) -> Value {
|
||||
let cfg = support::test_config(Path::new("/tmp/unused.sock"));
|
||||
serde_json::from_str(&build_body(&cfg, req).unwrap()).unwrap()
|
||||
}
|
||||
|
||||
fn user(text: &str) -> ChatMessage {
|
||||
ChatMessage::User {
|
||||
content: text.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn read_file_tool() -> ToolSchema {
|
||||
ToolSchema {
|
||||
name: "read_file".to_string(),
|
||||
description: "Read a text file.".to_string(),
|
||||
parameters: json!({"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_plain_request_carries_every_setting() {
|
||||
let req = ChatRequest {
|
||||
slot: 1,
|
||||
messages: vec![user("hi")],
|
||||
tools: vec![],
|
||||
thinking: true,
|
||||
};
|
||||
let want = json!({
|
||||
"model": "test-model",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"id_slot": 1,
|
||||
"cache_prompt": true,
|
||||
"stream": true,
|
||||
"return_progress": true,
|
||||
"timings_per_token": true,
|
||||
"reasoning_control": true,
|
||||
"max_tokens": 8192,
|
||||
"temperature": 0.6,
|
||||
"top_p": 0.95,
|
||||
"top_k": 20,
|
||||
"chat_template_kwargs": {"enable_thinking": true},
|
||||
});
|
||||
assert_eq!(body(&req), want, "no `tools` key when there are no tools");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn thinking_off_and_tools() {
|
||||
let req = ChatRequest {
|
||||
slot: 0,
|
||||
messages: vec![user("hi")],
|
||||
tools: vec![read_file_tool()],
|
||||
thinking: false,
|
||||
};
|
||||
let got = body(&req);
|
||||
assert_eq!(
|
||||
got["chat_template_kwargs"],
|
||||
json!({"enable_thinking": false})
|
||||
);
|
||||
let want_tools = json!([{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "read_file",
|
||||
"description": "Read a text file.",
|
||||
"parameters": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]},
|
||||
},
|
||||
}]);
|
||||
assert_eq!(got["tools"], want_tools);
|
||||
}
|
||||
|
||||
/// One message of every kind, and both shapes of assistant message. An example of each is
|
||||
/// here on purpose: they are rendered differently.
|
||||
#[test]
|
||||
fn every_kind_of_message() {
|
||||
let call = proto::ToolCall {
|
||||
id: "call_1".to_string(),
|
||||
name: "read_file".to_string(),
|
||||
arguments: r#"{"path":"/etc/hostname"}"#.to_string(),
|
||||
};
|
||||
let messages = vec![
|
||||
ChatMessage::System {
|
||||
content: "be careful".to_string(),
|
||||
},
|
||||
user("read it"),
|
||||
ChatMessage::Assistant {
|
||||
content: None,
|
||||
reasoning_content: Some("I should read the file.".to_string()),
|
||||
tool_calls: vec![call],
|
||||
},
|
||||
ChatMessage::Tool {
|
||||
tool_call_id: "call_1".to_string(),
|
||||
content: "straylight\n".to_string(),
|
||||
},
|
||||
ChatMessage::Assistant {
|
||||
content: Some("It says straylight.".to_string()),
|
||||
reasoning_content: None,
|
||||
tool_calls: vec![],
|
||||
},
|
||||
];
|
||||
let req = ChatRequest {
|
||||
slot: 0,
|
||||
messages,
|
||||
tools: vec![read_file_tool()],
|
||||
thinking: true,
|
||||
};
|
||||
let want = json!([
|
||||
{"role": "system", "content": "be careful"},
|
||||
{"role": "user", "content": "read it"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"reasoning_content": "I should read the file.",
|
||||
"tool_calls": [{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "read_file", "arguments": "{\"path\":\"/etc/hostname\"}"},
|
||||
}],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "straylight\n"},
|
||||
{"role": "assistant", "content": "It says straylight."},
|
||||
]);
|
||||
assert_eq!(body(&req)["messages"], want);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn values_come_from_the_config() {
|
||||
let mut cfg = support::test_config(Path::new("/tmp/unused.sock"));
|
||||
cfg.infer.model = "another-model".to_string();
|
||||
cfg.limits.max_tokens = 99;
|
||||
cfg.sampling.temperature = 0.25;
|
||||
cfg.sampling.top_p = 0.5;
|
||||
cfg.sampling.top_k = 7;
|
||||
let req = ChatRequest {
|
||||
slot: 3,
|
||||
messages: vec![user("hi")],
|
||||
tools: vec![],
|
||||
thinking: false,
|
||||
};
|
||||
let got: Value = serde_json::from_str(&build_body(&cfg, &req).unwrap()).unwrap();
|
||||
assert_eq!(got["model"], "another-model");
|
||||
assert_eq!(got["max_tokens"], 99);
|
||||
assert_eq!(got["temperature"], 0.25);
|
||||
assert_eq!(got["top_p"], 0.5);
|
||||
assert_eq!(got["top_k"], 7);
|
||||
assert_eq!(got["id_slot"], 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_same_input_gives_the_same_bytes() {
|
||||
let cfg = support::test_config(Path::new("/tmp/unused.sock"));
|
||||
let req = ChatRequest {
|
||||
slot: 0,
|
||||
messages: vec![user(
|
||||
"text with \"quotes\", a \\ backslash, a\nnewline and caf\u{e9} \u{1f4e6}",
|
||||
)],
|
||||
tools: vec![read_file_tool()],
|
||||
thinking: true,
|
||||
};
|
||||
let first = build_body(&cfg, &req).unwrap();
|
||||
for _ in 0..20 {
|
||||
assert_eq!(build_body(&cfg, &req).unwrap(), first);
|
||||
}
|
||||
let back: Value = serde_json::from_str(&first).unwrap();
|
||||
assert_eq!(
|
||||
back["messages"][0]["content"],
|
||||
"text with \"quotes\", a \\ backslash, a\nnewline and caf\u{e9} \u{1f4e6}"
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user