feat(agent): add agent/llm crate (LLM abstraction layer with providers, embeddings, fallback, skills, types)
zhenyi committed
Added agent/llm/types/capabilities.rs agent/llm/types/capabilities.rs +75 −0 View agent/llm/types/capabilities.rs
@@ -0,0 +1,75 @@
1use serde::{Deserialize, Serialize};
2
3/// What a provider+model pair supports.
4///
5/// Serializable so capability tables can be declared in configuration rather
6/// than hard-coded per provider.
7#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
8pub struct ProviderCapabilities {
9 #[serde(default)]
10 pub streaming: bool,
11 #[serde(default)]
12 pub tool_calling: bool,
13 #[serde(default)]
14 pub structured_output: bool,
15 #[serde(default)]
16 pub vision: bool,
17 #[serde(default)]
18 pub reasoning: bool,
19 #[serde(default)]
20 pub embeddings: bool,
21 #[serde(default)]
22 pub max_context_tokens: usize,
23 #[serde(default)]
24 pub max_output_tokens: usize,
25}
26
27impl ProviderCapabilities {
28 /// Whether `input_tokens` leaves room for `max_output_tokens` of output.
29 ///
30 /// A `max_context_tokens` of 0 means "unknown", and always fits — an
31 /// unconfigured capability table must not silently reject every request.
32 pub fn fits_context(&self, input_tokens: usize, output_tokens: usize) -> bool {
33 self.max_context_tokens == 0
34 || input_tokens.saturating_add(output_tokens) <= self.max_context_tokens
35 }
36
37 /// Remaining input budget after reserving room for `output_tokens`.
38 /// `None` when the context window is unknown.
39 pub fn remaining_input_budget(&self, output_tokens: usize) -> Option<usize> {
40 (self.max_context_tokens != 0)
41 .then(|| self.max_context_tokens.saturating_sub(output_tokens))
42 }
43}
44
45#[cfg(test)]
46mod tests {
47 use super::*;
48
49 #[test]
50 fn unknown_context_window_never_rejects() {
51 let caps = ProviderCapabilities::default();
52 assert!(caps.fits_context(usize::MAX, usize::MAX));
53 assert_eq!(caps.remaining_input_budget(10), None);
54 }
55
56 #[test]
57 fn known_context_window_enforces_the_limit() {
58 let caps = ProviderCapabilities {
59 max_context_tokens: 100,
60 ..Default::default()
61 };
62 assert!(caps.fits_context(90, 10));
63 assert!(!caps.fits_context(91, 10));
64 assert_eq!(caps.remaining_input_budget(10), Some(90));
65 }
66
67 #[test]
68 fn budget_saturates_rather_than_underflowing() {
69 let caps = ProviderCapabilities {
70 max_context_tokens: 10,
71 ..Default::default()
72 };
73 assert_eq!(caps.remaining_input_budget(999), Some(0));
74 }
75}
Added agent/llm/types/message.rs agent/llm/types/message.rs +212 −0 View agent/llm/types/message.rs
@@ -0,0 +1,212 @@
1use crate::types::tools::ToolCall;
2use serde::{Deserialize, Serialize};
3
4#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
5#[serde(rename_all = "snake_case")]
6pub enum MessageRole {
7 System,
8 User,
9 Assistant,
10 Tool,
11}
12
13#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
14#[serde(untagged)]
15pub enum MessageContent {
16 Text(String),
17 Parts(Vec<ContentPart>),
18}
19
20impl Default for MessageContent {
21 fn default() -> Self {
22 Self::Text(String::new())
23 }
24}
25
26#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
27#[serde(rename_all = "snake_case")]
28pub enum ContentPart {
29 Text(String),
30 Image { url: String, detail: ImageDetail },
31}
32
33#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
34#[serde(rename_all = "snake_case")]
35pub enum ImageDetail {
36 #[default]
37 Auto,
38 Low,
39 High,
40}
41
42#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
43pub struct ChatMessage {
44 pub role: MessageRole,
45 pub content: MessageContent,
46 #[serde(default, skip_serializing_if = "Option::is_none")]
47 pub name: Option<String>,
48 #[serde(default, skip_serializing_if = "Vec::is_empty")]
49 pub tool_calls: Vec<ToolCall>,
50 /// Required when `role` is [`MessageRole::Tool`]; see
51 /// [`ChatMessage::validate`].
52 #[serde(default, skip_serializing_if = "Option::is_none")]
53 pub tool_call_id: Option<String>,
54}
55
56impl ChatMessage {
57 fn text(role: MessageRole, text: impl Into<String>) -> Self {
58 Self {
59 role,
60 content: MessageContent::Text(text.into()),
61 name: None,
62 tool_calls: Vec::new(),
63 tool_call_id: None,
64 }
65 }
66
67 pub fn system(text: impl Into<String>) -> Self {
68 Self::text(MessageRole::System, text)
69 }
70
71 pub fn user(text: impl Into<String>) -> Self {
72 Self::text(MessageRole::User, text)
73 }
74
75 pub fn assistant(text: impl Into<String>) -> Self {
76 Self::text(MessageRole::Assistant, text)
77 }
78
79 /// A tool result. `tool_call_id` is mandatory for this role, so this
80 /// constructor makes it impossible to build one that omits it.
81 pub fn tool(tool_call_id: impl Into<String>, text: impl Into<String>) -> Self {
82 Self {
83 tool_call_id: Some(tool_call_id.into()),
84 ..Self::text(MessageRole::Tool, text)
85 }
86 }
87
88 pub fn with_name(mut self, name: impl Into<String>) -> Self {
89 self.name = Some(name.into());
90 self
91 }
92
93 pub fn with_tool_calls(mut self, tool_calls: Vec<ToolCall>) -> Self {
94 self.tool_calls = tool_calls;
95 self
96 }
97
98 /// Check role-dependent invariants the type system cannot express.
99 ///
100 /// Deserialization bypasses the constructors, so anything arriving from
101 /// outside the process should be validated before it reaches a provider.
102 pub fn validate(&self) -> Result<(), &'static str> {
103 match self.role {
104 MessageRole::Tool if self.tool_call_id.is_none() => {
105 Err("tool message requires tool_call_id")
106 }
107 MessageRole::Tool if !self.tool_calls.is_empty() => {
108 Err("tool message must not carry tool_calls")
109 }
110 MessageRole::Assistant | MessageRole::Tool => Ok(()),
111 _ if !self.tool_calls.is_empty() => Err("only assistant messages may carry tool_calls"),
112 _ => Ok(()),
113 }
114 }
115
116 /// Plain-text view of the content, concatenating text parts and skipping
117 /// images.
118 pub fn text_content(&self) -> String {
119 match &self.content {
120 MessageContent::Text(s) => s.clone(),
121 MessageContent::Parts(parts) => parts
122 .iter()
123 .filter_map(|p| match p {
124 ContentPart::Text(s) => Some(s.as_str()),
125 ContentPart::Image { .. } => None,
126 })
127 .collect(),
128 }
129 }
130}
131
132#[cfg(test)]
133mod tests {
134 use super::*;
135 use serde_json::json;
136
137 fn call() -> ToolCall {
138 ToolCall {
139 id: "c".into(),
140 name: "n".into(),
141 arguments: json!({}),
142 }
143 }
144
145 #[test]
146 fn tool_constructor_always_sets_the_id() {
147 let m = ChatMessage::tool("call_1", "42");
148 assert_eq!(m.tool_call_id.as_deref(), Some("call_1"));
149 assert!(m.validate().is_ok());
150 }
151
152 #[test]
153 fn validate_rejects_tool_message_without_id() {
154 let mut m = ChatMessage::tool("call_1", "42");
155 m.tool_call_id = None;
156 assert!(m.validate().is_err());
157 }
158
159 #[test]
160 fn validate_rejects_tool_calls_on_non_assistant_roles() {
161 assert!(
162 ChatMessage::user("hi")
163 .with_tool_calls(vec![call()])
164 .validate()
165 .is_err()
166 );
167 assert!(
168 ChatMessage::system("hi")
169 .with_tool_calls(vec![call()])
170 .validate()
171 .is_err()
172 );
173 assert!(
174 ChatMessage::assistant("")
175 .with_tool_calls(vec![call()])
176 .validate()
177 .is_ok()
178 );
179 }
180
181 #[test]
182 fn plain_messages_validate() {
183 assert!(ChatMessage::user("hi").validate().is_ok());
184 assert!(ChatMessage::system("hi").validate().is_ok());
185 }
186
187 #[test]
188 fn text_content_flattens_parts_and_skips_images() {
189 let m = ChatMessage {
190 role: MessageRole::User,
191 content: MessageContent::Parts(vec![
192 ContentPart::Text("a".into()),
193 ContentPart::Image {
194 url: "u".into(),
195 detail: ImageDetail::Low,
196 },
197 ContentPart::Text("b".into()),
198 ]),
199 name: None,
200 tool_calls: Vec::new(),
201 tool_call_id: None,
202 };
203 assert_eq!(m.text_content(), "ab");
204 }
205
206 #[test]
207 fn round_trips_through_json() {
208 let m = ChatMessage::user("hi").with_name("alice");
209 let back: ChatMessage = serde_json::from_str(&serde_json::to_string(&m).unwrap()).unwrap();
210 assert_eq!(back, m);
211 }
212}
Added agent/llm/types/request.rs agent/llm/types/request.rs +136 −0 View agent/llm/types/request.rs
@@ -0,0 +1,136 @@
1use crate::types::message::ChatMessage;
2use crate::types::tools::ToolDefinition;
3use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
6#[serde(rename_all = "snake_case")]
7pub enum ResponseFormat {
8 #[default]
9 Text,
10 Json,
11}
12
13/// A provider-agnostic chat completion request.
14///
15/// Serializable so requests can be recorded and replayed — the primary way to
16/// build regression tests against real provider traffic.
17///
18/// `#[non_exhaustive]` keeps future fields from being breaking changes; build
19/// values with [`ChatRequest::new`] plus the `with_*` methods, or from
20/// [`Default`].
21#[derive(Debug, Clone, Default, Serialize, Deserialize)]
22#[non_exhaustive]
23pub struct ChatRequest {
24 pub messages: Vec<ChatMessage>,
25 #[serde(default, skip_serializing_if = "Vec::is_empty")]
26 pub tools: Vec<ToolDefinition>,
27 #[serde(default, skip_serializing_if = "Option::is_none")]
28 pub max_tokens: Option<usize>,
29 #[serde(default, skip_serializing_if = "Option::is_none")]
30 pub temperature: Option<f32>,
31 #[serde(default, skip_serializing_if = "Option::is_none")]
32 pub top_p: Option<f32>,
33 #[serde(default, skip_serializing_if = "Option::is_none")]
34 pub response_format: Option<ResponseFormat>,
35 #[serde(default, skip_serializing_if = "Vec::is_empty")]
36 pub stop_sequences: Vec<String>,
37 #[serde(default, skip_serializing_if = "Option::is_none")]
38 pub provider_options: Option<serde_json::Value>,
39 #[serde(default, skip_serializing_if = "Option::is_none")]
40 pub model_override: Option<String>,
41}
42
43impl ChatRequest {
44 pub fn new(messages: Vec<ChatMessage>) -> Self {
45 Self {
46 messages,
47 ..Default::default()
48 }
49 }
50
51 pub fn with_tools(mut self, tools: Vec<ToolDefinition>) -> Self {
52 self.tools = tools;
53 self
54 }
55
56 pub fn with_max_tokens(mut self, max_tokens: usize) -> Self {
57 self.max_tokens = Some(max_tokens);
58 self
59 }
60
61 /// Set sampling temperature, clamped to the `0.0..=2.0` range every major
62 /// provider accepts. Out-of-range values are otherwise only rejected at
63 /// the API boundary, after the request has been paid for.
64 pub fn with_temperature(mut self, temperature: f32) -> Self {
65 self.temperature = Some(temperature.clamp(0.0, 2.0));
66 self
67 }
68
69 /// Set nucleus sampling cutoff, clamped to `0.0..=1.0`.
70 pub fn with_top_p(mut self, top_p: f32) -> Self {
71 self.top_p = Some(top_p.clamp(0.0, 1.0));
72 self
73 }
74
75 pub fn with_response_format(mut self, format: ResponseFormat) -> Self {
76 self.response_format = Some(format);
77 self
78 }
79
80 pub fn with_stop_sequences(mut self, stop: Vec<String>) -> Self {
81 self.stop_sequences = stop;
82 self
83 }
84
85 pub fn with_provider_options(mut self, options: serde_json::Value) -> Self {
86 self.provider_options = Some(options);
87 self
88 }
89
90 pub fn with_model_override(mut self, model: impl Into<String>) -> Self {
91 self.model_override = Some(model.into());
92 self
93 }
94}
95
96#[cfg(test)]
97mod tests {
98 use super::*;
99
100 #[test]
101 fn builder_sets_fields() {
102 let req = ChatRequest::new(vec![ChatMessage::user("hi")])
103 .with_max_tokens(64)
104 .with_model_override("gpt-4o");
105 assert_eq!(req.max_tokens, Some(64));
106 assert_eq!(req.model_override.as_deref(), Some("gpt-4o"));
107 }
108
109 #[test]
110 fn sampling_parameters_are_clamped() {
111 assert_eq!(
112 ChatRequest::default().with_temperature(9.0).temperature,
113 Some(2.0)
114 );
115 assert_eq!(
116 ChatRequest::default().with_temperature(-1.0).temperature,
117 Some(0.0)
118 );
119 assert_eq!(ChatRequest::default().with_top_p(5.0).top_p, Some(1.0));
120 }
121
122 #[test]
123 fn round_trips_through_json() {
124 let req = ChatRequest::new(vec![ChatMessage::user("hi")]).with_max_tokens(8);
125 let json = serde_json::to_string(&req).unwrap();
126 let back: ChatRequest = serde_json::from_str(&json).unwrap();
127 assert_eq!(back.messages.len(), 1);
128 assert_eq!(back.max_tokens, Some(8));
129 }
130
131 #[test]
132 fn empty_optionals_are_omitted_from_json() {
133 let json = serde_json::to_string(&ChatRequest::new(vec![])).unwrap();
134 assert_eq!(json, r#"{"messages":[]}"#);
135 }
136}
Added agent/llm/types/response.rs agent/llm/types/response.rs +115 −0 View agent/llm/types/response.rs
@@ -0,0 +1,115 @@
1use crate::token::TokenUsageStats;
2use crate::types::tools::ToolCall;
3use serde::{Deserialize, Serialize};
4
5/// A completed (non-streaming) chat response.
6///
7/// Serializable so responses can be recorded alongside their requests and
8/// replayed in tests.
9#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
10pub struct ChatResponse {
11 pub content: String,
12 #[serde(default, skip_serializing_if = "Vec::is_empty")]
13 pub tool_calls: Vec<ToolCall>,
14 pub finish_reason: FinishReason,
15 #[serde(default)]
16 pub usage: TokenUsageStats,
17 pub model: String,
18 pub provider: String,
19}
20
21impl ChatResponse {
22 pub fn new(
23 content: impl Into<String>,
24 provider: impl Into<String>,
25 model: impl Into<String>,
26 ) -> Self {
27 Self {
28 content: content.into(),
29 tool_calls: Vec::new(),
30 finish_reason: FinishReason::Stop,
31 usage: TokenUsageStats::default(),
32 model: model.into(),
33 provider: provider.into(),
34 }
35 }
36
37 pub fn with_tool_calls(mut self, tool_calls: Vec<ToolCall>) -> Self {
38 self.finish_reason = FinishReason::ToolUse;
39 self.tool_calls = tool_calls;
40 self
41 }
42
43 pub fn with_usage(mut self, usage: TokenUsageStats) -> Self {
44 self.usage = usage;
45 self
46 }
47
48 pub fn with_finish_reason(mut self, reason: FinishReason) -> Self {
49 self.finish_reason = reason;
50 self
51 }
52
53 /// Whether the model asked to invoke tools, meaning the agent loop should
54 /// execute them and continue rather than treat `content` as final.
55 pub fn requests_tools(&self) -> bool {
56 !self.tool_calls.is_empty()
57 }
58
59 /// Whether the response was cut short and `content` may be incomplete.
60 pub fn is_truncated(&self) -> bool {
61 matches!(
62 self.finish_reason,
63 FinishReason::MaxTokens | FinishReason::ContentFilter | FinishReason::Error
64 )
65 }
66}
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
69#[serde(rename_all = "snake_case")]
70pub enum FinishReason {
71 Stop,
72 ToolUse,
73 MaxTokens,
74 ContentFilter,
75 Error,
76}
77
78#[cfg(test)]
79mod tests {
80 use super::*;
81 use serde_json::json;
82
83 #[test]
84 fn tool_calls_imply_tool_use_finish_reason() {
85 let r = ChatResponse::new("", "openai", "gpt-4o").with_tool_calls(vec![ToolCall {
86 id: "c".into(),
87 name: "n".into(),
88 arguments: json!({}),
89 }]);
90 assert_eq!(r.finish_reason, FinishReason::ToolUse);
91 assert!(r.requests_tools());
92 assert!(!r.is_truncated());
93 }
94
95 #[test]
96 fn truncation_covers_incomplete_outcomes() {
97 let base = ChatResponse::new("x", "p", "m");
98 assert!(!base.clone().is_truncated());
99 for reason in [
100 FinishReason::MaxTokens,
101 FinishReason::ContentFilter,
102 FinishReason::Error,
103 ] {
104 assert!(base.clone().with_finish_reason(reason).is_truncated());
105 }
106 }
107
108 #[test]
109 fn round_trips_through_json() {
110 let r = ChatResponse::new("hi", "openai", "gpt-4o")
111 .with_usage(TokenUsageStats::new(3, 4).with_reasoning(1));
112 let back: ChatResponse = serde_json::from_str(&serde_json::to_string(&r).unwrap()).unwrap();
113 assert_eq!(back, r);
114 }
115}
Added agent/llm/types/stream.rs agent/llm/types/stream.rs +286 −0 View agent/llm/types/stream.rs
@@ -0,0 +1,286 @@
1use crate::error::LlmError;
2use crate::token::TokenUsageStats;
3use crate::types::response::FinishReason;
4use crate::types::tools::ToolCall;
5use std::collections::BTreeMap;
6
7#[derive(Debug, Clone, PartialEq)]
8pub enum StreamEvent {
9 TextDelta {
10 delta: String,
11 },
12 ReasoningDelta {
13 delta: String,
14 },
15 ToolCallDelta {
16 index: Option<usize>,
17 id: String,
18 name: Option<String>,
19 args_delta: String,
20 },
21 Usage(TokenUsageStats),
22 Done {
23 finish_reason: FinishReason,
24 },
25}
26
27/// Reassembles a streamed response into its final form.
28///
29/// Providers emit tool calls as fragments — the arguments arrive as partial
30/// JSON text spread over many events, and `name` typically appears only on the
31/// first fragment. Every provider would otherwise reimplement this same
32/// grouping/concatenation/parse state machine, so it lives here once.
33///
34/// Fragments are grouped by `index` when the provider supplies one (the
35/// reliable discriminator for parallel tool calls) and by `id` otherwise.
36#[derive(Debug, Default)]
37pub struct StreamAccumulator {
38 text: String,
39 reasoning: String,
40 calls: BTreeMap<CallKey, PartialToolCall>,
41 usage: Option<TokenUsageStats>,
42 finish_reason: Option<FinishReason>,
43}
44
45#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
46enum CallKey {
47 Index(usize),
48 Id(String),
49}
50
51#[derive(Debug, Default)]
52struct PartialToolCall {
53 id: String,
54 name: Option<String>,
55 args: String,
56}
57
58impl StreamAccumulator {
59 pub fn new() -> Self {
60 Self::default()
61 }
62
63 pub fn push(&mut self, event: &StreamEvent) {
64 match event {
65 StreamEvent::TextDelta { delta } => self.text.push_str(delta),
66 StreamEvent::ReasoningDelta { delta } => self.reasoning.push_str(delta),
67 StreamEvent::ToolCallDelta {
68 index,
69 id,
70 name,
71 args_delta,
72 } => {
73 let key = match index {
74 Some(i) => CallKey::Index(*i),
75 None => CallKey::Id(id.clone()),
76 };
77 let entry = self.calls.entry(key).or_default();
78 // Later fragments usually carry an empty id/name; never let
79 // them erase what the opening fragment established.
80 if !id.is_empty() {
81 entry.id = id.clone();
82 }
83 if let Some(name) = name
84 && !name.is_empty()
85 {
86 entry.name = Some(name.clone());
87 }
88 entry.args.push_str(args_delta);
89 }
90 StreamEvent::Usage(usage) => match &mut self.usage {
91 Some(existing) => existing.merge(usage),
92 None => self.usage = Some(usage.clone()),
93 },
94 StreamEvent::Done { finish_reason } => self.finish_reason = Some(*finish_reason),
95 }
96 }
97
98 pub fn text(&self) -> &str {
99 &self.text
100 }
101
102 pub fn reasoning(&self) -> &str {
103 &self.reasoning
104 }
105
106 pub fn usage(&self) -> Option<&TokenUsageStats> {
107 self.usage.as_ref()
108 }
109
110 pub fn finish_reason(&self) -> Option<FinishReason> {
111 self.finish_reason
112 }
113
114 /// Parse the accumulated tool-call fragments.
115 ///
116 /// Fails if a call never received a name or if its arguments did not form
117 /// valid JSON — a truncated stream is a real failure, not an empty call.
118 pub fn tool_calls(&self) -> Result<Vec<ToolCall>, LlmError> {
119 self.calls
120 .values()
121 .map(|p| {
122 let name = p.name.clone().ok_or_else(|| {
123 LlmError::Stream(format!("tool call '{}' never received a name", p.id))
124 })?;
125 let args = p.args.trim();
126 let arguments = if args.is_empty() {
127 serde_json::Value::Object(serde_json::Map::new())
128 } else {
129 serde_json::from_str(args).map_err(|e| {
130 LlmError::Stream(format!("tool call '{name}' has malformed arguments: {e}"))
131 })?
132 };
133 Ok(ToolCall {
134 id: p.id.clone(),
135 name,
136 arguments,
137 })
138 })
139 .collect()
140 }
141
142 /// Build the final response. `Done` must have been observed, otherwise the
143 /// stream ended early and the result would silently look complete.
144 pub fn finish(
145 self,
146 provider: impl Into<String>,
147 model: impl Into<String>,
148 ) -> Result<crate::types::response::ChatResponse, LlmError> {
149 let tool_calls = self.tool_calls()?;
150 let finish_reason = self
151 .finish_reason
152 .ok_or_else(|| LlmError::Stream("stream ended without a Done event".into()))?;
153 Ok(crate::types::response::ChatResponse {
154 content: self.text,
155 tool_calls,
156 finish_reason,
157 usage: self.usage.unwrap_or_default(),
158 model: model.into(),
159 provider: provider.into(),
160 })
161 }
162}
163
164#[cfg(test)]
165mod tests {
166 use super::*;
167 use serde_json::json;
168
169 fn delta(index: Option<usize>, id: &str, name: Option<&str>, args: &str) -> StreamEvent {
170 StreamEvent::ToolCallDelta {
171 index,
172 id: id.to_string(),
173 name: name.map(str::to_string),
174 args_delta: args.to_string(),
175 }
176 }
177
178 #[test]
179 fn concatenates_text_and_reasoning_separately() {
180 let mut acc = StreamAccumulator::new();
181 acc.push(&StreamEvent::TextDelta { delta: "he".into() });
182 acc.push(&StreamEvent::ReasoningDelta { delta: "hm".into() });
183 acc.push(&StreamEvent::TextDelta {
184 delta: "llo".into(),
185 });
186 assert_eq!(acc.text(), "hello");
187 assert_eq!(acc.reasoning(), "hm");
188 }
189
190 #[test]
191 fn reassembles_split_tool_call_arguments() {
192 let mut acc = StreamAccumulator::new();
193 acc.push(&delta(Some(0), "c1", Some("search"), r#"{"q":"#));
194 acc.push(&delta(Some(0), "", None, r#""rust"}"#));
195 let calls = acc.tool_calls().unwrap();
196 assert_eq!(calls.len(), 1);
197 assert_eq!(calls[0].id, "c1");
198 assert_eq!(calls[0].name, "search");
199 assert_eq!(calls[0].arguments, json!({ "q": "rust" }));
200 }
201
202 #[test]
203 fn keeps_parallel_calls_separate_by_index() {
204 let mut acc = StreamAccumulator::new();
205 acc.push(&delta(Some(0), "c1", Some("a"), r#"{"x":1}"#));
206 acc.push(&delta(Some(1), "c2", Some("b"), r#"{"y":"#));
207 acc.push(&delta(Some(1), "", None, "2}"));
208 let calls = acc.tool_calls().unwrap();
209 assert_eq!(calls.len(), 2);
210 assert_eq!(calls[0].name, "a");
211 assert_eq!(calls[1].arguments, json!({ "y": 2 }));
212 }
213
214 #[test]
215 fn groups_by_id_when_index_is_absent() {
216 let mut acc = StreamAccumulator::new();
217 acc.push(&delta(None, "c1", Some("a"), r#"{"x":"#));
218 acc.push(&delta(None, "c1", None, "1}"));
219 acc.push(&delta(None, "c2", Some("b"), "{}"));
220 assert_eq!(acc.tool_calls().unwrap().len(), 2);
221 }
222
223 #[test]
224 fn later_fragments_do_not_erase_id_or_name() {
225 let mut acc = StreamAccumulator::new();
226 acc.push(&delta(Some(0), "c1", Some("search"), "{"));
227 acc.push(&delta(Some(0), "", Some(""), "}"));
228 let calls = acc.tool_calls().unwrap();
229 assert_eq!(calls[0].id, "c1");
230 assert_eq!(calls[0].name, "search");
231 }
232
233 #[test]
234 fn empty_arguments_become_an_empty_object() {
235 let mut acc = StreamAccumulator::new();
236 acc.push(&delta(Some(0), "c1", Some("noop"), ""));
237 assert_eq!(acc.tool_calls().unwrap()[0].arguments, json!({}));
238 }
239
240 #[test]
241 fn truncated_arguments_are_reported_not_swallowed() {
242 let mut acc = StreamAccumulator::new();
243 acc.push(&delta(Some(0), "c1", Some("search"), r#"{"q":"ru"#));
244 assert!(matches!(acc.tool_calls(), Err(LlmError::Stream(_))));
245 }
246
247 #[test]
248 fn call_without_a_name_is_an_error() {
249 let mut acc = StreamAccumulator::new();
250 acc.push(&delta(Some(0), "c1", None, "{}"));
251 assert!(matches!(acc.tool_calls(), Err(LlmError::Stream(_))));
252 }
253
254 #[test]
255 fn usage_events_accumulate() {
256 let mut acc = StreamAccumulator::new();
257 acc.push(&StreamEvent::Usage(TokenUsageStats::new(5, 1)));
258 acc.push(&StreamEvent::Usage(TokenUsageStats::new(0, 2)));
259 assert_eq!(acc.usage().unwrap().output_tokens, 3);
260 }
261
262 #[test]
263 fn finish_builds_a_response() {
264 let mut acc = StreamAccumulator::new();
265 acc.push(&StreamEvent::TextDelta { delta: "hi".into() });
266 acc.push(&StreamEvent::Usage(TokenUsageStats::new(2, 1)));
267 acc.push(&StreamEvent::Done {
268 finish_reason: FinishReason::Stop,
269 });
270 let resp = acc.finish("openai", "gpt-4o").unwrap();
271 assert_eq!(resp.content, "hi");
272 assert_eq!(resp.finish_reason, FinishReason::Stop);
273 assert_eq!(resp.usage.input_tokens, 2);
274 assert_eq!(resp.provider, "openai");
275 }
276
277 #[test]
278 fn finish_rejects_a_stream_that_never_completed() {
279 let mut acc = StreamAccumulator::new();
280 acc.push(&StreamEvent::TextDelta { delta: "hi".into() });
281 assert!(matches!(
282 acc.finish("openai", "gpt-4o"),
283 Err(LlmError::Stream(_))
284 ));
285 }
286}
Added agent/llm/types/tools.rs agent/llm/types/tools.rs +51 −0 View agent/llm/types/tools.rs
@@ -0,0 +1,51 @@
1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
4pub struct ToolCall {
5 pub id: String,
6 pub name: String,
7 pub arguments: serde_json::Value,
8}
9
10/// How a tool may be scheduled relative to other tool calls in the same turn.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
12#[serde(rename_all = "snake_case")]
13pub enum ConcurrencyPolicy {
14 /// Safe to run in parallel with any other `Safe` tool.
15 #[default]
16 Safe,
17 /// Must run alone, with no other tool in flight.
18 Exclusive,
19 /// Must run in the order the model emitted it, but may overlap `Safe` work.
20 Ordered,
21}
22
23#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
24pub struct ToolDefinition {
25 pub name: String,
26 pub description: String,
27 /// JSON Schema for the tool's arguments.
28 pub parameters: serde_json::Value,
29 #[serde(default)]
30 pub concurrency: ConcurrencyPolicy,
31}
32
33impl ToolDefinition {
34 pub fn new(
35 name: impl Into<String>,
36 description: impl Into<String>,
37 parameters: serde_json::Value,
38 ) -> Self {
39 Self {
40 name: name.into(),
41 description: description.into(),
42 parameters,
43 concurrency: ConcurrencyPolicy::default(),
44 }
45 }
46
47 pub fn with_concurrency(mut self, concurrency: ConcurrencyPolicy) -> Self {
48 self.concurrency = concurrency;
49 self
50 }
51}
Added agent/llm/Cargo.toml agent/llm/Cargo.toml +24 −0 View agent/llm/Cargo.toml
@@ -0,0 +1,24 @@
1[package]
2name = "llm"
3version.workspace = true
4edition.workspace = true
5authors.workspace = true
6categories.workspace = true
7keywords.workspace = true
8repository.workspace = true
9homepage.workspace = true
10
11[lib]
12path = "lib.rs"
13name = "llm"
14[dependencies]
15thiserror.workspace = true
16serde.workspace = true
17serde_json.workspace = true
18async-trait.workspace = true
19futures.workspace = true
20tokio-stream.workspace = true
21tracing.workspace = true
22
23[dev-dependencies]
24tokio = { workspace = true, features = ["macros", "rt"] }
No newline at end of file
Added agent/llm/embed.rs agent/llm/embed.rs +7 −0 View agent/llm/embed.rs
@@ -0,0 +1,7 @@
1use crate::error::LlmError;
2
3#[async_trait::async_trait]
4pub trait EmbeddingProvider: Send + Sync {
5 async fn embed(&self, texts: Vec<String>) -> Result<Vec<Vec<f32>>, LlmError>;
6 fn embedding_dimensions(&self) -> usize;
7}
Added agent/llm/error.rs agent/llm/error.rs +200 −0 View agent/llm/error.rs
@@ -0,0 +1,200 @@
1use std::time::Duration;
2
3#[derive(Debug, thiserror::Error)]
4pub enum LlmError {
5 #[error("Agent error: {0}")]
6 Agent(String),
7
8 #[error("Provider error: {0}")]
9 Provider(String),
10
11 #[error("Token budget exceeded: used {used}, budget {budget}")]
12 TokenBudgetExceeded { used: usize, budget: usize },
13
14 #[error("Configuration error: {0}")]
15 Config(String),
16
17 #[error("Memory error: {0}")]
18 Memory(String),
19
20 #[error("Coordination error: {0}")]
21 Coordination(String),
22
23 #[error("Isolation error: {0}")]
24 Isolation(String),
25
26 #[error("Protocol error: {0}")]
27 Protocol(String),
28
29 #[error("IO error: {0}")]
30 Io(#[from] std::io::Error),
31
32 #[error("Rate limited by {provider}. Retry after {retry_after_secs:?}s")]
33 RateLimited {
34 provider: String,
35 retry_after_secs: Option<u64>,
36 },
37
38 #[error("Context length exceeded: {tokens_used} tokens, limit {limit}")]
39 ContextLengthExceeded { tokens_used: usize, limit: usize },
40
41 #[error("Content filtered by {provider}: {reason}")]
42 ContentFiltered { provider: String, reason: String },
43
44 #[error("Authentication failed for {provider}")]
45 AuthError { provider: String },
46
47 #[error("Model not found: {model} on {provider}")]
48 ModelNotFound { provider: String, model: String },
49
50 #[error("Network error: {0}")]
51 Network(String),
52
53 #[error("Serialization error: {0}")]
54 Serialization(#[from] serde_json::Error),
55
56 #[error("{provider} API error: {status} - {message}")]
57 ApiError {
58 provider: String,
59 status: u16,
60 message: String,
61 },
62
63 #[error("Streaming error: {0}")]
64 Stream(String),
65
66 #[error("Provider not found: {0}")]
67 ProviderNotFound(String),
68 #[error("Missing required parameter '{parameter}' for skill '{skill}'")]
69 SkillMissingParameter { skill: String, parameter: String },
70
71 #[error("Skill not found: {0}")]
72 SkillNotFound(String),
73
74 #[error("Skill rendering failed: {0}")]
75 SkillRenderFailed(String),
76}
77
78impl LlmError {
79 /// Whether retrying the same request against another provider (or the same
80 /// one after a delay) could plausibly succeed.
81 ///
82 /// Deliberately excludes `ContextLengthExceeded`, `ContentFiltered`, and
83 /// `AuthError`: those reproduce deterministically, so retrying only burns
84 /// quota. `ApiError` is retryable for 5xx and 429 only.
85 pub fn is_retryable(&self) -> bool {
86 match self {
87 Self::RateLimited { .. } | Self::Network(_) => true,
88 Self::ApiError { status, .. } => *status == 429 || (500..600).contains(status),
89 Self::Io(e) => matches!(
90 e.kind(),
91 std::io::ErrorKind::TimedOut
92 | std::io::ErrorKind::Interrupted
93 | std::io::ErrorKind::ConnectionReset
94 | std::io::ErrorKind::ConnectionAborted
95 | std::io::ErrorKind::BrokenPipe
96 ),
97 _ => false,
98 }
99 }
100
101 /// Provider-advertised delay before retrying, when one was supplied.
102 pub fn retry_after(&self) -> Option<Duration> {
103 match self {
104 Self::RateLimited {
105 retry_after_secs, ..
106 } => retry_after_secs.map(Duration::from_secs),
107 _ => None,
108 }
109 }
110
111 /// The provider that produced this error, when attributable.
112 pub fn provider(&self) -> Option<&str> {
113 match self {
114 Self::RateLimited { provider, .. }
115 | Self::ContentFiltered { provider, .. }
116 | Self::AuthError { provider }
117 | Self::ModelNotFound { provider, .. }
118 | Self::ApiError { provider, .. }
119 | Self::ProviderNotFound(provider) => Some(provider),
120 _ => None,
121 }
122 }
123}
124
125#[cfg(test)]
126mod tests {
127 use super::*;
128
129 #[test]
130 fn rate_limit_and_transport_errors_are_retryable() {
131 assert!(
132 LlmError::RateLimited {
133 provider: "p".into(),
134 retry_after_secs: Some(3),
135 }
136 .is_retryable()
137 );
138 assert!(LlmError::Network("reset".into()).is_retryable());
139 }
140
141 #[test]
142 fn api_errors_retry_on_5xx_and_429_only() {
143 let err = |status| LlmError::ApiError {
144 provider: "p".into(),
145 status,
146 message: String::new(),
147 };
148 assert!(err(429).is_retryable());
149 assert!(err(500).is_retryable());
150 assert!(err(503).is_retryable());
151 assert!(!err(400).is_retryable());
152 assert!(!err(404).is_retryable());
153 }
154
155 #[test]
156 fn deterministic_failures_are_not_retryable() {
157 assert!(
158 !LlmError::ContextLengthExceeded {
159 tokens_used: 10,
160 limit: 5,
161 }
162 .is_retryable()
163 );
164 assert!(
165 !LlmError::AuthError {
166 provider: "p".into()
167 }
168 .is_retryable()
169 );
170 assert!(
171 !LlmError::ContentFiltered {
172 provider: "p".into(),
173 reason: "x".into(),
174 }
175 .is_retryable()
176 );
177 }
178
179 #[test]
180 fn retry_after_is_exposed_as_duration() {
181 let e = LlmError::RateLimited {
182 provider: "p".into(),
183 retry_after_secs: Some(30),
184 };
185 assert_eq!(e.retry_after(), Some(Duration::from_secs(30)));
186 assert_eq!(LlmError::Network("x".into()).retry_after(), None);
187 }
188
189 #[test]
190 fn provider_is_attributed_where_known() {
191 assert_eq!(
192 LlmError::AuthError {
193 provider: "openai".into()
194 }
195 .provider(),
196 Some("openai")
197 );
198 assert_eq!(LlmError::Stream("x".into()).provider(), None);
199 }
200}
Added agent/llm/fallback.rs agent/llm/fallback.rs +387 −0 View agent/llm/fallback.rs
@@ -0,0 +1,387 @@
1use crate::error::LlmError;
2use crate::provider::LlmProvider;
3use crate::types::{ChatRequest, ChatResponse, ProviderCapabilities, StreamEvent};
4use futures::StreamExt;
5use std::pin::Pin;
6use std::sync::Arc;
7use tokio_stream::Stream;
8
9type EventStream = Pin<Box<dyn Stream<Item = Result<StreamEvent, LlmError>> + Send>>;
10
11pub struct FallbackTarget {
12 pub provider: Arc<dyn LlmProvider>,
13 pub model: String,
14}
15
16/// Routes to a secondary provider when the primary fails retryably.
17///
18/// Failover triggers on [`LlmError::is_retryable`] rather than on rate limits
19/// alone, so transient network faults and provider 5xx are covered too.
20pub struct FallbackProvider {
21 primary: Arc<dyn LlmProvider>,
22 fallback: Option<FallbackTarget>,
23}
24
25impl FallbackProvider {
26 pub fn new(primary: Arc<dyn LlmProvider>, fallback: Option<FallbackTarget>) -> Self {
27 Self { primary, fallback }
28 }
29
30 /// Point a request at the fallback model.
31 ///
32 /// An explicit `model_override` from the caller is replaced — the fallback
33 /// provider almost certainly cannot serve the primary's model name — but
34 /// the discarded value is logged so the substitution is never silent.
35 fn retarget(mut request: ChatRequest, model: &str) -> ChatRequest {
36 if let Some(previous) = &request.model_override
37 && previous != model
38 {
39 tracing::debug!(
40 discarded_model_override = %previous,
41 fallback_model = %model,
42 "replacing caller model_override for fallback provider",
43 );
44 }
45 request.model_override = Some(model.to_string());
46 request
47 }
48
49 fn log_failover(&self, fb: &FallbackTarget, err: &LlmError, stage: &'static str) {
50 tracing::warn!(
51 primary = %self.primary.provider_id(),
52 primary_model = %self.primary.model_id(),
53 fallback = %fb.provider.provider_id(),
54 fallback_model = %fb.model,
55 stage,
56 error = %err,
57 "primary failed retryably; falling back",
58 );
59 }
60}
61
62#[async_trait::async_trait]
63impl LlmProvider for FallbackProvider {
64 fn provider_id(&self) -> &str {
65 self.primary.provider_id()
66 }
67
68 fn model_id(&self) -> &str {
69 self.primary.model_id()
70 }
71
72 fn capabilities(&self) -> ProviderCapabilities {
73 // Intersect: advertising a capability the fallback lacks would let a
74 // request be built that the fallback cannot serve.
75 let Some(fb) = &self.fallback else {
76 return self.primary.capabilities();
77 };
78 let (p, f) = (self.primary.capabilities(), fb.provider.capabilities());
79 ProviderCapabilities {
80 streaming: p.streaming && f.streaming,
81 tool_calling: p.tool_calling && f.tool_calling,
82 structured_output: p.structured_output && f.structured_output,
83 vision: p.vision && f.vision,
84 reasoning: p.reasoning && f.reasoning,
85 embeddings: p.embeddings && f.embeddings,
86 max_context_tokens: p.max_context_tokens.min(f.max_context_tokens),
87 max_output_tokens: p.max_output_tokens.min(f.max_output_tokens),
88 }
89 }
90
91 async fn chat(&self, request: ChatRequest) -> Result<ChatResponse, LlmError> {
92 let Some(fb) = &self.fallback else {
93 return self.primary.chat(request).await;
94 };
95 match self.primary.chat(request.clone()).await {
96 Err(err) if err.is_retryable() => {
97 self.log_failover(fb, &err, "request");
98 fb.provider.chat(Self::retarget(request, &fb.model)).await
99 }
100 other => other,
101 }
102 }
103
104 async fn chat_stream(&self, request: ChatRequest) -> Result<EventStream, LlmError> {
105 let Some(fb) = &self.fallback else {
106 return self.primary.chat_stream(request).await;
107 };
108
109 let mut stream = match self.primary.chat_stream(request.clone()).await {
110 Err(err) if err.is_retryable() => {
111 self.log_failover(fb, &err, "stream_open");
112 return fb
113 .provider
114 .chat_stream(Self::retarget(request, &fb.model))
115 .await;
116 }
117 other => other?,
118 };
119
120 // A stream can also fail *after* opening. Peek one event: until
121 // something has been yielded to the caller, a retryable error is still
122 // safely recoverable by restarting on the fallback.
123 let Some(first) = stream.next().await else {
124 return Ok(Box::pin(tokio_stream::empty()));
125 };
126
127 match first {
128 Err(err) if err.is_retryable() => {
129 self.log_failover(fb, &err, "stream_first_event");
130 fb.provider
131 .chat_stream(Self::retarget(request, &fb.model))
132 .await
133 }
134 // Put the peeked event back at the head of the stream.
135 first => Ok(Box::pin(
136 futures::stream::once(async move { first }).chain(stream),
137 )),
138 }
139 }
140
141 fn count_tokens(&self, request: &ChatRequest) -> usize {
142 self.primary.count_tokens(request)
143 }
144}
145
146#[cfg(test)]
147mod tests {
148 use super::*;
149 use crate::types::{ChatMessage, FinishReason};
150 use std::sync::Mutex;
151 use std::sync::atomic::{AtomicUsize, Ordering};
152
153 /// A provider whose responses are scripted per call.
154 struct Scripted {
155 id: String,
156 caps: ProviderCapabilities,
157 calls: AtomicUsize,
158 /// Error to return from `chat`/`chat_stream`, if any.
159 fail_with: Mutex<Option<LlmError>>,
160 /// When set, `chat_stream` opens successfully then yields this error.
161 fail_mid_stream: Mutex<Option<LlmError>>,
162 seen_model_override: Mutex<Option<String>>,
163 }
164
165 impl Scripted {
166 fn new(id: &str) -> Self {
167 Self {
168 id: id.to_string(),
169 caps: ProviderCapabilities::default(),
170 calls: AtomicUsize::new(0),
171 fail_with: Mutex::new(None),
172 fail_mid_stream: Mutex::new(None),
173 seen_model_override: Mutex::new(None),
174 }
175 }
176
177 fn failing(id: &str, err: LlmError) -> Self {
178 let s = Self::new(id);
179 *s.fail_with.lock().unwrap() = Some(err);
180 s
181 }
182
183 fn failing_mid_stream(id: &str, err: LlmError) -> Self {
184 let s = Self::new(id);
185 *s.fail_mid_stream.lock().unwrap() = Some(err);
186 s
187 }
188
189 fn with_caps(mut self, caps: ProviderCapabilities) -> Self {
190 self.caps = caps;
191 self
192 }
193
194 fn take_error(&self, slot: &Mutex<Option<LlmError>>) -> Option<LlmError> {
195 slot.lock().unwrap().take()
196 }
197 }
198
199 #[async_trait::async_trait]
200 impl LlmProvider for Scripted {
201 fn provider_id(&self) -> &str {
202 &self.id
203 }
204 fn model_id(&self) -> &str {
205 "m"
206 }
207 fn capabilities(&self) -> ProviderCapabilities {
208 self.caps
209 }
210
211 async fn chat(&self, request: ChatRequest) -> Result<ChatResponse, LlmError> {
212 self.calls.fetch_add(1, Ordering::SeqCst);
213 *self.seen_model_override.lock().unwrap() = request.model_override.clone();
214 match self.take_error(&self.fail_with) {
215 Some(err) => Err(err),
216 None => Ok(ChatResponse::new("ok", &self.id, "m")),
217 }
218 }
219
220 async fn chat_stream(&self, request: ChatRequest) -> Result<EventStream, LlmError> {
221 self.calls.fetch_add(1, Ordering::SeqCst);
222 *self.seen_model_override.lock().unwrap() = request.model_override.clone();
223 if let Some(err) = self.take_error(&self.fail_with) {
224 return Err(err);
225 }
226 let id = self.id.clone();
227 if let Some(err) = self.take_error(&self.fail_mid_stream) {
228 return Ok(Box::pin(tokio_stream::iter(vec![Err(err)])));
229 }
230 Ok(Box::pin(tokio_stream::iter(vec![
231 Ok(StreamEvent::TextDelta { delta: id }),
232 Ok(StreamEvent::Done {
233 finish_reason: FinishReason::Stop,
234 }),
235 ])))
236 }
237 }
238
239 fn rate_limited() -> LlmError {
240 LlmError::RateLimited {
241 provider: "primary".into(),
242 retry_after_secs: Some(1),
243 }
244 }
245
246 fn target(provider: Arc<dyn LlmProvider>) -> FallbackTarget {
247 FallbackTarget {
248 provider,
249 model: "fallback-model".into(),
250 }
251 }
252
253 fn request() -> ChatRequest {
254 ChatRequest::new(vec![ChatMessage::user("hi")])
255 }
256
257 async fn collect(stream: EventStream) -> Vec<Result<StreamEvent, LlmError>> {
258 stream.collect::<Vec<_>>().await
259 }
260
261 #[tokio::test]
262 async fn falls_back_on_network_error_not_just_rate_limits() {
263 let secondary = Arc::new(Scripted::new("secondary"));
264 let fb = FallbackProvider::new(
265 Arc::new(Scripted::failing(
266 "primary",
267 LlmError::Network("reset".into()),
268 )),
269 Some(target(secondary.clone())),
270 );
271 let resp = fb.chat(request()).await.unwrap();
272 assert_eq!(resp.provider, "secondary");
273 assert_eq!(secondary.calls.load(Ordering::SeqCst), 1);
274 }
275
276 #[tokio::test]
277 async fn falls_back_on_server_side_api_error() {
278 let secondary = Arc::new(Scripted::new("secondary"));
279 let err = LlmError::ApiError {
280 provider: "primary".into(),
281 status: 503,
282 message: "unavailable".into(),
283 };
284 let fb = FallbackProvider::new(
285 Arc::new(Scripted::failing("primary", err)),
286 Some(target(secondary.clone())),
287 );
288 assert_eq!(fb.chat(request()).await.unwrap().provider, "secondary");
289 }
290
291 #[tokio::test]
292 async fn does_not_fall_back_on_deterministic_errors() {
293 let secondary = Arc::new(Scripted::new("secondary"));
294 let err = LlmError::ContextLengthExceeded {
295 tokens_used: 10,
296 limit: 5,
297 };
298 let fb = FallbackProvider::new(
299 Arc::new(Scripted::failing("primary", err)),
300 Some(target(secondary.clone())),
301 );
302 assert!(fb.chat(request()).await.is_err());
303 assert_eq!(
304 secondary.calls.load(Ordering::SeqCst),
305 0,
306 "a deterministic failure must not burn fallback quota"
307 );
308 }
309
310 #[tokio::test]
311 async fn retargets_the_fallback_model() {
312 let secondary = Arc::new(Scripted::new("secondary"));
313 let fb = FallbackProvider::new(
314 Arc::new(Scripted::failing("primary", rate_limited())),
315 Some(target(secondary.clone())),
316 );
317 fb.chat(request().with_model_override("primary-only-model"))
318 .await
319 .unwrap();
320 assert_eq!(
321 secondary.seen_model_override.lock().unwrap().as_deref(),
322 Some("fallback-model")
323 );
324 }
325
326 #[tokio::test]
327 async fn falls_back_when_the_stream_fails_after_opening() {
328 let secondary = Arc::new(Scripted::new("secondary"));
329 let fb = FallbackProvider::new(
330 Arc::new(Scripted::failing_mid_stream("primary", rate_limited())),
331 Some(target(secondary.clone())),
332 );
333 let events = collect(fb.chat_stream(request()).await.unwrap()).await;
334 assert!(
335 matches!(&events[0], Ok(StreamEvent::TextDelta { delta }) if delta == "secondary"),
336 "expected the fallback stream, got {events:?}"
337 );
338 }
339
340 #[tokio::test]
341 async fn healthy_primary_stream_is_passed_through_intact() {
342 let secondary = Arc::new(Scripted::new("secondary"));
343 let fb = FallbackProvider::new(
344 Arc::new(Scripted::new("primary")),
345 Some(target(secondary.clone())),
346 );
347 let events = collect(fb.chat_stream(request()).await.unwrap()).await;
348 assert_eq!(events.len(), 2, "buffered first event must not be dropped");
349 assert!(matches!(&events[0], Ok(StreamEvent::TextDelta { delta }) if delta == "primary"));
350 assert!(matches!(&events[1], Ok(StreamEvent::Done { .. })));
351 assert_eq!(secondary.calls.load(Ordering::SeqCst), 0);
352 }
353
354 #[tokio::test]
355 async fn capabilities_are_intersected_with_the_fallback() {
356 let primary = Scripted::new("primary").with_caps(ProviderCapabilities {
357 streaming: true,
358 vision: true,
359 max_context_tokens: 200_000,
360 ..Default::default()
361 });
362 let secondary = Scripted::new("secondary").with_caps(ProviderCapabilities {
363 streaming: true,
364 vision: false,
365 max_context_tokens: 8_000,
366 ..Default::default()
367 });
368 let fb = FallbackProvider::new(Arc::new(primary), Some(target(Arc::new(secondary))));
369 let caps = fb.capabilities();
370 assert!(caps.streaming);
371 assert!(
372 !caps.vision,
373 "must not advertise a capability the fallback lacks"
374 );
375 assert_eq!(caps.max_context_tokens, 8_000);
376 }
377
378 #[tokio::test]
379 async fn without_a_fallback_errors_propagate() {
380 let fb =
381 FallbackProvider::new(Arc::new(Scripted::failing("primary", rate_limited())), None);
382 assert!(matches!(
383 fb.chat(request()).await,
384 Err(LlmError::RateLimited { .. })
385 ));
386 }
387}
Added agent/llm/lib.rs agent/llm/lib.rs +9 −0 View agent/llm/lib.rs
@@ -0,0 +1,9 @@
1pub mod embed;
2pub mod error;
3pub mod fallback;
4pub mod provider;
5pub mod reasoning;
6pub mod register;
7pub mod skill;
8pub mod token;
9pub mod types;
Added agent/llm/provider.rs agent/llm/provider.rs +171 −0 View agent/llm/provider.rs
@@ -0,0 +1,171 @@
1use crate::error::LlmError;
2use crate::token::{
3 IMAGE_TOKENS_HIGH, IMAGE_TOKENS_LOW, MESSAGE_OVERHEAD_TOKENS, estimate_json_tokens,
4 estimate_text_tokens,
5};
6use crate::types::*;
7use std::pin::Pin;
8use tokio_stream::Stream;
9
10#[async_trait::async_trait]
11pub trait LlmProvider: Send + Sync + 'static {
12 fn provider_id(&self) -> &str;
13
14 fn model_id(&self) -> &str;
15
16 fn capabilities(&self) -> ProviderCapabilities;
17
18 async fn chat(&self, request: ChatRequest) -> Result<ChatResponse, LlmError>;
19
20 async fn chat_stream(
21 &self,
22 request: ChatRequest,
23 ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamEvent, LlmError>> + Send>>, LlmError>;
24
25 /// Estimate the input tokens this request will consume.
26 ///
27 /// The default implementation covers messages, tool-call payloads, and the
28 /// tool schemas — tool definitions are real input tokens and routinely run
29 /// to thousands, so omitting them badly under-counts tool-heavy requests.
30 ///
31 /// This is a script-aware heuristic, not a tokenizer: it is suitable for
32 /// budgeting, routing, and logging, but not for billing. Providers that
33 /// ship a real tokenizer should override it.
34 fn count_tokens(&self, request: &ChatRequest) -> usize {
35 let messages: usize = request
36 .messages
37 .iter()
38 .map(|m| {
39 let content = match &m.content {
40 MessageContent::Text(s) => estimate_text_tokens(s),
41 MessageContent::Parts(parts) => parts
42 .iter()
43 .map(|p| match p {
44 ContentPart::Text(s) => estimate_text_tokens(s),
45 // Real cost scales with pixel dimensions, which a
46 // URL does not reveal; over-estimate rather than
47 // overflow the context at the provider.
48 ContentPart::Image { detail, .. } => match detail {
49 ImageDetail::Low => IMAGE_TOKENS_LOW,
50 ImageDetail::High | ImageDetail::Auto => IMAGE_TOKENS_HIGH,
51 },
52 })
53 .sum(),
54 };
55 let tool_calls: usize = m
56 .tool_calls
57 .iter()
58 .map(|tc| estimate_text_tokens(&tc.name) + estimate_json_tokens(&tc.arguments))
59 .sum();
60 let name = m.name.as_deref().map_or(0, estimate_text_tokens);
61 content + tool_calls + name + MESSAGE_OVERHEAD_TOKENS
62 })
63 .sum();
64
65 let tools: usize = request
66 .tools
67 .iter()
68 .map(|t| {
69 estimate_text_tokens(&t.name)
70 + estimate_text_tokens(&t.description)
71 + estimate_json_tokens(&t.parameters)
72 })
73 .sum();
74
75 messages + tools
76 }
77}
78
79#[cfg(test)]
80mod tests {
81 use super::*;
82 use crate::error::LlmError;
83 use serde_json::json;
84
85 struct Stub;
86
87 #[async_trait::async_trait]
88 impl LlmProvider for Stub {
89 fn provider_id(&self) -> &str {
90 "stub"
91 }
92 fn model_id(&self) -> &str {
93 "stub-1"
94 }
95 fn capabilities(&self) -> ProviderCapabilities {
96 ProviderCapabilities::default()
97 }
98 async fn chat(&self, _: ChatRequest) -> Result<ChatResponse, LlmError> {
99 Err(LlmError::Provider("stub".into()))
100 }
101 async fn chat_stream(
102 &self,
103 _: ChatRequest,
104 ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamEvent, LlmError>> + Send>>, LlmError>
105 {
106 Err(LlmError::Provider("stub".into()))
107 }
108 }
109
110 fn user(text: &str) -> ChatMessage {
111 ChatMessage::user(text)
112 }
113
114 #[test]
115 fn tool_schemas_are_counted() {
116 let base = ChatRequest::new(vec![user("hi")]);
117 let with_tool = ChatRequest::new(vec![user("hi")]).with_tools(vec![ToolDefinition {
118 name: "search".into(),
119 description: "Search the web for a query string".into(),
120 parameters: json!({
121 "type": "object",
122 "properties": { "query": { "type": "string" } },
123 "required": ["query"]
124 }),
125 concurrency: ConcurrencyPolicy::default(),
126 }]);
127 assert!(
128 Stub.count_tokens(&with_tool) > Stub.count_tokens(&base),
129 "tool definitions must contribute to the input estimate"
130 );
131 }
132
133 #[test]
134 fn image_detail_changes_the_estimate() {
135 let img = |detail| {
136 ChatRequest::new(vec![ChatMessage {
137 role: MessageRole::User,
138 content: MessageContent::Parts(vec![ContentPart::Image {
139 url: "https://example.invalid/a.png".into(),
140 detail,
141 }]),
142 name: None,
143 tool_calls: Vec::new(),
144 tool_call_id: None,
145 }])
146 };
147 let low = Stub.count_tokens(&img(ImageDetail::Low));
148 let high = Stub.count_tokens(&img(ImageDetail::High));
149 assert!(high > low, "high detail must not be priced as low detail");
150 }
151
152 #[test]
153 fn cjk_is_not_inflated_by_byte_counting() {
154 let zh = "这是一段中文测试文本内容"; // 12 chars / 36 bytes
155 let req = ChatRequest::new(vec![user(zh)]);
156 assert_eq!(Stub.count_tokens(&req), 12 + MESSAGE_OVERHEAD_TOKENS);
157 }
158
159 #[test]
160 fn tool_call_arguments_are_counted() {
161 let mut msg = ChatMessage::assistant("");
162 msg.tool_calls = vec![ToolCall {
163 id: "c1".into(),
164 name: "search".into(),
165 arguments: json!({ "query": "a fairly long search query here" }),
166 }];
167 let with_call = ChatRequest::new(vec![msg]);
168 let without = ChatRequest::new(vec![ChatMessage::assistant("")]);
169 assert!(Stub.count_tokens(&with_call) > Stub.count_tokens(&without));
170 }
171}
Added agent/llm/reasoning.rs agent/llm/reasoning.rs +12 −0 View agent/llm/reasoning.rs
@@ -0,0 +1,12 @@
1use crate::error::LlmError;
2use crate::provider::LlmProvider;
3use crate::types::{ChatRequest, ChatResponse};
4
5#[async_trait::async_trait]
6pub trait ReasoningProvider: LlmProvider {
7 async fn chat_with_reasoning(
8 &self,
9 request: ChatRequest,
10 reasoning_budget: usize,
11 ) -> Result<(ChatResponse, String), LlmError>;
12}
Added agent/llm/register.rs agent/llm/register.rs +195 −0 View agent/llm/register.rs
@@ -0,0 +1,195 @@
1use crate::provider::LlmProvider;
2use std::collections::HashMap;
3use std::sync::Arc;
4
5/// Identifies a concrete provider+model pair.
6///
7/// Keying the registry on `provider_id` alone would let `openai/gpt-4o` and
8/// `openai/o3` silently displace one another, so both halves participate.
9#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
10pub struct ProviderKey {
11 pub provider_id: String,
12 pub model_id: String,
13}
14
15impl ProviderKey {
16 pub fn new(provider_id: impl Into<String>, model_id: impl Into<String>) -> Self {
17 Self {
18 provider_id: provider_id.into(),
19 model_id: model_id.into(),
20 }
21 }
22}
23
24impl std::fmt::Display for ProviderKey {
25 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26 write!(f, "{}/{}", self.provider_id, self.model_id)
27 }
28}
29
30#[derive(Default)]
31pub struct ProviderRegistry {
32 providers: HashMap<ProviderKey, Arc<dyn LlmProvider>>,
33}
34
35impl ProviderRegistry {
36 pub fn new() -> Self {
37 Self {
38 providers: HashMap::new(),
39 }
40 }
41
42 /// Register a provider under its `(provider_id, model_id)` pair, returning
43 /// the provider it displaced, if any.
44 pub fn register(&mut self, provider: Arc<dyn LlmProvider>) -> Option<Arc<dyn LlmProvider>> {
45 let key = ProviderKey::new(provider.provider_id(), provider.model_id());
46 self.providers.insert(key, provider)
47 }
48
49 /// Look up an exact provider+model pair.
50 ///
51 /// Returns an owned `Arc` rather than a borrow so callers can hold it
52 /// across an `.await` without keeping the registry borrowed.
53 pub fn get(&self, provider_id: &str, model_id: &str) -> Option<Arc<dyn LlmProvider>> {
54 self.providers
55 .get(&ProviderKey::new(provider_id, model_id))
56 .cloned()
57 }
58
59 /// Any registered model for `provider_id`.
60 ///
61 /// Ties are broken by `model_id` ordering so repeated calls agree with each
62 /// other rather than following `HashMap` iteration order.
63 pub fn any_for_provider(&self, provider_id: &str) -> Option<Arc<dyn LlmProvider>> {
64 self.providers
65 .iter()
66 .filter(|(k, _)| k.provider_id == provider_id)
67 .min_by(|(a, _), (b, _)| a.model_id.cmp(&b.model_id))
68 .map(|(_, v)| Arc::clone(v))
69 }
70
71 pub fn contains(&self, provider_id: &str, model_id: &str) -> bool {
72 self.providers
73 .contains_key(&ProviderKey::new(provider_id, model_id))
74 }
75
76 pub fn remove(&mut self, provider_id: &str, model_id: &str) -> Option<Arc<dyn LlmProvider>> {
77 self.providers
78 .remove(&ProviderKey::new(provider_id, model_id))
79 }
80
81 /// All registered `(provider_id, model_id)` pairs, sorted for stable output.
82 pub fn keys(&self) -> Vec<ProviderKey> {
83 let mut keys: Vec<_> = self.providers.keys().cloned().collect();
84 keys.sort();
85 keys
86 }
87
88 /// Distinct provider ids, sorted and deduplicated.
89 pub fn provider_ids(&self) -> Vec<String> {
90 let mut ids: Vec<_> = self
91 .providers
92 .keys()
93 .map(|k| k.provider_id.clone())
94 .collect();
95 ids.sort();
96 ids.dedup();
97 ids
98 }
99
100 pub fn len(&self) -> usize {
101 self.providers.len()
102 }
103
104 pub fn is_empty(&self) -> bool {
105 self.providers.is_empty()
106 }
107}
108
109#[cfg(test)]
110mod tests {
111 use super::*;
112 use crate::error::LlmError;
113 use crate::types::{ChatRequest, ChatResponse, ProviderCapabilities, StreamEvent};
114 use std::pin::Pin;
115 use tokio_stream::Stream;
116
117 struct Stub {
118 provider: String,
119 model: String,
120 }
121
122 impl Stub {
123 fn arc(provider: &str, model: &str) -> Arc<dyn LlmProvider> {
124 Arc::new(Self {
125 provider: provider.to_string(),
126 model: model.to_string(),
127 })
128 }
129 }
130
131 #[async_trait::async_trait]
132 impl LlmProvider for Stub {
133 fn provider_id(&self) -> &str {
134 &self.provider
135 }
136 fn model_id(&self) -> &str {
137 &self.model
138 }
139 fn capabilities(&self) -> ProviderCapabilities {
140 ProviderCapabilities::default()
141 }
142 async fn chat(&self, _: ChatRequest) -> Result<ChatResponse, LlmError> {
143 Err(LlmError::Provider("stub".into()))
144 }
145 async fn chat_stream(
146 &self,
147 _: ChatRequest,
148 ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamEvent, LlmError>> + Send>>, LlmError>
149 {
150 Err(LlmError::Provider("stub".into()))
151 }
152 }
153
154 #[test]
155 fn models_of_same_provider_coexist() {
156 let mut reg = ProviderRegistry::new();
157 reg.register(Stub::arc("openai", "gpt-4o"));
158 reg.register(Stub::arc("openai", "o3"));
159
160 assert_eq!(reg.len(), 2);
161 assert!(reg.get("openai", "gpt-4o").is_some());
162 assert!(reg.get("openai", "o3").is_some());
163 assert_eq!(reg.provider_ids(), vec!["openai"]);
164 }
165
166 #[test]
167 fn re_registering_same_pair_reports_displacement() {
168 let mut reg = ProviderRegistry::new();
169 assert!(reg.register(Stub::arc("openai", "o3")).is_none());
170 assert!(reg.register(Stub::arc("openai", "o3")).is_some());
171 assert_eq!(reg.len(), 1);
172 }
173
174 #[test]
175 fn lookup_miss_and_removal() {
176 let mut reg = ProviderRegistry::new();
177 reg.register(Stub::arc("openai", "o3"));
178 assert!(reg.get("openai", "nope").is_none());
179 assert!(reg.get("anthropic", "o3").is_none());
180 assert!(reg.remove("openai", "o3").is_some());
181 assert!(reg.is_empty());
182 }
183
184 #[test]
185 fn any_for_provider_is_deterministic() {
186 let mut reg = ProviderRegistry::new();
187 reg.register(Stub::arc("openai", "gpt-4o"));
188 reg.register(Stub::arc("openai", "o3"));
189 for _ in 0..32 {
190 let picked = reg.any_for_provider("openai").unwrap();
191 assert_eq!(picked.model_id(), "gpt-4o");
192 }
193 assert!(reg.any_for_provider("missing").is_none());
194 }
195}
Added agent/llm/skill.rs agent/llm/skill.rs +256 −0 View agent/llm/skill.rs
@@ -0,0 +1,256 @@
1//! Skill system — reusable prompt templates that agents can invoke.
2
3use crate::error::LlmError;
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct Skill {
9 pub name: String,
10 pub description: String,
11 pub template: String,
12 #[serde(default)]
13 pub parameters: Vec<SkillParameter>,
14}
15
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct SkillParameter {
18 pub name: String,
19 pub description: String,
20 #[serde(default)]
21 pub required: bool,
22 pub default: Option<String>,
23}
24
25impl Skill {
26 /// Render the template by substituting `{{name}}` placeholders.
27 ///
28 /// Substitution is a **single left-to-right pass**: text written to the
29 /// output is never rescanned, so a parameter value that itself contains
30 /// `{{other}}` is emitted literally rather than expanded. This is what
31 /// stops caller-supplied data from being reinterpreted as template syntax,
32 /// and it makes the result independent of `params` iteration order.
33 ///
34 /// Resolution for a placeholder `{{k}}`:
35 /// 1. `k` is declared and present in `params` → that value
36 /// 2. `k` is declared with a `default` → the default
37 /// 3. `k` is declared and optional → empty string
38 /// 4. `k` is undeclared → emitted verbatim, including the braces
39 ///
40 /// Keys in `params` not declared on the skill are ignored: only the
41 /// declared parameter set is substitutable.
42 pub fn render(&self, params: &HashMap<String, String>) -> Result<String, LlmError> {
43 for param in &self.parameters {
44 if param.required && !params.contains_key(¶m.name) && param.default.is_none() {
45 return Err(LlmError::SkillMissingParameter {
46 skill: self.name.clone(),
47 parameter: param.name.clone(),
48 });
49 }
50 }
51
52 let mut out = String::with_capacity(self.template.len());
53 let mut rest = self.template.as_str();
54
55 while let Some(open) = rest.find("{{") {
56 let after_open = &rest[open + 2..];
57 let Some(close) = after_open.find("}}") else {
58 break; // Unclosed brace: the remainder is literal text.
59 };
60 let end = open + 2 + close + 2;
61
62 match self.resolve(after_open[..close].trim(), params) {
63 Some(value) => {
64 out.push_str(&rest[..open]);
65 out.push_str(value);
66 }
67 None => out.push_str(&rest[..end]),
68 }
69 rest = &rest[end..];
70 }
71
72 out.push_str(rest);
73 Ok(out)
74 }
75
76 fn resolve<'a>(&'a self, key: &str, params: &'a HashMap<String, String>) -> Option<&'a str> {
77 let declared = self.parameters.iter().find(|p| p.name == key)?;
78 Some(
79 params
80 .get(key)
81 .or(declared.default.as_ref())
82 .map_or("", String::as_str),
83 )
84 }
85}
86
87#[derive(Debug, Default)]
88pub struct SkillRegistry {
89 skills: HashMap<String, Skill>,
90}
91
92impl SkillRegistry {
93 pub fn new() -> Self {
94 Self {
95 skills: HashMap::new(),
96 }
97 }
98
99 /// Register a skill, returning the previous skill of the same name if one
100 /// was displaced.
101 pub fn register(&mut self, skill: Skill) -> Option<Skill> {
102 self.skills.insert(skill.name.clone(), skill)
103 }
104
105 /// Get a skill by name.
106 pub fn get(&self, name: &str) -> Option<&Skill> {
107 self.skills.get(name)
108 }
109
110 /// Look up and render a skill in one step.
111 pub fn render(&self, name: &str, params: &HashMap<String, String>) -> Result<String, LlmError> {
112 self.get(name)
113 .ok_or_else(|| LlmError::SkillNotFound(name.to_string()))?
114 .render(params)
115 }
116
117 /// List all skill names.
118 pub fn names(&self) -> Vec<String> {
119 self.skills.keys().cloned().collect()
120 }
121
122 /// Number of registered skills.
123 pub fn count(&self) -> usize {
124 self.skills.len()
125 }
126
127 pub fn is_empty(&self) -> bool {
128 self.skills.is_empty()
129 }
130}
131
132#[cfg(test)]
133mod tests {
134 use super::*;
135
136 fn param(name: &str, required: bool, default: Option<&str>) -> SkillParameter {
137 SkillParameter {
138 name: name.to_string(),
139 description: String::new(),
140 required,
141 default: default.map(str::to_string),
142 }
143 }
144
145 fn skill(template: &str, parameters: Vec<SkillParameter>) -> Skill {
146 Skill {
147 name: "s".to_string(),
148 description: String::new(),
149 template: template.to_string(),
150 parameters,
151 }
152 }
153
154 #[test]
155 fn substitutes_declared_parameters() {
156 let s = skill("hello {{who}}!", vec![param("who", true, None)]);
157 let params = HashMap::from([("who".to_string(), "world".to_string())]);
158 assert_eq!(s.render(¶ms).unwrap(), "hello world!");
159 }
160
161 #[test]
162 fn does_not_expand_placeholders_injected_through_values() {
163 let s = skill(
164 "{{a}}|{{b}}",
165 vec![param("a", true, None), param("b", true, None)],
166 );
167 let params = HashMap::from([
168 ("a".to_string(), "{{b}}".to_string()),
169 ("b".to_string(), "SECRET".to_string()),
170 ]);
171 // The `{{b}}` arriving as a's *value* must survive verbatim; the old
172 // two-pass implementation expanded it to "SECRET|SECRET".
173 assert_eq!(s.render(¶ms).unwrap(), "{{b}}|SECRET");
174 }
175
176 #[test]
177 fn render_is_deterministic_regardless_of_map_order() {
178 let s = skill(
179 "{{a}}{{b}}{{c}}",
180 vec![
181 param("a", true, None),
182 param("b", true, None),
183 param("c", true, None),
184 ],
185 );
186 let params = HashMap::from([
187 ("a".to_string(), "{{c}}".to_string()),
188 ("b".to_string(), "{{a}}".to_string()),
189 ("c".to_string(), "z".to_string()),
190 ]);
191 let first = s.render(¶ms).unwrap();
192 for _ in 0..64 {
193 assert_eq!(s.render(¶ms).unwrap(), first);
194 }
195 assert_eq!(first, "{{c}}{{a}}z");
196 }
197
198 #[test]
199 fn undeclared_placeholder_is_left_verbatim() {
200 let s = skill("{{known}} {{unknown}}", vec![param("known", true, None)]);
201 let params = HashMap::from([
202 ("known".to_string(), "ok".to_string()),
203 ("unknown".to_string(), "injected".to_string()),
204 ]);
205 assert_eq!(s.render(¶ms).unwrap(), "ok {{unknown}}");
206 }
207
208 #[test]
209 fn falls_back_to_default_then_empty() {
210 let s = skill(
211 "[{{d}}][{{o}}]",
212 vec![param("d", false, Some("dv")), param("o", false, None)],
213 );
214 assert_eq!(s.render(&HashMap::new()).unwrap(), "[dv][]");
215 }
216
217 #[test]
218 fn missing_required_parameter_errors() {
219 let s = skill("{{x}}", vec![param("x", true, None)]);
220 assert!(matches!(
221 s.render(&HashMap::new()),
222 Err(LlmError::SkillMissingParameter { .. })
223 ));
224 }
225
226 #[test]
227 fn unclosed_placeholder_is_passed_through() {
228 let s = skill("a {{b", vec![param("b", false, None)]);
229 assert_eq!(s.render(&HashMap::new()).unwrap(), "a {{b");
230 }
231
232 #[test]
233 fn multibyte_template_does_not_split_char_boundaries() {
234 let s = skill("前缀{{名字}}后缀", vec![param("名字", true, None)]);
235 let params = HashMap::from([("名字".to_string(), "世界".to_string())]);
236 assert_eq!(s.render(¶ms).unwrap(), "前缀世界后缀");
237 }
238
239 #[test]
240 fn registry_render_reports_unknown_skill() {
241 let reg = SkillRegistry::new();
242 assert!(matches!(
243 reg.render("nope", &HashMap::new()),
244 Err(LlmError::SkillNotFound(_))
245 ));
246 }
247
248 #[test]
249 fn registry_register_returns_displaced_skill() {
250 let mut reg = SkillRegistry::new();
251 assert!(reg.register(skill("v1", vec![])).is_none());
252 let old = reg.register(skill("v2", vec![])).expect("displaced");
253 assert_eq!(old.template, "v1");
254 assert_eq!(reg.count(), 1);
255 }
256}
Added agent/llm/token.rs agent/llm/token.rs +161 −0 View agent/llm/token.rs
@@ -0,0 +1,161 @@
1use serde::{Deserialize, Serialize};
2
3#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
4pub struct TokenUsageStats {
5 pub input_tokens: usize,
6 pub output_tokens: usize,
7 pub reasoning_tokens: Option<usize>,
8}
9
10impl TokenUsageStats {
11 pub fn new(input: usize, output: usize) -> Self {
12 Self {
13 input_tokens: input,
14 output_tokens: output,
15 reasoning_tokens: None,
16 }
17 }
18
19 pub fn with_reasoning(mut self, tokens: usize) -> Self {
20 self.reasoning_tokens = Some(tokens);
21 self
22 }
23
24 pub fn total(&self) -> usize {
25 self.input_tokens
26 .saturating_add(self.output_tokens)
27 .saturating_add(self.reasoning_tokens.unwrap_or(0))
28 }
29
30 pub fn merge(&mut self, other: &TokenUsageStats) {
31 self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
32 self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
33 match (self.reasoning_tokens, other.reasoning_tokens) {
34 (Some(a), Some(b)) => self.reasoning_tokens = Some(a.saturating_add(b)),
35 (None, Some(b)) => self.reasoning_tokens = Some(b),
36 _ => {}
37 }
38 }
39}
40
41/// Average characters per token for Latin-script text.
42const LATIN_CHARS_PER_TOKEN: usize = 4;
43
44/// Tokens for a low-detail image (OpenAI's documented flat cost).
45pub const IMAGE_TOKENS_LOW: usize = 85;
46
47/// Conservative upper estimate for a high-detail image (base + 4 tiles).
48///
49/// The real cost depends on pixel dimensions, which are not knowable from a
50/// URL. Over-estimating is the safe direction: it trips a local budget check
51/// rather than letting the request fail at the provider with a hard context
52/// overflow.
53pub const IMAGE_TOKENS_HIGH: usize = 765;
54
55/// Rough per-message framing overhead (role markers, separators).
56pub const MESSAGE_OVERHEAD_TOKENS: usize = 4;
57
58/// Estimate the token count of a string.
59///
60/// CJK characters count as roughly one token each; other scripts at
61/// [`LATIN_CHARS_PER_TOKEN`] characters per token. This counts **characters,
62/// not bytes** — counting UTF-8 bytes over-estimates CJK text by ~3x, which
63/// causes spurious budget rejections for Chinese/Japanese/Korean input.
64///
65/// This is a heuristic for budgeting and logging only. It is **not** accurate
66/// enough for billing and must not be the sole gate on a hard context limit;
67/// prefer a provider-supplied tokenizer wherever one is available.
68pub fn estimate_text_tokens(text: &str) -> usize {
69 let (cjk, other) = text.chars().fold((0usize, 0usize), |(cjk, other), c| {
70 if is_cjk(c) {
71 (cjk + 1, other)
72 } else {
73 (cjk, other + 1)
74 }
75 });
76 cjk + other.div_ceil(LATIN_CHARS_PER_TOKEN)
77}
78
79/// Estimate tokens for a JSON value (tool schemas, tool-call arguments).
80pub fn estimate_json_tokens(value: &serde_json::Value) -> usize {
81 serde_json::to_string(value)
82 .map(|s| estimate_text_tokens(&s))
83 .unwrap_or(0)
84}
85
86/// Characters that are typically one token each rather than ~4.
87fn is_cjk(c: char) -> bool {
88 matches!(c as u32,
89 0x2E80..=0x303F // CJK radicals and punctuation
90 | 0x3040..=0x30FF // Hiragana, Katakana
91 | 0x3400..=0x4DBF // CJK Unified Ext A
92 | 0x4E00..=0x9FFF // CJK Unified
93 | 0xA960..=0xA97F // Hangul Jamo Ext A
94 | 0xAC00..=0xD7FF // Hangul syllables
95 | 0xF900..=0xFAFF // CJK compatibility
96 | 0xFF00..=0xFFEF // Halfwidth and fullwidth forms
97 | 0x20000..=0x3FFFF // CJK Unified Ext B and beyond
98 )
99}
100
101#[cfg(test)]
102mod tests {
103 use super::*;
104
105 #[test]
106 fn total_and_merge_accumulate_reasoning() {
107 let mut a = TokenUsageStats::new(10, 5).with_reasoning(3);
108 assert_eq!(a.total(), 18);
109 a.merge(&TokenUsageStats::new(1, 1).with_reasoning(1));
110 assert_eq!(
111 a,
112 TokenUsageStats {
113 input_tokens: 11,
114 output_tokens: 6,
115 reasoning_tokens: Some(4),
116 }
117 );
118 }
119
120 #[test]
121 fn merge_adopts_reasoning_when_absent_locally() {
122 let mut a = TokenUsageStats::new(1, 1);
123 a.merge(&TokenUsageStats::new(0, 0).with_reasoning(7));
124 assert_eq!(a.reasoning_tokens, Some(7));
125 }
126
127 #[test]
128 fn merge_keeps_reasoning_when_other_has_none() {
129 let mut a = TokenUsageStats::new(1, 1).with_reasoning(7);
130 a.merge(&TokenUsageStats::new(1, 1));
131 assert_eq!(a.reasoning_tokens, Some(7));
132 }
133
134 #[test]
135 fn total_saturates_instead_of_overflowing() {
136 assert_eq!(TokenUsageStats::new(usize::MAX, 10).total(), usize::MAX);
137 }
138
139 #[test]
140 fn cjk_counts_by_character_not_byte() {
141 // 12 chars, 36 UTF-8 bytes. Byte-based `len() / 4` treats this as 9
142 // tokens' worth of budget for text that really costs ~12.
143 let zh = "这是一段中文测试文本内容";
144 assert_eq!(zh.chars().count(), 12);
145 assert_eq!(zh.len(), 36);
146 assert_eq!(estimate_text_tokens(zh), 12);
147 }
148
149 #[test]
150 fn latin_uses_four_chars_per_token() {
151 assert_eq!(estimate_text_tokens("abcdefgh"), 2);
152 assert_eq!(estimate_text_tokens("abcde"), 2); // rounds up
153 assert_eq!(estimate_text_tokens(""), 0);
154 }
155
156 #[test]
157 fn mixed_script_counts_each_run_appropriately() {
158 // 4 latin chars (1 token) + 2 CJK (2 tokens)
159 assert_eq!(estimate_text_tokens("abcd中文"), 3);
160 }
161}
Added agent/llm/types.rs agent/llm/types.rs +13 −0 View agent/llm/types.rs
@@ -0,0 +1,13 @@
1pub mod capabilities;
2pub mod message;
3pub mod request;
4pub mod response;
5pub mod stream;
6pub mod tools;
7
8pub use capabilities::*;
9pub use message::*;
10pub use request::*;
11pub use response::*;
12pub use stream::*;
13pub use tools::*;