use candid::Principal;
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use serde_json::{Map, json};
use std::collections::HashMap;
use crate::{Json, json::normalize_strict_schema};
pub use ic_auth_types::{ByteArrayB64, ByteBufB64, Xid};
mod completion;
mod content;
mod document;
mod resource;
mod text;
pub use completion::*;
pub use content::*;
pub use document::*;
pub use resource::*;
pub use text::*;
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct AgentInput {
pub name: String,
pub prompt: String,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub resources: Vec<Resource>,
#[serde(skip_serializing_if = "Option::is_none")]
pub topics: Option<Vec<String>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub meta: Option<RequestMeta>,
}
impl AgentInput {
pub fn new(name: String, prompt: String) -> Self {
Self {
name,
prompt,
resources: Vec::new(),
topics: None,
meta: None,
}
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub enum PromptCommand {
#[default]
Ping,
Plain {
prompt: String,
},
Command {
command: String,
prompt: String,
},
}
impl From<String> for PromptCommand {
fn from(prompt: String) -> Self {
let trimmed = prompt.trim();
if trimmed.is_empty() {
return Self::Ping;
}
let Some(stripped) = trimmed.strip_prefix('/') else {
return Self::Plain { prompt };
};
let command_end = stripped.find(char::is_whitespace).unwrap_or(stripped.len());
let command = stripped[..command_end].to_ascii_lowercase();
if command.is_empty() {
return Self::Plain { prompt };
}
if command == "ping" {
return Self::Ping;
}
Self::Command { command, prompt }
}
}
impl PromptCommand {
pub fn command_argument(&self) -> Option<&str> {
let Self::Command { command, prompt } = self else {
return None;
};
let trimmed = prompt.trim();
let Some(stripped) = trimmed.strip_prefix('/') else {
return Some(trimmed);
};
let command_end = stripped.find(char::is_whitespace).unwrap_or(stripped.len());
if !stripped[..command_end].eq_ignore_ascii_case(command) {
return Some(trimmed);
}
Some(stripped[command_end..].trim())
}
}
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct AgentOutput {
pub content: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub thoughts: Option<String>,
pub usage: Usage,
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub tools_usage: HashMap<String, Usage>,
#[serde(skip_serializing_if = "Option::is_none")]
pub failed_reason: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub tool_calls: Vec<ToolCall>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub chat_history: Vec<Message>,
#[serde(skip)]
pub raw_history: Vec<Json>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub artifacts: Vec<Resource>,
#[serde(skip_serializing_if = "Option::is_none")]
pub conversation: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub session: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
}
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct PartialAgentOutput {
pub content: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub thoughts: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub failed_reason: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub conversation: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub session: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
}
impl AgentOutput {
pub fn into_tool_output(self) -> ToolOutput<Json> {
let AgentOutput {
content,
thoughts,
usage,
tools_usage,
failed_reason,
artifacts,
conversation,
session,
model,
..
} = self;
let failed_reason = failed_reason.filter(|reason| !reason.trim().is_empty());
let has_metadata = thoughts.is_some()
|| failed_reason.is_some()
|| conversation.is_some()
|| session.is_some()
|| model.is_some();
let is_error = failed_reason.as_ref().map(|_| true);
let output = if has_metadata {
json!(PartialAgentOutput {
content,
thoughts,
failed_reason,
conversation,
session,
model,
})
} else {
serde_json::from_str::<Json>(&content).unwrap_or(Json::String(content))
};
ToolOutput {
output,
is_error,
artifacts,
usage,
tools_usage,
}
}
}
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct ToolInput<T> {
pub name: String,
pub args: T,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub resources: Vec<Resource>,
#[serde(skip_serializing_if = "Option::is_none")]
pub meta: Option<RequestMeta>,
}
impl<T> ToolInput<T> {
pub fn new(name: String, args: T) -> Self {
Self {
name,
args,
resources: Vec::new(),
meta: None,
}
}
}
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct ToolOutput<T> {
pub output: T,
#[serde(skip_serializing_if = "Option::is_none")]
pub is_error: Option<bool>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub artifacts: Vec<Resource>,
pub usage: Usage,
#[serde(default, skip_serializing_if = "HashMap::is_empty")]
pub tools_usage: HashMap<String, Usage>,
}
impl<T> ToolOutput<T> {
pub fn new(output: T) -> Self {
Self {
output,
is_error: None,
artifacts: Vec::new(),
usage: Usage::default(),
tools_usage: HashMap::new(),
}
}
}
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct RequestMeta {
#[serde(skip_serializing_if = "Option::is_none")]
pub engine: Option<Principal>,
#[serde(skip_serializing_if = "Option::is_none")]
pub user: Option<String>,
#[serde(flatten)]
#[serde(skip_serializing_if = "Map::is_empty")]
pub extra: Map<String, Json>,
}
impl RequestMeta {
pub fn get_extra_as<T>(&self, key: &str) -> Option<T>
where
T: DeserializeOwned,
{
self.extra
.get(key)
.and_then(|value| T::deserialize(value).ok())
}
}
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
pub struct Usage {
pub input_tokens: u64,
pub output_tokens: u64,
#[serde(default)]
pub cached_tokens: u64,
pub requests: u64,
}
impl Usage {
pub fn accumulate(&mut self, other: &Usage) {
self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
self.cached_tokens = self.cached_tokens.saturating_add(other.cached_tokens);
self.requests = self.requests.saturating_add(other.requests);
}
}
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct ToolCall {
pub name: String,
pub args: Json,
#[serde(skip_serializing_if = "Option::is_none")]
pub result: Option<ToolOutput<Json>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub call_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub remote_id: Option<Principal>,
}
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct Function {
pub definition: FunctionDefinition,
pub supported_resource_tags: Vec<String>,
}
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct FunctionDefinition {
pub name: String,
pub description: String,
pub parameters: Json,
#[serde(skip_serializing_if = "Option::is_none")]
pub strict: Option<bool>,
}
impl FunctionDefinition {
pub fn name_with_prefix(mut self, prefix: &str) -> Self {
self.name = format!("{}{}", prefix, self.name);
self
}
pub fn normalize_strict_parameters(mut self) -> Self {
if self.strict.unwrap_or_default() {
self.parameters = normalize_strict_schema(self.parameters);
}
self
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_support::resource;
#[test]
fn test_agent_and_tool_constructors_default_optional_fields() {
let agent = AgentInput::new("planner".into(), "summarize this".into());
assert_eq!(agent.name, "planner");
assert_eq!(agent.prompt, "summarize this");
assert!(agent.resources.is_empty());
assert!(agent.topics.is_none());
assert!(agent.meta.is_none());
let tool = ToolInput::new("sum".into(), json!({"x": 1, "y": 2}));
assert_eq!(tool.name, "sum");
assert_eq!(tool.args, json!({"x": 1, "y": 2}));
assert!(tool.resources.is_empty());
assert!(tool.meta.is_none());
let output = ToolOutput::new(json!("ok"));
assert_eq!(output.output, json!("ok"));
assert!(output.artifacts.is_empty());
assert_eq!(output.usage.requests, 0);
assert!(output.tools_usage.is_empty());
}
#[test]
fn test_prompt_command_from_string_variants() {
assert_eq!(PromptCommand::from("".to_string()), PromptCommand::Ping);
assert_eq!(
PromptCommand::from(" /PING ".to_string()),
PromptCommand::Ping
);
assert_eq!(
PromptCommand::from("hello".to_string()),
PromptCommand::Plain {
prompt: "hello".into(),
}
);
assert_eq!(
PromptCommand::from("/Status show details".to_string()),
PromptCommand::Command {
command: "status".into(),
prompt: "/Status show details".into(),
}
);
assert_eq!(
PromptCommand::from("/help".to_string()),
PromptCommand::Command {
command: "help".into(),
prompt: "/help".into(),
}
);
assert_eq!(
PromptCommand::from("/ping now".to_string()),
PromptCommand::Ping
);
assert_eq!(
PromptCommand::from("/".to_string()),
PromptCommand::Plain { prompt: "/".into() }
);
assert_eq!(
PromptCommand::from("/ arg".to_string()),
PromptCommand::Plain {
prompt: "/ arg".into(),
}
);
let stop = PromptCommand::from("/stop 停止当前任务,保留会话".to_string());
assert_eq!(stop.command_argument(), Some("停止当前任务,保留会话"));
let manual = PromptCommand::Command {
command: "cancel".into(),
prompt: "取消当前任务".into(),
};
assert_eq!(manual.command_argument(), Some("取消当前任务"));
assert_eq!(PromptCommand::Ping.command_argument(), None);
let command = PromptCommand::from("/ÄRGER now".to_string());
assert_eq!(
command,
PromptCommand::Command {
command: "Ärger".into(),
prompt: "/ÄRGER now".into(),
}
);
assert_eq!(command.command_argument(), Some("now"));
}
#[test]
fn test_agent_output_into_tool_output_handles_json_plain_text_and_metadata() {
let mut tools_usage = HashMap::new();
tools_usage.insert(
"sum".into(),
Usage {
requests: 1,
..Default::default()
},
);
let output = AgentOutput {
content: r#"{"ok":true}"#.into(),
usage: Usage {
input_tokens: 2,
output_tokens: 1,
requests: 1,
..Default::default()
},
tools_usage: tools_usage.clone(),
artifacts: vec![resource(7, &["text"])],
..Default::default()
}
.into_tool_output();
assert_eq!(output.output, json!({"ok": true}));
assert_eq!(output.artifacts.len(), 1);
assert_eq!(output.artifacts[0]._id, 7);
assert_eq!(output.usage.input_tokens, 2);
assert_eq!(output.tools_usage.get("sum").unwrap().requests, 1);
let output = AgentOutput {
content: "not-json".into(),
thoughts: Some("thinking".into()),
session: Some("session-1".into()),
model: Some("test-model".into()),
..Default::default()
}
.into_tool_output();
assert_eq!(
output.output,
json!({
"content": "not-json",
"thoughts": "thinking",
"session": "session-1",
"model": "test-model"
})
);
let output = AgentOutput {
content: "still-not-json".into(),
..Default::default()
}
.into_tool_output();
assert_eq!(output.output, json!("still-not-json"));
}
#[test]
fn test_agent_output_into_tool_output_normalizes_blank_failed_reason() {
let output = AgentOutput {
content: r#"{"ok":true}"#.into(),
failed_reason: Some(" ".into()),
..Default::default()
}
.into_tool_output();
assert_eq!(output.is_error, None);
assert_eq!(output.output, json!({"ok": true}));
let output = AgentOutput {
content: "boom".into(),
failed_reason: Some("boom".into()),
..Default::default()
}
.into_tool_output();
assert_eq!(output.is_error, Some(true));
assert_eq!(output.output.get("failed_reason").unwrap(), &json!("boom"));
}
#[test]
fn test_request_meta_get_extra_as_and_usage_accumulate() {
let mut extra = Map::new();
extra.insert("numbers".into(), json!([1, 2, 3]));
extra.insert("flag".into(), json!(true));
let meta = RequestMeta {
extra,
..Default::default()
};
assert_eq!(
meta.get_extra_as::<Vec<u64>>("numbers"),
Some(vec![1, 2, 3])
);
assert_eq!(meta.get_extra_as::<bool>("flag"), Some(true));
assert_eq!(meta.get_extra_as::<String>("missing"), None);
let mut usage = Usage {
input_tokens: u64::MAX - 1,
output_tokens: 2,
cached_tokens: 3,
requests: u64::MAX,
};
let other = Usage {
input_tokens: 10,
output_tokens: 5,
cached_tokens: u64::MAX,
requests: 1,
};
usage.accumulate(&other);
assert_eq!(usage.input_tokens, u64::MAX);
assert_eq!(usage.output_tokens, 7);
assert_eq!(usage.cached_tokens, u64::MAX);
assert_eq!(usage.requests, u64::MAX);
}
#[test]
fn test_request_meta_extra_flatten_serde() {
let meta = RequestMeta {
engine: None,
user: None,
extra: Map::new(),
};
let v = serde_json::to_value(&meta).unwrap();
assert_eq!(v, serde_json::json!({}));
let mut extra = Map::new();
extra.insert("foo".into(), serde_json::json!("bar"));
extra.insert("n".into(), serde_json::json!(1));
extra.insert("obj".into(), serde_json::json!({"x": true}));
let meta2 = RequestMeta {
engine: Some(Principal::from_text("aaaaa-aa").unwrap()),
user: Some("alice".into()),
extra,
};
let v2 = serde_json::to_value(&meta2).unwrap();
assert_eq!(v2.get("engine").unwrap(), "aaaaa-aa");
assert_eq!(v2.get("user").unwrap(), "alice");
assert_eq!(v2.get("foo").unwrap(), "bar");
assert_eq!(v2.get("n").unwrap(), 1);
assert_eq!(v2.get("obj").unwrap(), &serde_json::json!({"x": true}));
assert!(v2.get("extra").is_none());
let input = serde_json::json!({
"engine": "aaaaa-aa",
"user": "bob",
"k1": "v1",
"k2": 2,
"nested": {"a": 1}
});
let back: RequestMeta = serde_json::from_value(input).unwrap();
assert_eq!(back.engine.unwrap().to_text(), "aaaaa-aa");
assert_eq!(back.user.as_deref(), Some("bob"));
assert_eq!(back.extra.get("k1").unwrap(), "v1");
assert_eq!(back.extra.get("k2").unwrap(), 2);
assert_eq!(
back.extra.get("nested").unwrap(),
&serde_json::json!({"a": 1})
);
let back2: RequestMeta = serde_json::from_value(v2).unwrap();
assert_eq!(back2.engine.unwrap().to_text(), "aaaaa-aa");
assert_eq!(back2.user.as_deref(), Some("alice"));
assert_eq!(back2.extra.get("foo").unwrap(), "bar");
assert_eq!(back2.extra.get("n").unwrap(), 1);
assert_eq!(
back2.extra.get("obj").unwrap(),
&serde_json::json!({"x": true})
);
}
#[test]
fn test_function_definition_helpers() {
let definition = FunctionDefinition {
name: "search".into(),
description: "Find documents".into(),
parameters: json!({
"type": "object",
"properties": {},
"required": [],
"additionalProperties": false
}),
strict: Some(true),
}
.name_with_prefix("tool_");
assert_eq!(definition.name, "tool_search");
assert_eq!(definition.description, "Find documents");
}
}