anda_core 0.11.6

//! Core data models and traits for the AI agent system.
//!
//! This module defines the fundamental data structures and interfaces used throughout the AI agent system.
//! It includes:
//! - Core message and conversation structures ([`AgentOutput`], [`Message`], [`ToolCall`]).
//! - Function definition and tooling support ([`FunctionDefinition`]).
//! - Knowledge and document handling ([`Document`], [`Documents`]).
//! - Completion request and response structures ([`CompletionRequest`], [`Embedding`]).
//! - Core AI capabilities traits ([`CompletionFeatures`], [`EmbeddingFeatures`]).

use candid::Principal;
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use serde_json::{Map, json};
use std::collections::BTreeMap;

use crate::Json;
pub use ic_auth_types::{ByteArrayB64, ByteBufB64, Xid};

mod completion;
mod resource;

pub use completion::*;
pub use resource::*;

/// Represents a request to an agent for processing.
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct AgentInput {
    /// agent name, use default agent if empty.
    pub name: String,

    /// agent prompt or message.
    pub prompt: String,

    /// The resources to process by the agent.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub resources: Vec<Resource>,

    /// The topics for the agent request.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub topics: Option<Vec<String>>,

    /// The metadata for the agent request
    #[serde(skip_serializing_if = "Option::is_none")]
    pub meta: Option<RequestMeta>,
}

impl AgentInput {
    /// Creates a new agent input with the given name and prompt.
    pub fn new(name: String, prompt: String) -> Self {
        Self {
            name,
            prompt,
            resources: Vec::new(),
            topics: None,
            meta: None,
        }
    }
}

/// Represents the output of an agent execution.
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct AgentOutput {
    /// The output content from the agent, may be empty.
    pub content: String,

    /// The intermediate thinking process of the agent, may be empty if the agent does not provide it or the execution is already finished.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub thinking: Option<String>,

    /// The usage statistics for the agent execution.
    pub usage: Usage,

    /// Indicates failure reason if present, None means successful execution.
    /// Should be None when finish_reason is "stop" or "tool_calls".
    #[serde(skip_serializing_if = "Option::is_none")]
    pub failed_reason: Option<String>,

    /// Tool calls returned by the LLM function calling.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub tool_calls: Vec<ToolCall>,

    /// The history of the conversation.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub chat_history: Vec<Message>,

    /// raw_history is the model specialized history of the conversation,
    /// will be included in `ctx.completion` response,
    /// but not be included in the engine response.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub raw_history: Vec<Json>,

    /// A collection of artifacts generated by the agent during the execution of the task.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub artifacts: Vec<Resource>,

    /// The conversation ID.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub conversation: Option<u64>,

    /// The model used by the agent.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub model: Option<String>,
}

fn deserialize_content<'de, D>(deserializer: D) -> Result<Vec<ContentPart>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let value = Json::deserialize(deserializer)?;
    match value {
        Json::String(s) => Ok(vec![ContentPart::Text { text: s }]),
        Json::Array(_) => Vec::<ContentPart>::deserialize(value).map_err(serde::de::Error::custom),
        _ => Err(serde::de::Error::custom(
            "expected a string or array for content",
        )),
    }
}

/// Represents a message send to LLM for completion.
#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, Eq)]
pub struct Message {
    /// Message role: "system", "user", "assistant", "tool".
    pub role: String,

    /// The content of the message
    #[serde(default, deserialize_with = "deserialize_content")]
    pub content: Vec<ContentPart>,

    /// An optional name for the participant. Provides the model information to differentiate between participants of the same role.
    /// This field is not used by the model.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,

    /// The user ID of the message sender.
    /// This field is not used by the model.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub user: Option<Principal>,

    /// The timestamp of the message.
    /// This field is not used by the model.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub timestamp: Option<u64>,
}

impl Message {
    pub fn text(&self) -> Option<String> {
        let mut texts: Vec<&str> = Vec::new();
        for part in &self.content {
            if let ContentPart::Text { text } = part {
                texts.push(text);
            }
        }
        if texts.is_empty() {
            return None;
        }
        Some(texts.join("\n"))
    }

    pub fn tool_calls(&self) -> Vec<ToolCall> {
        let mut tool_calls: Vec<ToolCall> = Vec::new();
        for part in &self.content {
            if let ContentPart::ToolCall {
                name,
                args,
                call_id,
            } = part
            {
                tool_calls.push(ToolCall {
                    name: name.clone(),
                    args: args.clone(),
                    call_id: call_id.clone(),
                    result: None,
                    remote_id: None,
                });
            }
        }
        tool_calls
    }

    pub fn prune_content(&mut self) -> usize {
        let original_len = self.content.len();
        self.content.retain(|part| {
            matches!(
                part,
                ContentPart::Text { .. }
                    | ContentPart::Reasoning { .. }
                    | ContentPart::Action { .. }
            )
        });
        let pruned = original_len - self.content.len();
        if pruned > 0 {
            self.content.push(ContentPart::Text {
                text: format!(
                    "[{} items (tool calls or files) pruned due to limits]",
                    pruned
                ),
            });
        }
        pruned
    }
}

#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
#[serde(tag = "type", rename_all_fields = "camelCase")]
pub enum ContentPart {
    Text {
        text: String,
    },
    Reasoning {
        text: String,
    },
    FileData {
        file_uri: String,

        #[serde(skip_serializing_if = "Option::is_none")]
        mime_type: Option<String>,
    },
    InlineData {
        mime_type: String,
        data: ByteBufB64,
    },
    ToolCall {
        name: String,
        args: Json,

        #[serde(skip_serializing_if = "Option::is_none")]
        call_id: Option<String>,
    },
    ToolOutput {
        name: String,
        output: Json,

        #[serde(skip_serializing_if = "Option::is_none")]
        call_id: Option<String>,

        #[serde(skip_serializing_if = "Option::is_none")]
        remote_id: Option<Principal>,
    },
    Action {
        name: String,
        payload: Json,

        #[serde(skip_serializing_if = "Option::is_none")]
        recipients: Option<Vec<Principal>>,

        #[serde(skip_serializing_if = "Option::is_none")]
        signature: Option<ByteBufB64>,
    },
    #[serde(untagged)]
    Any(Json),
}

/// Converts a content part with inline data to a data URL string.
///
/// https://developer.mozilla.org/en-US/docs/Web/URI/Reference/Schemes/data
pub fn part_to_data_url(data: &ByteBufB64, mime_type: Option<&String>) -> String {
    format!(
        "data:{};base64,{}",
        mime_type.map(|m| m.as_str()).unwrap_or(""),
        data.to_base64()
    )
}

impl<'de> Deserialize<'de> for ContentPart {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        let value = Json::deserialize(deserializer)?;
        match &value {
            Json::String(s) => Ok(ContentPart::Text { text: s.clone() }),
            Json::Object(map)
                if matches!(
                    map.get("type").and_then(|t| t.as_str()),
                    Some(
                        "Text"
                            | "Reasoning"
                            | "FileData"
                            | "InlineData"
                            | "ToolCall"
                            | "ToolOutput"
                            | "Action"
                    )
                ) =>
            {
                #[derive(Deserialize)]
                #[serde(tag = "type", rename_all_fields = "camelCase")]
                enum Helper {
                    Text {
                        text: String,
                    },
                    Reasoning {
                        text: String,
                    },
                    FileData {
                        file_uri: String,
                        mime_type: Option<String>,
                    },
                    InlineData {
                        mime_type: String,
                        data: ByteBufB64,
                    },
                    ToolCall {
                        name: String,
                        args: Json,
                        call_id: Option<String>,
                    },
                    ToolOutput {
                        name: String,
                        output: Json,
                        call_id: Option<String>,
                        remote_id: Option<Principal>,
                    },
                    Action {
                        name: String,
                        payload: Json,
                        recipients: Option<Vec<Principal>>,
                        signature: Option<ByteBufB64>,
                    },
                }

                match serde_json::from_value::<Helper>(value) {
                    Ok(h) => Ok(match h {
                        Helper::Text { text } => ContentPart::Text { text },
                        Helper::Reasoning { text } => ContentPart::Reasoning { text },
                        Helper::FileData {
                            file_uri,
                            mime_type,
                        } => ContentPart::FileData {
                            file_uri,
                            mime_type,
                        },
                        Helper::InlineData { mime_type, data } => {
                            ContentPart::InlineData { mime_type, data }
                        }
                        Helper::ToolCall {
                            name,
                            args,
                            call_id,
                        } => ContentPart::ToolCall {
                            name,
                            args,
                            call_id,
                        },
                        Helper::ToolOutput {
                            name,
                            output,
                            call_id,
                            remote_id,
                        } => ContentPart::ToolOutput {
                            name,
                            output,
                            call_id,
                            remote_id,
                        },
                        Helper::Action {
                            name,
                            payload,
                            recipients,
                            signature,
                        } => ContentPart::Action {
                            name,
                            payload,
                            recipients,
                            signature,
                        },
                    }),
                    Err(_) => Err(serde::de::Error::custom("invalid ContentPart")),
                }
            }
            _ => Ok(ContentPart::Any(value)),
        }
    }
}

impl From<String> for ContentPart {
    fn from(text: String) -> Self {
        Self::Text { text }
    }
}

impl From<Json> for ContentPart {
    fn from(val: Json) -> Self {
        if let Json::Object(map) = &val
            && let Some(t) = map.get("type").and_then(|x| x.as_str())
        {
            match t {
                "Text" | "Reasoning" | "FileData" | "InlineData" | "ToolCall" | "ToolOutput"
                | "Action" | "Any" => {
                    if let Ok(part) = serde_json::from_value::<ContentPart>(val.clone()) {
                        return part;
                    }
                }
                _ => {}
            }
        }

        ContentPart::Any(val)
    }
}

impl From<Resource> for ContentPart {
    fn from(res: Resource) -> Self {
        if let Some(data) = res.blob {
            match String::from_utf8(data.0) {
                Ok(text) => ContentPart::Text { text },
                Err(v) => ContentPart::InlineData {
                    mime_type: res
                        .mime_type
                        .unwrap_or_else(|| "application/octet-stream".to_string()),
                    data: v.into_bytes().into(),
                },
            }
        } else if let Some(file_uri) = res.uri {
            ContentPart::FileData {
                file_uri,
                mime_type: res.mime_type,
            }
        } else {
            ContentPart::Text {
                text: serde_json::to_string(&res).unwrap_or_default(),
            }
        }
    }
}

/// Represents a request to a tool for processing.
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct ToolInput<T> {
    /// tool name.
    pub name: String,

    /// arguments in JSON format.
    pub args: T,

    /// The resources to process by the tool.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub resources: Vec<Resource>,

    /// The metadata for the tool request.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub meta: Option<RequestMeta>,
}

impl<T> ToolInput<T> {
    /// Creates a new tool input with the given name and arguments.
    pub fn new(name: String, args: T) -> Self {
        Self {
            name,
            args,
            resources: Vec::new(),
            meta: None,
        }
    }
}

/// Represents the output of a tool execution.
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct ToolOutput<T> {
    /// The output from the tool.
    pub output: T,

    /// A collection of artifacts generated by the tool execution.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub artifacts: Vec<Resource>,

    /// The usage statistics for the tool execution.
    pub usage: Usage,
}

impl<T> ToolOutput<T> {
    /// Creates a new tool output with the given output value.
    pub fn new(output: T) -> Self {
        Self {
            output,
            artifacts: Vec::new(),
            usage: Usage::default(),
        }
    }
}

/// Represents the metadata for an agent or tool request.
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct RequestMeta {
    /// The target engine principal for the request.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub engine: Option<Principal>,

    /// Gets the username from request context.
    /// Note: This is not verified and should not be used as a trusted identifier.
    /// For example, if triggered by a bot of X platform, this might be the username
    /// of the user interacting with the bot.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub user: Option<String>,

    /// Extra metadata key-value pairs.
    #[serde(flatten)]
    #[serde(skip_serializing_if = "Map::is_empty")]
    pub extra: Map<String, Json>,
}

impl RequestMeta {
    /// Gets an extra metadata value by key and deserializes it to the specified type.
    pub fn get_extra_as<T>(&self, key: &str) -> Option<T>
    where
        T: DeserializeOwned,
    {
        self.extra
            .get(key)
            .and_then(|value| serde_json::from_value(value.clone()).ok())
    }
}

/// Represents the usage statistics for the agent or tool execution.
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
pub struct Usage {
    /// input tokens sent to the LLM
    pub input_tokens: u64,

    /// output tokens received from the LLM
    pub output_tokens: u64,

    /// cached tokens used in the execution.
    #[serde(default)]
    pub cached_tokens: u64,

    /// number of requests made to agents and tools
    pub requests: u64,
}

impl Usage {
    /// Accumulates the usage statistics from another usage object.
    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);
    }
}

/// Represents a tool call response with it's ID, function name, and arguments.
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct ToolCall {
    /// tool function name.
    pub name: String,

    /// tool function  arguments.
    pub args: Json,

    /// The result of the tool call, auto processed by agents engine, if available.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<ToolOutput<Json>>,

    /// tool call id.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub call_id: Option<String>,

    /// The remote engine id where tool running
    #[serde(skip_serializing_if = "Option::is_none")]
    pub remote_id: Option<Principal>,
}

/// Represents a function definition with its metadata.
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct Function {
    /// Definition of the function.
    pub definition: FunctionDefinition,

    /// The tags of resource that this function supports.
    pub supported_resource_tags: Vec<String>,
}

/// Defines a callable function with its metadata and schema.
#[derive(Debug, Clone, Default, Deserialize, Serialize)]
pub struct FunctionDefinition {
    /// Name of the function.
    pub name: String,

    /// Description of what the function does.
    pub description: String,

    /// JSON schema defining the function's parameters.
    pub parameters: Json,

    /// Whether to enable strict schema adherence when generating the function call. If set to true, the model will follow the exact schema defined in the parameters field. Only a subset of JSON Schema is supported when strict is true.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub strict: Option<bool>,
}

impl FunctionDefinition {
    /// Modifies the function name with a prefix.
    pub fn name_with_prefix(mut self, prefix: &str) -> Self {
        self.name = format!("{}{}", prefix, self.name);
        self
    }
}

#[deprecated(note = "Use `estimate_tokens` instead.")]
pub fn evaluate_tokens(content: &str) -> usize {
    content.len() / 3
}

/// Estimate the number of tokens in the given content in the simplest way.
pub fn estimate_tokens(content: &str) -> usize {
    content.len() / 3
}

/// A document with metadata and content.
#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
pub struct Document {
    /// The metadata of the document.
    pub metadata: BTreeMap<String, Json>,

    /// The content of the document.
    pub content: Json,
}

impl Document {
    /// Creates a new text document with the given ID and text content.
    pub fn from_text(id: &str, text: &str) -> Self {
        Self {
            metadata: BTreeMap::from([
                ("id".to_string(), id.into()),
                ("type".to_string(), "Text".into()),
            ]),
            content: text.into(),
        }
    }
}

impl From<&Resource> for Document {
    fn from(res: &Resource) -> Self {
        let mut metadata = BTreeMap::from([
            ("id".to_string(), res._id.into()),
            ("type".to_string(), "Resource".into()),
        ]);
        if let Json::Object(mut val) = json!(res) {
            val.remove("blob");
            metadata.extend(val);
        };

        Self {
            metadata,
            content: Json::Null,
        }
    }
}

/// Collection of knowledge documents.
#[derive(Clone, Debug)]
pub struct Documents {
    /// The tag of the document collection. Defaults to "documents".
    tag: String,
    /// The documents in the collection.
    docs: Vec<Document>,
}

impl Default for Documents {
    fn default() -> Self {
        Self {
            tag: "documents".to_string(),
            docs: Vec::new(),
        }
    }
}

impl Documents {
    /// Creates a new document collection.
    pub fn new(tag: String, docs: Vec<Document>) -> Self {
        Self { tag, docs }
    }

    /// Sets the tag of the document collection.
    pub fn with_tag(self, tag: String) -> Self {
        Self { tag, ..self }
    }

    /// Returns the tag of the document collection.
    pub fn tag(&self) -> &str {
        &self.tag
    }

    /// Converts the document collection to a message.
    pub fn to_message(&self, rfc3339_datetime: &str) -> Option<Message> {
        if self.docs.is_empty() {
            return None;
        }

        Some(Message {
            role: "user".into(),
            content: vec![
                format!("Current Datetime: {}\n\n---\n\n{}", rfc3339_datetime, self).into(),
            ],
            name: Some("$system".into()),
            ..Default::default()
        })
    }

    /// Appends a document to the collection.
    pub fn append(&mut self, doc: Document) {
        self.docs.push(doc);
    }
}

impl From<Vec<String>> for Documents {
    fn from(texts: Vec<String>) -> Self {
        let mut docs = Vec::new();
        for (i, text) in texts.into_iter().enumerate() {
            docs.push(Document {
                content: text.into(),
                metadata: BTreeMap::from([
                    ("_id".to_string(), i.into()),
                    ("type".to_string(), "Text".into()),
                ]),
            });
        }
        Self {
            docs,
            ..Default::default()
        }
    }
}

impl From<Vec<Document>> for Documents {
    fn from(docs: Vec<Document>) -> Self {
        Self {
            docs,
            ..Default::default()
        }
    }
}

impl std::ops::Deref for Documents {
    type Target = Vec<Document>;

    fn deref(&self) -> &Self::Target {
        &self.docs
    }
}

impl std::ops::DerefMut for Documents {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.docs
    }
}

impl AsRef<Vec<Document>> for Documents {
    fn as_ref(&self) -> &Vec<Document> {
        &self.docs
    }
}

impl std::fmt::Display for Document {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        json!(self).fmt(f)
    }
}

impl std::fmt::Display for Documents {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if self.docs.is_empty() {
            return Ok(());
        }
        writeln!(f, "<{}>", self.tag)?;
        for doc in &self.docs {
            writeln!(f, "{}", doc)?;
        }
        write!(f, "</{}>", self.tag)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_prompt() {
        let documents: Documents = vec![
            Document {
                metadata: BTreeMap::from([("_id".to_string(), 1.into())]),
                content: "Test document 1.".into(),
            },
            Document {
                metadata: BTreeMap::from([
                    ("_id".to_string(), 2.into()),
                    ("key".to_string(), "value".into()),
                    ("a".to_string(), "b".into()),
                ]),
                content: "Test document 2.".into(),
            },
        ]
        .into();
        // println!("{}", documents);

        let s = documents.to_string();
        let lines: Vec<&str> = s.lines().collect();
        assert_eq!(lines[0], "<documents>");
        assert_eq!(lines[3], "</documents>");

        let doc1: Json = serde_json::from_str(lines[1]).unwrap();
        assert_eq!(doc1.get("content").unwrap(), "Test document 1.");
        assert_eq!(doc1.get("metadata").unwrap().get("_id").unwrap(), 1);

        let doc2: Json = serde_json::from_str(lines[2]).unwrap();
        assert_eq!(doc2.get("content").unwrap(), "Test document 2.");
        assert_eq!(doc2.get("metadata").unwrap().get("_id").unwrap(), 2);
        assert_eq!(doc2.get("metadata").unwrap().get("key").unwrap(), "value");
        assert_eq!(doc2.get("metadata").unwrap().get("a").unwrap(), "b");

        let documents = documents.with_tag("my_docs".to_string());
        let s = documents.to_string();
        let lines: Vec<&str> = s.lines().collect();
        assert_eq!(lines[0], "<my_docs>");
        assert_eq!(lines[3], "</my_docs>");

        let doc1: Json = serde_json::from_str(lines[1]).unwrap();
        assert_eq!(doc1.get("content").unwrap(), "Test document 1.");
        assert_eq!(doc1.get("metadata").unwrap().get("_id").unwrap(), 1);

        let doc2: Json = serde_json::from_str(lines[2]).unwrap();
        assert_eq!(doc2.get("content").unwrap(), "Test document 2.");
        assert_eq!(doc2.get("metadata").unwrap().get("_id").unwrap(), 2);
        assert_eq!(doc2.get("metadata").unwrap().get("key").unwrap(), "value");
        assert_eq!(doc2.get("metadata").unwrap().get("a").unwrap(), "b");
    }

    #[test]
    fn test_content_part_text_serde_and_from() {
        let part: ContentPart = "hello".to_string().into();
        assert_eq!(
            part,
            ContentPart::Text {
                text: "hello".into()
            }
        );

        // serde round-trip
        let v = serde_json::to_value(&part).unwrap();
        assert_eq!(v.get("type").unwrap(), "Text");
        assert_eq!(v.get("text").unwrap(), "hello");

        let back: ContentPart = serde_json::from_value(v.clone()).unwrap();
        assert_eq!(back, part);
        let back: ContentPart = v.into();
        assert_eq!(back, part);

        let part: Vec<ContentPart> = serde_json::from_str(
            r#"
            [
                "hello",
                {
                    "type": "Text",
                    "text": "world"
                }
            ]
            "#,
        )
        .unwrap();
        assert_eq!(
            part,
            vec![
                ContentPart::Text {
                    text: "hello".into()
                },
                ContentPart::Text {
                    text: "world".into()
                }
            ]
        );
    }

    #[test]
    fn test_content_part_filedata_serde_optional() {
        // mime_type = None -> 不序列化
        let part = ContentPart::FileData {
            file_uri: "gs://bucket/file".into(),
            mime_type: None,
        };
        let v = serde_json::to_value(&part).unwrap();
        assert_eq!(v.get("type").unwrap(), "FileData");
        // 字段采用 camelCase
        assert_eq!(v.get("fileUri").unwrap(), "gs://bucket/file");
        assert!(v.get("mimeType").is_none());

        // mime_type = Some -> 出现
        let part2 = ContentPart::FileData {
            file_uri: "gs://bucket/file2".into(),
            mime_type: Some("image/png".into()),
        };
        let v2 = serde_json::to_value(&part2).unwrap();
        assert_eq!(v2.get("type").unwrap(), "FileData");
        assert_eq!(v2.get("fileUri").unwrap(), "gs://bucket/file2");
        assert_eq!(v2.get("mimeType").unwrap(), "image/png");

        // 反序列化校验
        let back: ContentPart = serde_json::from_value(v2.clone()).unwrap();
        assert_eq!(back, part2);
        let back: ContentPart = v2.into();
        assert_eq!(back, part2);
    }

    #[test]
    fn test_content_part_inlinedata_serde() {
        let part = ContentPart::InlineData {
            mime_type: "text/plain".into(),
            data: b"hello".to_vec().into(),
        };
        let v = serde_json::to_value(&part).unwrap();
        assert_eq!(v.get("type").unwrap(), "InlineData");
        assert_eq!(v.get("mimeType").unwrap(), "text/plain");
        assert_eq!(v.get("data").unwrap(), "aGVsbG8=");

        let back: ContentPart = serde_json::from_value(v.clone()).unwrap();
        assert_eq!(back, part);
        let back: ContentPart = v.into();
        assert_eq!(back, part);
    }

    #[test]
    fn test_content_part_any_serde() {
        let v = json!({
            "type": "text/plain",
            "data": "aGVsbG8=",
        });
        let part: ContentPart = v.clone().into();
        assert_eq!(part, ContentPart::Any(v));
        let v2 = serde_json::to_value(&part).unwrap();
        assert_eq!(v2.get("type").unwrap(), "text/plain");
        assert_eq!(v2.get("data").unwrap(), "aGVsbG8=");

        let part = ContentPart::Any(json!({
            "data": "aGVsbG8=",
        }));
        let v2 = serde_json::to_value(&part).unwrap();
        assert!(v2.get("type").is_none());
        assert_eq!(v2.get("data").unwrap(), "aGVsbG8=");
    }

    #[test]
    fn test_content_part_toolcall_and_tooloutput_serde() {
        let call = ContentPart::ToolCall {
            name: "sum".into(),
            args: serde_json::json!({"x":1, "y":2}),
            call_id: None,
        };
        let v_call = serde_json::to_value(&call).unwrap();
        assert_eq!(v_call.get("type").unwrap(), "ToolCall");
        assert_eq!(v_call.get("name").unwrap(), "sum");
        assert_eq!(
            v_call.get("args").unwrap(),
            &serde_json::json!({"x":1, "y":2})
        );
        // callId 省略
        assert!(v_call.get("callId").is_none());
        let back_call: ContentPart = serde_json::from_value(v_call.clone()).unwrap();
        assert_eq!(back_call, call);
        let back: ContentPart = v_call.into();
        assert_eq!(back, call);

        let out = ContentPart::ToolOutput {
            name: "sum".into(),
            output: serde_json::json!({"result":3}),
            call_id: Some("c1".into()),
            remote_id: None,
        };
        let v_out = serde_json::to_value(&out).unwrap();
        assert_eq!(v_out.get("type").unwrap(), "ToolOutput");
        assert_eq!(v_out.get("name").unwrap(), "sum");
        assert_eq!(
            v_out.get("output").unwrap(),
            &serde_json::json!({"result":3})
        );
        // callId 存在
        assert_eq!(v_out.get("callId").unwrap(), "c1");
        let back_out: ContentPart = serde_json::from_value(v_out.clone()).unwrap();
        assert_eq!(back_out, out);
        let back: ContentPart = v_out.into();
        assert_eq!(back, out);
    }

    #[test]
    fn test_message_tool_calls_extract_from_content_parts() {
        let parts = vec![
            ContentPart::Text {
                text: "hello".into(),
            },
            ContentPart::ToolCall {
                name: "sum".into(),
                args: serde_json::json!({"x":1, "y": 2}),
                call_id: Some("abc".into()),
            },
            ContentPart::ToolCall {
                name: "echo".into(),
                args: serde_json::json!({"text":"hi"}),
                call_id: None,
            },
        ];
        let msg = Message {
            role: "assistant".into(),
            content: parts,
            ..Default::default()
        };
        println!("{:#?}", json!(msg));

        let calls = msg.tool_calls();
        assert_eq!(calls.len(), 2);
        assert_eq!(calls[0].name, "sum");
        assert_eq!(calls[0].args, serde_json::json!({"x":1, "y":2}));
        assert_eq!(calls[1].name, "echo");
        assert_eq!(calls[1].args, serde_json::json!({"text":"hi"}));
    }

    #[test]
    fn test_message_content_deserialize_from_string() {
        // content as a plain string
        let msg: Message = serde_json::from_value(serde_json::json!({
            "role": "user",
            "content": "hello world"
        }))
        .unwrap();
        assert_eq!(msg.role, "user");
        assert_eq!(msg.content.len(), 1);
        assert_eq!(
            msg.content[0],
            ContentPart::Text {
                text: "hello world".into()
            }
        );

        // content as an array still works
        let msg2: Message = serde_json::from_value(serde_json::json!({
            "role": "assistant",
            "content": [{"type": "Text", "text": "hi"}]
        }))
        .unwrap();
        assert_eq!(msg2.content.len(), 1);
        assert_eq!(msg2.content[0], ContentPart::Text { text: "hi".into() });

        // missing content defaults to empty vec
        let msg3: Message = serde_json::from_value(serde_json::json!({
            "role": "system"
        }))
        .unwrap();
        assert!(msg3.content.is_empty());
    }

    #[test]
    fn test_request_meta_extra_flatten_serde() {
        // empty extra should not serialize
        let meta = RequestMeta {
            engine: None,
            user: None,
            extra: Map::new(),
        };
        let v = serde_json::to_value(&meta).unwrap();
        assert_eq!(v, serde_json::json!({}));

        // extra should be flattened into the top-level object
        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());

        // deserialization: unknown fields go into extra
        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})
        );

        // round-trip (field-by-field)
        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})
        );
    }
}