765 lines
25 KiB
Rust
765 lines
25 KiB
Rust
//! Provider-neutral exact-endpoint LLM transport and bounded wire envelope.
|
|
|
|
#![allow(clippy::missing_errors_doc)] // Public methods share the exhaustive typed LlmError boundary.
|
|
|
|
use crate::config::LlmConnection;
|
|
use crate::types::{
|
|
BoundedText, BoundedVec, MAX_BODY_BYTES, MAX_CONVERSATION_MESSAGES, MAX_IDENTIFIER_BYTES,
|
|
MAX_MESSAGE_BYTES, MAX_TOOL_CALLS, MessageRole, ProposedToolCall,
|
|
};
|
|
use libremetaverse_types::compat::CancellationToken;
|
|
use serde::Deserialize;
|
|
use serde_json::{Map, Value, json};
|
|
use std::collections::{BTreeMap, BTreeSet};
|
|
use std::error::Error;
|
|
use std::fmt;
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicU64, Ordering};
|
|
use std::time::{Duration, Instant};
|
|
use tokio::sync::Semaphore;
|
|
|
|
const MAX_TOOLS: usize = 64;
|
|
const MAX_SCHEMA_PROPERTIES: usize = 128;
|
|
const MAX_SCHEMA_DEPTH: usize = 16;
|
|
const MAX_SCHEMA_NODES: usize = 1_024;
|
|
|
|
#[derive(Clone, Debug, Eq, PartialEq)]
|
|
pub enum ContentPart {
|
|
Text(BoundedText<MAX_MESSAGE_BYTES>),
|
|
Image {
|
|
url: BoundedText<MAX_BODY_BYTES>,
|
|
detail: ImageDetail,
|
|
},
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
|
pub enum ImageDetail {
|
|
Auto,
|
|
Low,
|
|
High,
|
|
}
|
|
|
|
#[derive(Clone, Debug, Eq, PartialEq)]
|
|
pub struct CompletionMessage {
|
|
pub role: MessageRole,
|
|
pub content: BoundedVec<ContentPart, 16>,
|
|
pub tool_call_id: Option<BoundedText<MAX_IDENTIFIER_BYTES>>,
|
|
pub proposed_calls: BoundedVec<ProposedToolCall, MAX_TOOL_CALLS>,
|
|
}
|
|
|
|
impl CompletionMessage {
|
|
pub fn text(role: MessageRole, text: impl Into<String>) -> Result<Self, LlmError> {
|
|
let mut content = BoundedVec::new();
|
|
content
|
|
.try_push(
|
|
"completion_message.content",
|
|
ContentPart::Text(BoundedText::new("completion_message.text", text)?),
|
|
)
|
|
.map_err(LlmError::Boundary)?;
|
|
Ok(Self {
|
|
role,
|
|
content,
|
|
tool_call_id: None,
|
|
proposed_calls: BoundedVec::new(),
|
|
})
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Debug, Eq, PartialEq)]
|
|
pub enum ToolSchema {
|
|
Object {
|
|
properties: BTreeMap<String, ToolSchema>,
|
|
required: BTreeSet<String>,
|
|
additional_properties: bool,
|
|
},
|
|
String,
|
|
Integer,
|
|
Number,
|
|
Boolean,
|
|
Array {
|
|
items: Box<ToolSchema>,
|
|
max_items: usize,
|
|
},
|
|
}
|
|
|
|
impl ToolSchema {
|
|
pub fn validate_schema(&self) -> Result<(), LlmError> {
|
|
let mut remaining_nodes = MAX_SCHEMA_NODES;
|
|
self.validate_schema_inner(0, &mut remaining_nodes)
|
|
}
|
|
|
|
fn validate_schema_inner(
|
|
&self,
|
|
depth: usize,
|
|
remaining_nodes: &mut usize,
|
|
) -> Result<(), LlmError> {
|
|
if depth > MAX_SCHEMA_DEPTH || *remaining_nodes == 0 {
|
|
return Err(LlmError::InvalidToolSchema);
|
|
}
|
|
*remaining_nodes -= 1;
|
|
match self {
|
|
Self::Object {
|
|
properties,
|
|
required,
|
|
..
|
|
} => {
|
|
if properties.len() > MAX_SCHEMA_PROPERTIES
|
|
|| required.len() > MAX_SCHEMA_PROPERTIES
|
|
|| required.iter().any(|name| !properties.contains_key(name))
|
|
|| properties
|
|
.keys()
|
|
.any(|name| name.is_empty() || name.len() > MAX_IDENTIFIER_BYTES)
|
|
{
|
|
return Err(LlmError::InvalidToolSchema);
|
|
}
|
|
for schema in properties.values() {
|
|
schema.validate_schema_inner(depth + 1, remaining_nodes)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
Self::Array { items, max_items } => {
|
|
if *max_items == 0 || *max_items > 256 {
|
|
return Err(LlmError::InvalidToolSchema);
|
|
}
|
|
items.validate_schema_inner(depth + 1, remaining_nodes)
|
|
}
|
|
_ => Ok(()),
|
|
}
|
|
}
|
|
|
|
pub fn validate_value(&self, value: &Value) -> Result<(), LlmError> {
|
|
match self {
|
|
Self::Object {
|
|
properties,
|
|
required,
|
|
additional_properties,
|
|
} => {
|
|
let object = value.as_object().ok_or(LlmError::InvalidToolArguments)?;
|
|
if required.iter().any(|name| !object.contains_key(name))
|
|
|| (!additional_properties
|
|
&& object.keys().any(|name| !properties.contains_key(name)))
|
|
{
|
|
return Err(LlmError::InvalidToolArguments);
|
|
}
|
|
for (name, value) in object {
|
|
if let Some(schema) = properties.get(name) {
|
|
schema.validate_value(value)?;
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
Self::String if value.is_string() => Ok(()),
|
|
Self::Integer if value.as_i64().is_some() || value.as_u64().is_some() => Ok(()),
|
|
Self::Number if value.is_number() => Ok(()),
|
|
Self::Boolean if value.is_boolean() => Ok(()),
|
|
Self::Array { items, max_items } => {
|
|
let values = value.as_array().ok_or(LlmError::InvalidToolArguments)?;
|
|
if values.len() > *max_items {
|
|
return Err(LlmError::InvalidToolArguments);
|
|
}
|
|
for value in values {
|
|
items.validate_value(value)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
_ => Err(LlmError::InvalidToolArguments),
|
|
}
|
|
}
|
|
|
|
fn wire_value(&self) -> Value {
|
|
match self {
|
|
Self::Object {
|
|
properties,
|
|
required,
|
|
additional_properties,
|
|
} => {
|
|
let properties = properties
|
|
.iter()
|
|
.map(|(name, schema)| (name.clone(), schema.wire_value()))
|
|
.collect::<Map<_, _>>();
|
|
json!({
|
|
"type": "object",
|
|
"properties": properties,
|
|
"required": required,
|
|
"additionalProperties": additional_properties
|
|
})
|
|
}
|
|
Self::String => json!({"type":"string"}),
|
|
Self::Integer => json!({"type":"integer"}),
|
|
Self::Number => json!({"type":"number"}),
|
|
Self::Boolean => json!({"type":"boolean"}),
|
|
Self::Array { items, max_items } => {
|
|
json!({"type":"array", "items":items.wire_value(), "maxItems":max_items})
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Debug, Eq, PartialEq)]
|
|
pub struct ToolDefinition {
|
|
pub name: BoundedText<MAX_IDENTIFIER_BYTES>,
|
|
pub description: BoundedText<MAX_MESSAGE_BYTES>,
|
|
pub schema: ToolSchema,
|
|
pub mutating: bool,
|
|
}
|
|
|
|
impl ToolDefinition {
|
|
pub fn validate(&self) -> Result<(), LlmError> {
|
|
if !self
|
|
.name
|
|
.chars()
|
|
.all(|character| character.is_ascii_alphanumeric() || matches!(character, '_' | '-'))
|
|
{
|
|
return Err(LlmError::InvalidToolSchema);
|
|
}
|
|
self.schema.validate_schema()
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Debug, Eq, PartialEq)]
|
|
pub struct Usage {
|
|
pub prompt_tokens: Option<u64>,
|
|
pub completion_tokens: Option<u64>,
|
|
pub total_tokens: Option<u64>,
|
|
}
|
|
|
|
#[derive(Clone, Debug, Eq, PartialEq)]
|
|
pub struct Completion {
|
|
pub request_id: u64,
|
|
pub correlation_id: BoundedText<MAX_IDENTIFIER_BYTES>,
|
|
pub message: CompletionMessage,
|
|
pub usage: Option<Usage>,
|
|
pub latency: Duration,
|
|
pub attempts: usize,
|
|
}
|
|
|
|
#[derive(Clone, Debug, Eq, PartialEq)]
|
|
pub struct LlmTransportLimits {
|
|
pub connect_timeout: Duration,
|
|
pub request_timeout: Duration,
|
|
pub read_idle_timeout: Duration,
|
|
pub pool_idle_timeout: Duration,
|
|
pub total_timeout: Duration,
|
|
pub max_prompt_bytes: usize,
|
|
pub max_response_bytes: usize,
|
|
pub max_concurrent_requests: usize,
|
|
pub max_retries: usize,
|
|
pub max_retry_delay: Duration,
|
|
}
|
|
|
|
impl Default for LlmTransportLimits {
|
|
fn default() -> Self {
|
|
Self {
|
|
connect_timeout: Duration::from_secs(10),
|
|
request_timeout: Duration::from_mins(1),
|
|
read_idle_timeout: Duration::from_secs(15),
|
|
pool_idle_timeout: Duration::from_secs(30),
|
|
total_timeout: Duration::from_secs(90),
|
|
max_prompt_bytes: 1024 * 1024,
|
|
max_response_bytes: 1024 * 1024,
|
|
max_concurrent_requests: 8,
|
|
max_retries: 2,
|
|
max_retry_delay: Duration::from_secs(5),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl LlmTransportLimits {
|
|
pub fn validate(&self) -> Result<(), LlmError> {
|
|
if self.connect_timeout.is_zero()
|
|
|| self.request_timeout.is_zero()
|
|
|| self.read_idle_timeout.is_zero()
|
|
|| self.pool_idle_timeout.is_zero()
|
|
|| self.total_timeout.is_zero()
|
|
|| self.max_prompt_bytes == 0
|
|
|| self.max_prompt_bytes > MAX_BODY_BYTES
|
|
|| self.max_response_bytes == 0
|
|
|| self.max_response_bytes > MAX_BODY_BYTES
|
|
|| self.max_concurrent_requests == 0
|
|
|| self.max_concurrent_requests > 256
|
|
|| self.max_retries > 8
|
|
|| self.max_retry_delay > Duration::from_mins(1)
|
|
{
|
|
return Err(LlmError::UnsafeLimits);
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Debug, Eq, PartialEq)]
|
|
pub enum LlmError {
|
|
Boundary(crate::types::BoundaryError),
|
|
UnsafeLimits,
|
|
InvalidToolSchema,
|
|
InvalidToolArguments,
|
|
PromptTooLarge,
|
|
ResponseTooLarge,
|
|
Cancelled,
|
|
Timeout,
|
|
Transport,
|
|
RedirectRefused,
|
|
HttpStatus(u16),
|
|
MalformedJson,
|
|
UnsupportedResponse,
|
|
DuplicateToolCallId,
|
|
}
|
|
|
|
impl fmt::Display for LlmError {
|
|
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
match self {
|
|
Self::Boundary(error) => write!(formatter, "LLM boundary rejected input: {error}"),
|
|
Self::UnsafeLimits => formatter.write_str("unsafe LLM transport limits"),
|
|
Self::InvalidToolSchema => formatter.write_str("invalid registered tool schema"),
|
|
Self::InvalidToolArguments => formatter.write_str("tool arguments do not match schema"),
|
|
Self::PromptTooLarge => formatter.write_str("LLM prompt exceeds its byte bound"),
|
|
Self::ResponseTooLarge => formatter.write_str("LLM response exceeds its byte bound"),
|
|
Self::Cancelled => formatter.write_str("LLM request cancelled"),
|
|
Self::Timeout => formatter.write_str("LLM request timed out"),
|
|
Self::Transport => formatter.write_str("LLM transport failed"),
|
|
Self::RedirectRefused => formatter.write_str("LLM endpoint redirect refused"),
|
|
Self::HttpStatus(status) => write!(formatter, "LLM endpoint returned HTTP {status}"),
|
|
Self::MalformedJson => formatter.write_str("LLM endpoint returned malformed JSON"),
|
|
Self::UnsupportedResponse => formatter.write_str("unsupported LLM response shape"),
|
|
Self::DuplicateToolCallId => {
|
|
formatter.write_str("LLM response repeated a tool-call ID")
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Error for LlmError {}
|
|
|
|
impl From<crate::types::BoundaryError> for LlmError {
|
|
fn from(value: crate::types::BoundaryError) -> Self {
|
|
Self::Boundary(value)
|
|
}
|
|
}
|
|
|
|
pub struct LlmClient {
|
|
connection: LlmConnection,
|
|
client: reqwest::Client,
|
|
limits: LlmTransportLimits,
|
|
slots: Arc<Semaphore>,
|
|
next_request_id: AtomicU64,
|
|
}
|
|
|
|
impl fmt::Debug for LlmClient {
|
|
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
formatter
|
|
.debug_struct("LlmClient")
|
|
.field("endpoint", &self.connection.endpoint_url)
|
|
.field("limits", &self.limits)
|
|
.field("available_slots", &self.slots.available_permits())
|
|
.finish_non_exhaustive()
|
|
}
|
|
}
|
|
|
|
impl LlmClient {
|
|
pub fn new(connection: LlmConnection, limits: LlmTransportLimits) -> Result<Self, LlmError> {
|
|
limits.validate()?;
|
|
let client = reqwest::Client::builder()
|
|
.redirect(reqwest::redirect::Policy::none())
|
|
.connect_timeout(limits.connect_timeout)
|
|
.timeout(limits.request_timeout)
|
|
.read_timeout(limits.read_idle_timeout)
|
|
.pool_idle_timeout(limits.pool_idle_timeout)
|
|
.pool_max_idle_per_host(limits.max_concurrent_requests)
|
|
.build()
|
|
.map_err(|_| LlmError::Transport)?;
|
|
Ok(Self {
|
|
connection,
|
|
client,
|
|
slots: Arc::new(Semaphore::new(limits.max_concurrent_requests)),
|
|
limits,
|
|
next_request_id: AtomicU64::new(1),
|
|
})
|
|
}
|
|
|
|
pub async fn complete(
|
|
&self,
|
|
messages: &[CompletionMessage],
|
|
tools: &[ToolDefinition],
|
|
cancellation: &CancellationToken,
|
|
) -> Result<Completion, LlmError> {
|
|
if messages.is_empty()
|
|
|| messages.len() > MAX_CONVERSATION_MESSAGES
|
|
|| tools.len() > MAX_TOOLS
|
|
{
|
|
return Err(LlmError::PromptTooLarge);
|
|
}
|
|
for tool in tools {
|
|
tool.validate()?;
|
|
}
|
|
let body = request_body(messages, tools)?;
|
|
if body.len() > self.limits.max_prompt_bytes {
|
|
return Err(LlmError::PromptTooLarge);
|
|
}
|
|
let request_id = self.next_request_id.fetch_add(1, Ordering::Relaxed);
|
|
let correlation_id = BoundedText::new(
|
|
"llm.correlation_id",
|
|
format!("agent-request-{request_id:016x}"),
|
|
)?;
|
|
let started = Instant::now();
|
|
let operation = async {
|
|
let permit = tokio::select! {
|
|
() = cancellation.cancelled() => return Err(LlmError::Cancelled),
|
|
permit = Arc::clone(&self.slots).acquire_owned() => permit.map_err(|_| LlmError::Cancelled)?,
|
|
};
|
|
let result = self
|
|
.complete_with_retries(request_id, correlation_id, body, cancellation, started)
|
|
.await;
|
|
drop(permit);
|
|
result
|
|
};
|
|
tokio::time::timeout(self.limits.total_timeout, operation)
|
|
.await
|
|
.map_err(|_| LlmError::Timeout)?
|
|
}
|
|
|
|
async fn complete_with_retries(
|
|
&self,
|
|
request_id: u64,
|
|
correlation_id: BoundedText<MAX_IDENTIFIER_BYTES>,
|
|
body: Vec<u8>,
|
|
cancellation: &CancellationToken,
|
|
started: Instant,
|
|
) -> Result<Completion, LlmError> {
|
|
for attempt in 0..=self.limits.max_retries {
|
|
match self
|
|
.complete_once(
|
|
request_id,
|
|
correlation_id.clone(),
|
|
&body,
|
|
cancellation,
|
|
started,
|
|
attempt + 1,
|
|
)
|
|
.await
|
|
{
|
|
Ok(completion) => return Ok(completion),
|
|
Err(AttemptError { error, retry_after })
|
|
if attempt < self.limits.max_retries && transient(&error) =>
|
|
{
|
|
let delay = retry_after
|
|
.unwrap_or_else(|| deterministic_backoff(request_id, attempt))
|
|
.min(self.limits.max_retry_delay);
|
|
tokio::select! {
|
|
() = cancellation.cancelled() => return Err(LlmError::Cancelled),
|
|
() = tokio::time::sleep(delay) => {}
|
|
}
|
|
}
|
|
Err(AttemptError { error, .. }) => return Err(error),
|
|
}
|
|
}
|
|
Err(LlmError::Transport)
|
|
}
|
|
|
|
async fn complete_once(
|
|
&self,
|
|
request_id: u64,
|
|
correlation_id: BoundedText<MAX_IDENTIFIER_BYTES>,
|
|
body: &[u8],
|
|
cancellation: &CancellationToken,
|
|
started: Instant,
|
|
attempts: usize,
|
|
) -> Result<Completion, AttemptError> {
|
|
let response = tokio::select! {
|
|
() = cancellation.cancelled() => return Err(AttemptError::new(LlmError::Cancelled)),
|
|
response = self.client
|
|
.post(self.connection.endpoint_url.expose_url())
|
|
.bearer_auth(self.connection.api_key.expose_secret())
|
|
.header("content-type", "application/json")
|
|
.header("accept", "application/json")
|
|
.header("x-correlation-id", correlation_id.as_str())
|
|
.body(body.to_vec())
|
|
.send() => response.map_err(|error| AttemptError::new(classify_reqwest(&error)))?,
|
|
};
|
|
let status = response.status();
|
|
if status.is_redirection() {
|
|
return Err(AttemptError::new(LlmError::RedirectRefused));
|
|
}
|
|
if !status.is_success() {
|
|
let retry_after = response
|
|
.headers()
|
|
.get("retry-after")
|
|
.and_then(|value| value.to_str().ok())
|
|
.and_then(|value| value.parse::<u64>().ok())
|
|
.map(Duration::from_secs);
|
|
return Err(AttemptError {
|
|
error: LlmError::HttpStatus(status.as_u16()),
|
|
retry_after,
|
|
});
|
|
}
|
|
let bytes = read_bounded(response, self.limits.max_response_bytes, cancellation).await?;
|
|
parse_completion(
|
|
request_id,
|
|
correlation_id,
|
|
&bytes,
|
|
started.elapsed(),
|
|
attempts,
|
|
)
|
|
.map_err(AttemptError::new)
|
|
}
|
|
}
|
|
|
|
struct AttemptError {
|
|
error: LlmError,
|
|
retry_after: Option<Duration>,
|
|
}
|
|
|
|
impl AttemptError {
|
|
const fn new(error: LlmError) -> Self {
|
|
Self {
|
|
error,
|
|
retry_after: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn read_bounded(
|
|
mut response: reqwest::Response,
|
|
maximum: usize,
|
|
cancellation: &CancellationToken,
|
|
) -> Result<Vec<u8>, AttemptError> {
|
|
if response
|
|
.content_length()
|
|
.is_some_and(|length| usize::try_from(length).map_or(true, |length| length > maximum))
|
|
{
|
|
return Err(AttemptError::new(LlmError::ResponseTooLarge));
|
|
}
|
|
let mut body = Vec::with_capacity(
|
|
response
|
|
.content_length()
|
|
.and_then(|value| usize::try_from(value).ok())
|
|
.unwrap_or(0)
|
|
.min(maximum),
|
|
);
|
|
loop {
|
|
let chunk = tokio::select! {
|
|
() = cancellation.cancelled() => return Err(AttemptError::new(LlmError::Cancelled)),
|
|
chunk = response.chunk() => chunk.map_err(|error| AttemptError::new(classify_reqwest(&error)))?,
|
|
};
|
|
let Some(chunk) = chunk else { break };
|
|
if body.len().saturating_add(chunk.len()) > maximum {
|
|
return Err(AttemptError::new(LlmError::ResponseTooLarge));
|
|
}
|
|
body.extend_from_slice(&chunk);
|
|
}
|
|
Ok(body)
|
|
}
|
|
|
|
fn request_body(
|
|
messages: &[CompletionMessage],
|
|
tools: &[ToolDefinition],
|
|
) -> Result<Vec<u8>, LlmError> {
|
|
let messages = messages.iter().map(message_wire_value).collect::<Vec<_>>();
|
|
let tools = tools
|
|
.iter()
|
|
.map(|tool| {
|
|
json!({
|
|
"type":"function",
|
|
"function": {
|
|
"name": tool.name.as_str(),
|
|
"description": tool.description.as_str(),
|
|
"parameters": tool.schema.wire_value()
|
|
}
|
|
})
|
|
})
|
|
.collect::<Vec<_>>();
|
|
serde_json::to_vec(&json!({"messages":messages, "tools":tools}))
|
|
.map_err(|_| LlmError::MalformedJson)
|
|
}
|
|
|
|
fn message_wire_value(message: &CompletionMessage) -> Value {
|
|
let role = match message.role {
|
|
MessageRole::System => "system",
|
|
MessageRole::Avatar => "user",
|
|
MessageRole::Agent => "assistant",
|
|
MessageRole::Tool => "tool",
|
|
};
|
|
let content = message.content.as_slice().iter().map(|part| match part {
|
|
ContentPart::Text(text) => json!({"type":"text", "text":text.as_str()}),
|
|
ContentPart::Image { url, detail } => json!({
|
|
"type":"image_url",
|
|
"image_url":{"url":url.as_str(), "detail":match detail { ImageDetail::Auto=>"auto", ImageDetail::Low=>"low", ImageDetail::High=>"high"}}
|
|
}),
|
|
}).collect::<Vec<_>>();
|
|
let mut value = json!({"role":role,"content":content});
|
|
if let Some(call_id) = &message.tool_call_id {
|
|
value["tool_call_id"] = Value::String(call_id.as_str().to_owned());
|
|
}
|
|
if !message.proposed_calls.is_empty() {
|
|
value["tool_calls"] = Value::Array(message.proposed_calls.as_slice().iter().map(|call| json!({
|
|
"id":call.call_id.as_str(), "type":"function", "function":{"name":call.name.as_str(),"arguments":call.arguments_json.as_str()}
|
|
})).collect());
|
|
}
|
|
value
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct WireResponse {
|
|
choices: Vec<WireChoice>,
|
|
usage: Option<WireUsage>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct WireChoice {
|
|
message: WireMessage,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct WireMessage {
|
|
content: Option<String>,
|
|
#[serde(default)]
|
|
tool_calls: Vec<WireToolCall>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct WireToolCall {
|
|
id: String,
|
|
#[serde(rename = "type")]
|
|
kind: Option<String>,
|
|
function: WireFunction,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct WireFunction {
|
|
name: String,
|
|
arguments: String,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct WireUsage {
|
|
#[serde(rename = "prompt_tokens")]
|
|
prompt: Option<u64>,
|
|
#[serde(rename = "completion_tokens")]
|
|
completion: Option<u64>,
|
|
#[serde(rename = "total_tokens")]
|
|
total: Option<u64>,
|
|
}
|
|
|
|
fn parse_completion(
|
|
request_id: u64,
|
|
correlation_id: BoundedText<MAX_IDENTIFIER_BYTES>,
|
|
body: &[u8],
|
|
latency: Duration,
|
|
attempts: usize,
|
|
) -> Result<Completion, LlmError> {
|
|
let wire: WireResponse = serde_json::from_slice(body).map_err(|_| LlmError::MalformedJson)?;
|
|
let choice = wire
|
|
.choices
|
|
.into_iter()
|
|
.next()
|
|
.ok_or(LlmError::UnsupportedResponse)?;
|
|
let mut proposed_calls = BoundedVec::new();
|
|
let mut ids = BTreeSet::new();
|
|
for call in choice.message.tool_calls {
|
|
if call.kind.as_deref().is_some_and(|kind| kind != "function") {
|
|
return Err(LlmError::UnsupportedResponse);
|
|
}
|
|
if !ids.insert(call.id.clone()) {
|
|
return Err(LlmError::DuplicateToolCallId);
|
|
}
|
|
proposed_calls
|
|
.try_push(
|
|
"completion.tool_calls",
|
|
ProposedToolCall::from_model_output(
|
|
call.id,
|
|
call.function.name,
|
|
call.function.arguments,
|
|
)?,
|
|
)
|
|
.map_err(LlmError::Boundary)?;
|
|
}
|
|
let mut content = BoundedVec::new();
|
|
if let Some(text) = choice.message.content {
|
|
content
|
|
.try_push(
|
|
"completion.content",
|
|
ContentPart::Text(BoundedText::new_allow_empty("completion.text", text)?),
|
|
)
|
|
.map_err(LlmError::Boundary)?;
|
|
}
|
|
if content.is_empty() && proposed_calls.is_empty() {
|
|
return Err(LlmError::UnsupportedResponse);
|
|
}
|
|
Ok(Completion {
|
|
request_id,
|
|
correlation_id,
|
|
message: CompletionMessage {
|
|
role: MessageRole::Agent,
|
|
content,
|
|
tool_call_id: None,
|
|
proposed_calls,
|
|
},
|
|
usage: wire.usage.map(|usage| Usage {
|
|
prompt_tokens: usage.prompt,
|
|
completion_tokens: usage.completion,
|
|
total_tokens: usage.total,
|
|
}),
|
|
latency,
|
|
attempts,
|
|
})
|
|
}
|
|
|
|
fn classify_reqwest(error: &reqwest::Error) -> LlmError {
|
|
if error.is_timeout() {
|
|
LlmError::Timeout
|
|
} else {
|
|
LlmError::Transport
|
|
}
|
|
}
|
|
|
|
fn transient(error: &LlmError) -> bool {
|
|
matches!(error, LlmError::Transport | LlmError::Timeout)
|
|
|| matches!(
|
|
error,
|
|
LlmError::HttpStatus(408 | 425 | 429 | 500 | 502 | 503 | 504)
|
|
)
|
|
}
|
|
|
|
fn deterministic_backoff(request_id: u64, attempt: usize) -> Duration {
|
|
let exponent = u32::try_from(attempt).unwrap_or(8).min(8);
|
|
let base = 100_u64.saturating_mul(1_u64 << exponent);
|
|
let jitter = request_id
|
|
.wrapping_mul(6_364_136_223_846_793_005)
|
|
.rotate_left(exponent)
|
|
% 97;
|
|
Duration::from_millis(base.saturating_add(jitter))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::config::AgentConfig;
|
|
|
|
#[tokio::test]
|
|
async fn total_timeout_includes_waiting_for_a_concurrency_slot() {
|
|
let limits = LlmTransportLimits {
|
|
total_timeout: Duration::from_millis(20),
|
|
max_concurrent_requests: 1,
|
|
..LlmTransportLimits::default()
|
|
};
|
|
let connection = AgentConfig::offline("http://127.0.0.1:9/exact", "test-key")
|
|
.expect("test config")
|
|
.llm;
|
|
let client = LlmClient::new(connection, limits).expect("test client");
|
|
let permit = Arc::clone(&client.slots)
|
|
.acquire_owned()
|
|
.await
|
|
.expect("semaphore open");
|
|
assert_eq!(
|
|
client
|
|
.complete(
|
|
&[CompletionMessage::text(MessageRole::Avatar, "hello").expect("message")],
|
|
&[],
|
|
&CancellationToken::default(),
|
|
)
|
|
.await
|
|
.expect_err("slot wait must time out"),
|
|
LlmError::Timeout
|
|
);
|
|
drop(permit);
|
|
}
|
|
}
|