anda_core/
model.rs

1//! Core data models and traits for the AI agent system.
2//!
3//! This module defines the fundamental data structures and interfaces used throughout the AI agent system.
4//! It includes:
5//! - Core message and conversation structures ([`AgentOutput`], [`Message`], [`ToolCall`]).
6//! - Function definition and tooling support ([`FunctionDefinition`]).
7//! - Knowledge and document handling ([`Document`], [`Documents`]).
8//! - Completion request and response structures ([`CompletionRequest`], [`Embedding`]).
9//! - Core AI capabilities traits ([`CompletionFeatures`], [`EmbeddingFeatures`]).
10
11use candid::Principal;
12use serde::{Deserialize, Serialize};
13use serde_json::{Map, json};
14use std::collections::BTreeMap;
15
16use crate::Json;
17pub use ic_auth_types::{ByteArrayB64, ByteBufB64, Xid};
18
19mod completion;
20mod embedding;
21mod resource;
22
23pub use completion::*;
24pub use embedding::*;
25pub use resource::*;
26
27/// Represents a request to an agent for processing.
28#[derive(Debug, Clone, Default, Deserialize, Serialize)]
29pub struct AgentInput {
30    /// agent name, use default agent if empty.
31    pub name: String,
32
33    /// agent prompt or message.
34    pub prompt: String,
35
36    /// The resources to process by the agent.
37    #[serde(default, skip_serializing_if = "Vec::is_empty")]
38    pub resources: Vec<Resource>,
39
40    /// The topics for the agent request.
41    #[serde(skip_serializing_if = "Option::is_none")]
42    pub topics: Option<Vec<String>>,
43
44    /// The metadata for the agent request
45    #[serde(skip_serializing_if = "Option::is_none")]
46    pub meta: Option<RequestMeta>,
47}
48
49impl AgentInput {
50    /// Creates a new agent input with the given name and prompt.
51    pub fn new(name: String, prompt: String) -> Self {
52        Self {
53            name,
54            prompt,
55            resources: Vec::new(),
56            topics: None,
57            meta: None,
58        }
59    }
60}
61
62/// Represents the output of an agent execution.
63#[derive(Debug, Clone, Default, Deserialize, Serialize)]
64pub struct AgentOutput {
65    /// The output content from the agent, may be empty.
66    pub content: String,
67
68    /// The usage statistics for the agent execution.
69    pub usage: Usage,
70
71    /// Indicates failure reason if present, None means successful execution.
72    /// Should be None when finish_reason is "stop" or "tool_calls".
73    #[serde(skip_serializing_if = "Option::is_none")]
74    pub failed_reason: Option<String>,
75
76    /// Tool calls returned by the LLM function calling.
77    #[serde(default, skip_serializing_if = "Vec::is_empty")]
78    pub tool_calls: Vec<ToolCall>,
79
80    /// The history of the conversation.
81    #[serde(default, skip_serializing_if = "Vec::is_empty")]
82    pub chat_history: Vec<Message>,
83
84    /// raw_history is the model specialized history of the conversation,
85    /// will be included in `ctx.completion` response,
86    /// but not be included in the engine response.
87    #[serde(default, skip_serializing_if = "Vec::is_empty")]
88    pub raw_history: Vec<Json>,
89
90    /// A collection of artifacts generated by the agent during the execution of the task.
91    #[serde(default, skip_serializing_if = "Vec::is_empty")]
92    pub artifacts: Vec<Resource>,
93
94    /// The conversation ID.
95    #[serde(skip_serializing_if = "Option::is_none")]
96    pub conversation: Option<u64>,
97
98    /// The model used by the agent.
99    #[serde(skip_serializing_if = "Option::is_none")]
100    pub model: Option<String>,
101}
102
103fn deserialize_content<'de, D>(deserializer: D) -> Result<Vec<ContentPart>, D::Error>
104where
105    D: serde::Deserializer<'de>,
106{
107    let value = Json::deserialize(deserializer)?;
108    match value {
109        Json::String(s) => Ok(vec![ContentPart::Text { text: s }]),
110        Json::Array(_) => Vec::<ContentPart>::deserialize(value).map_err(serde::de::Error::custom),
111        _ => Err(serde::de::Error::custom(
112            "expected a string or array for content",
113        )),
114    }
115}
116
117/// Represents a message send to LLM for completion.
118#[derive(Debug, Clone, Default, Deserialize, Serialize)]
119pub struct Message {
120    /// Message role: "system", "user", "assistant", "tool".
121    pub role: String,
122
123    /// The content of the message
124    #[serde(default, deserialize_with = "deserialize_content")]
125    pub content: Vec<ContentPart>,
126
127    /// An optional name for the participant. Provides the model information to differentiate between participants of the same role.
128    /// This field is not used by the model.
129    #[serde(skip_serializing_if = "Option::is_none")]
130    pub name: Option<String>,
131
132    /// The user ID of the message sender.
133    /// This field is not used by the model.
134    #[serde(skip_serializing_if = "Option::is_none")]
135    pub user: Option<Principal>,
136
137    /// The timestamp of the message.
138    /// This field is not used by the model.
139    #[serde(skip_serializing_if = "Option::is_none")]
140    pub timestamp: Option<u64>,
141}
142
143impl Message {
144    pub fn text(&self) -> Option<String> {
145        let mut texts: Vec<&str> = Vec::new();
146        for part in &self.content {
147            if let ContentPart::Text { text } = part {
148                texts.push(text);
149            }
150        }
151        if texts.is_empty() {
152            return None;
153        }
154        Some(texts.join("\n"))
155    }
156
157    pub fn tool_calls(&self) -> Vec<ToolCall> {
158        let mut tool_calls: Vec<ToolCall> = Vec::new();
159        for part in &self.content {
160            if let ContentPart::ToolCall {
161                name,
162                args,
163                call_id,
164            } = part
165            {
166                tool_calls.push(ToolCall {
167                    name: name.clone(),
168                    args: args.clone(),
169                    call_id: call_id.clone(),
170                    result: None,
171                    remote_id: None,
172                });
173            }
174        }
175        tool_calls
176    }
177}
178
179#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
180#[serde(tag = "type", rename_all_fields = "camelCase")]
181pub enum ContentPart {
182    Text {
183        text: String,
184    },
185    Reasoning {
186        text: String,
187    },
188    FileData {
189        file_uri: String,
190
191        #[serde(skip_serializing_if = "Option::is_none")]
192        mime_type: Option<String>,
193    },
194    InlineData {
195        mime_type: String,
196        data: ByteBufB64,
197    },
198    ToolCall {
199        name: String,
200        args: Json,
201
202        #[serde(skip_serializing_if = "Option::is_none")]
203        call_id: Option<String>,
204    },
205    ToolOutput {
206        name: String,
207        output: Json,
208
209        #[serde(skip_serializing_if = "Option::is_none")]
210        call_id: Option<String>,
211
212        #[serde(skip_serializing_if = "Option::is_none")]
213        remote_id: Option<Principal>,
214    },
215    Action {
216        name: String,
217        payload: Json,
218
219        #[serde(skip_serializing_if = "Option::is_none")]
220        recipients: Option<Vec<Principal>>,
221
222        #[serde(skip_serializing_if = "Option::is_none")]
223        signature: Option<ByteBufB64>,
224    },
225    Any(Json),
226}
227
228impl<'de> Deserialize<'de> for ContentPart {
229    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
230    where
231        D: serde::Deserializer<'de>,
232    {
233        let value = Json::deserialize(deserializer)?;
234        match &value {
235            Json::String(s) => Ok(ContentPart::Text { text: s.clone() }),
236            Json::Object(map)
237                if matches!(
238                    map.get("type").and_then(|t| t.as_str()),
239                    Some(
240                        "Text"
241                            | "Reasoning"
242                            | "FileData"
243                            | "InlineData"
244                            | "ToolCall"
245                            | "ToolOutput"
246                            | "Action"
247                    )
248                ) =>
249            {
250                #[derive(Deserialize)]
251                #[serde(tag = "type", rename_all_fields = "camelCase")]
252                enum Helper {
253                    Text {
254                        text: String,
255                    },
256                    Reasoning {
257                        text: String,
258                    },
259                    FileData {
260                        file_uri: String,
261                        mime_type: Option<String>,
262                    },
263                    InlineData {
264                        mime_type: String,
265                        data: ByteBufB64,
266                    },
267                    ToolCall {
268                        name: String,
269                        args: Json,
270                        call_id: Option<String>,
271                    },
272                    ToolOutput {
273                        name: String,
274                        output: Json,
275                        call_id: Option<String>,
276                        remote_id: Option<Principal>,
277                    },
278                    Action {
279                        name: String,
280                        payload: Json,
281                        recipients: Option<Vec<Principal>>,
282                        signature: Option<ByteBufB64>,
283                    },
284                }
285
286                match serde_json::from_value::<Helper>(value) {
287                    Ok(h) => Ok(match h {
288                        Helper::Text { text } => ContentPart::Text { text },
289                        Helper::Reasoning { text } => ContentPart::Reasoning { text },
290                        Helper::FileData {
291                            file_uri,
292                            mime_type,
293                        } => ContentPart::FileData {
294                            file_uri,
295                            mime_type,
296                        },
297                        Helper::InlineData { mime_type, data } => {
298                            ContentPart::InlineData { mime_type, data }
299                        }
300                        Helper::ToolCall {
301                            name,
302                            args,
303                            call_id,
304                        } => ContentPart::ToolCall {
305                            name,
306                            args,
307                            call_id,
308                        },
309                        Helper::ToolOutput {
310                            name,
311                            output,
312                            call_id,
313                            remote_id,
314                        } => ContentPart::ToolOutput {
315                            name,
316                            output,
317                            call_id,
318                            remote_id,
319                        },
320                        Helper::Action {
321                            name,
322                            payload,
323                            recipients,
324                            signature,
325                        } => ContentPart::Action {
326                            name,
327                            payload,
328                            recipients,
329                            signature,
330                        },
331                    }),
332                    Err(_) => Err(serde::de::Error::custom("invalid ContentPart")),
333                }
334            }
335            _ => Ok(ContentPart::Any(value)),
336        }
337    }
338}
339
340impl From<String> for ContentPart {
341    fn from(text: String) -> Self {
342        Self::Text { text }
343    }
344}
345
346impl From<Json> for ContentPart {
347    fn from(val: Json) -> Self {
348        if let Json::Object(map) = &val
349            && let Some(t) = map.get("type").and_then(|x| x.as_str())
350        {
351            match t {
352                "Text" | "Reasoning" | "FileData" | "InlineData" | "ToolCall" | "ToolOutput"
353                | "Action" | "Any" => {
354                    if let Ok(part) = serde_json::from_value::<ContentPart>(val.clone()) {
355                        return part;
356                    }
357                }
358                _ => {}
359            }
360        }
361
362        ContentPart::Any(val)
363    }
364}
365
366/// Represents a request to a tool for processing.
367#[derive(Debug, Clone, Default, Deserialize, Serialize)]
368pub struct ToolInput<T> {
369    /// tool name.
370    pub name: String,
371
372    /// arguments in JSON format.
373    pub args: T,
374
375    /// The resources to process by the tool.
376    #[serde(default, skip_serializing_if = "Vec::is_empty")]
377    pub resources: Vec<Resource>,
378
379    /// The metadata for the tool request.
380    #[serde(skip_serializing_if = "Option::is_none")]
381    pub meta: Option<RequestMeta>,
382}
383
384impl<T> ToolInput<T> {
385    /// Creates a new tool input with the given name and arguments.
386    pub fn new(name: String, args: T) -> Self {
387        Self {
388            name,
389            args,
390            resources: Vec::new(),
391            meta: None,
392        }
393    }
394}
395
396/// Represents the output of a tool execution.
397#[derive(Debug, Clone, Default, Deserialize, Serialize)]
398pub struct ToolOutput<T> {
399    /// The output from the tool.
400    pub output: T,
401
402    /// A collection of artifacts generated by the tool execution.
403    #[serde(default, skip_serializing_if = "Vec::is_empty")]
404    pub artifacts: Vec<Resource>,
405
406    /// The usage statistics for the tool execution.
407    pub usage: Usage,
408}
409
410impl<T> ToolOutput<T> {
411    /// Creates a new tool output with the given output value.
412    pub fn new(output: T) -> Self {
413        Self {
414            output,
415            artifacts: Vec::new(),
416            usage: Usage::default(),
417        }
418    }
419}
420
421/// Represents the metadata for an agent or tool request.
422#[derive(Debug, Clone, Default, Deserialize, Serialize)]
423pub struct RequestMeta {
424    /// The target engine principal for the request.
425    #[serde(skip_serializing_if = "Option::is_none")]
426    pub engine: Option<Principal>,
427
428    /// Gets the username from request context.
429    /// Note: This is not verified and should not be used as a trusted identifier.
430    /// For example, if triggered by a bot of X platform, this might be the username
431    /// of the user interacting with the bot.
432    #[serde(skip_serializing_if = "Option::is_none")]
433    pub user: Option<String>,
434
435    /// Extra metadata key-value pairs.
436    #[serde(flatten)]
437    #[serde(skip_serializing_if = "Map::is_empty")]
438    pub extra: Map<String, Json>,
439}
440
441/// Represents the usage statistics for the agent or tool execution.
442#[derive(Clone, Debug, Default, Deserialize, Serialize)]
443pub struct Usage {
444    /// input tokens sent to the LLM
445    pub input_tokens: u64,
446
447    /// output tokens received from the LLM
448    pub output_tokens: u64,
449
450    /// number of requests made to agents and tools
451    pub requests: u64,
452}
453
454impl Usage {
455    /// Accumulates the usage statistics from another usage object.
456    pub fn accumulate(&mut self, other: &Usage) {
457        self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
458        self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
459        self.requests = self.requests.saturating_add(other.requests);
460    }
461}
462
463/// Represents a tool call response with it's ID, function name, and arguments.
464#[derive(Debug, Clone, Default, Deserialize, Serialize)]
465pub struct ToolCall {
466    /// tool function name.
467    pub name: String,
468
469    /// tool function  arguments.
470    pub args: Json,
471
472    /// The result of the tool call, auto processed by agents engine, if available.
473    #[serde(skip_serializing_if = "Option::is_none")]
474    pub result: Option<ToolOutput<Json>>,
475
476    /// tool call id.
477    #[serde(skip_serializing_if = "Option::is_none")]
478    pub call_id: Option<String>,
479
480    /// The remote engine id where tool running
481    #[serde(skip_serializing_if = "Option::is_none")]
482    pub remote_id: Option<Principal>,
483}
484
485/// Represents a function definition with its metadata.
486#[derive(Debug, Clone, Default, Deserialize, Serialize)]
487pub struct Function {
488    /// Definition of the function.
489    pub definition: FunctionDefinition,
490
491    /// The tags of resource that this function supports.
492    pub supported_resource_tags: Vec<String>,
493}
494
495/// Defines a callable function with its metadata and schema.
496#[derive(Debug, Clone, Default, Deserialize, Serialize)]
497pub struct FunctionDefinition {
498    /// Name of the function.
499    pub name: String,
500
501    /// Description of what the function does.
502    pub description: String,
503
504    /// JSON schema defining the function's parameters.
505    pub parameters: Json,
506
507    /// 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.
508    #[serde(skip_serializing_if = "Option::is_none")]
509    pub strict: Option<bool>,
510}
511
512impl FunctionDefinition {
513    /// Modifies the function name with a prefix.
514    pub fn name_with_prefix(mut self, prefix: &str) -> Self {
515        self.name = format!("{}{}", prefix, self.name);
516        self
517    }
518}
519
520/// Returns the number of tokens in the given content in the simplest way.
521pub fn evaluate_tokens(content: &str) -> usize {
522    content.len() / 3
523}
524
525/// A document with metadata and content.
526#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
527pub struct Document {
528    /// The metadata of the document.
529    pub metadata: BTreeMap<String, Json>,
530
531    /// The content of the document.
532    pub content: Json,
533}
534
535impl Document {
536    /// Creates a new text document with the given ID and text content.
537    pub fn from_text(id: &str, text: &str) -> Self {
538        Self {
539            metadata: BTreeMap::from([
540                ("id".to_string(), id.into()),
541                ("type".to_string(), "Text".into()),
542            ]),
543            content: text.into(),
544        }
545    }
546}
547
548impl From<&Resource> for Document {
549    fn from(res: &Resource) -> Self {
550        let mut metadata = BTreeMap::from([
551            ("id".to_string(), res._id.into()),
552            ("type".to_string(), "Resource".into()),
553        ]);
554        if let Json::Object(mut val) = json!(res) {
555            val.remove("blob");
556            metadata.extend(val);
557        };
558
559        Self {
560            metadata,
561            content: Json::Null,
562        }
563    }
564}
565
566/// Collection of knowledge documents.
567#[derive(Clone, Debug)]
568pub struct Documents {
569    /// The tag of the document collection. Defaults to "documents".
570    tag: String,
571    /// The documents in the collection.
572    docs: Vec<Document>,
573}
574
575impl Default for Documents {
576    fn default() -> Self {
577        Self {
578            tag: "documents".to_string(),
579            docs: Vec::new(),
580        }
581    }
582}
583
584impl Documents {
585    /// Creates a new document collection.
586    pub fn new(tag: String, docs: Vec<Document>) -> Self {
587        Self { tag, docs }
588    }
589
590    /// Sets the tag of the document collection.
591    pub fn with_tag(self, tag: String) -> Self {
592        Self { tag, ..self }
593    }
594
595    /// Returns the tag of the document collection.
596    pub fn tag(&self) -> &str {
597        &self.tag
598    }
599
600    /// Converts the document collection to a message.
601    pub fn to_message(&self, rfc3339_datetime: &str) -> Option<Message> {
602        if self.docs.is_empty() {
603            return None;
604        }
605
606        Some(Message {
607            role: "user".into(),
608            content: vec![
609                format!("Current Datetime: {}\n\n---\n\n{}", rfc3339_datetime, self).into(),
610            ],
611            name: Some("$system".into()),
612            ..Default::default()
613        })
614    }
615
616    /// Appends a document to the collection.
617    pub fn append(&mut self, doc: Document) {
618        self.docs.push(doc);
619    }
620}
621
622impl From<Vec<String>> for Documents {
623    fn from(texts: Vec<String>) -> Self {
624        let mut docs = Vec::new();
625        for (i, text) in texts.into_iter().enumerate() {
626            docs.push(Document {
627                content: text.into(),
628                metadata: BTreeMap::from([
629                    ("_id".to_string(), i.into()),
630                    ("type".to_string(), "Text".into()),
631                ]),
632            });
633        }
634        Self {
635            docs,
636            ..Default::default()
637        }
638    }
639}
640
641impl From<Vec<Document>> for Documents {
642    fn from(docs: Vec<Document>) -> Self {
643        Self {
644            docs,
645            ..Default::default()
646        }
647    }
648}
649
650impl std::ops::Deref for Documents {
651    type Target = Vec<Document>;
652
653    fn deref(&self) -> &Self::Target {
654        &self.docs
655    }
656}
657
658impl std::ops::DerefMut for Documents {
659    fn deref_mut(&mut self) -> &mut Self::Target {
660        &mut self.docs
661    }
662}
663
664impl AsRef<Vec<Document>> for Documents {
665    fn as_ref(&self) -> &Vec<Document> {
666        &self.docs
667    }
668}
669
670impl std::fmt::Display for Document {
671    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
672        json!(self).fmt(f)
673    }
674}
675
676impl std::fmt::Display for Documents {
677    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
678        if self.docs.is_empty() {
679            return Ok(());
680        }
681        writeln!(f, "<{}>", self.tag)?;
682        for doc in &self.docs {
683            writeln!(f, "{}", doc)?;
684        }
685        write!(f, "</{}>", self.tag)
686    }
687}
688
689#[cfg(test)]
690mod tests {
691    use super::*;
692
693    #[test]
694    fn test_prompt() {
695        let documents: Documents = vec![
696            Document {
697                metadata: BTreeMap::from([("_id".to_string(), 1.into())]),
698                content: "Test document 1.".into(),
699            },
700            Document {
701                metadata: BTreeMap::from([
702                    ("_id".to_string(), 2.into()),
703                    ("key".to_string(), "value".into()),
704                    ("a".to_string(), "b".into()),
705                ]),
706                content: "Test document 2.".into(),
707            },
708        ]
709        .into();
710        // println!("{}", documents);
711
712        let s = documents.to_string();
713        let lines: Vec<&str> = s.lines().collect();
714        assert_eq!(lines[0], "<documents>");
715        assert_eq!(lines[3], "</documents>");
716
717        let doc1: Json = serde_json::from_str(lines[1]).unwrap();
718        assert_eq!(doc1.get("content").unwrap(), "Test document 1.");
719        assert_eq!(doc1.get("metadata").unwrap().get("_id").unwrap(), 1);
720
721        let doc2: Json = serde_json::from_str(lines[2]).unwrap();
722        assert_eq!(doc2.get("content").unwrap(), "Test document 2.");
723        assert_eq!(doc2.get("metadata").unwrap().get("_id").unwrap(), 2);
724        assert_eq!(doc2.get("metadata").unwrap().get("key").unwrap(), "value");
725        assert_eq!(doc2.get("metadata").unwrap().get("a").unwrap(), "b");
726
727        let documents = documents.with_tag("my_docs".to_string());
728        let s = documents.to_string();
729        let lines: Vec<&str> = s.lines().collect();
730        assert_eq!(lines[0], "<my_docs>");
731        assert_eq!(lines[3], "</my_docs>");
732
733        let doc1: Json = serde_json::from_str(lines[1]).unwrap();
734        assert_eq!(doc1.get("content").unwrap(), "Test document 1.");
735        assert_eq!(doc1.get("metadata").unwrap().get("_id").unwrap(), 1);
736
737        let doc2: Json = serde_json::from_str(lines[2]).unwrap();
738        assert_eq!(doc2.get("content").unwrap(), "Test document 2.");
739        assert_eq!(doc2.get("metadata").unwrap().get("_id").unwrap(), 2);
740        assert_eq!(doc2.get("metadata").unwrap().get("key").unwrap(), "value");
741        assert_eq!(doc2.get("metadata").unwrap().get("a").unwrap(), "b");
742    }
743
744    #[test]
745    fn test_content_part_text_serde_and_from() {
746        let part: ContentPart = "hello".to_string().into();
747        assert_eq!(
748            part,
749            ContentPart::Text {
750                text: "hello".into()
751            }
752        );
753
754        // serde round-trip
755        let v = serde_json::to_value(&part).unwrap();
756        assert_eq!(v.get("type").unwrap(), "Text");
757        assert_eq!(v.get("text").unwrap(), "hello");
758
759        let back: ContentPart = serde_json::from_value(v.clone()).unwrap();
760        assert_eq!(back, part);
761        let back: ContentPart = v.into();
762        assert_eq!(back, part);
763
764        let part: Vec<ContentPart> = serde_json::from_str(
765            r#"
766            [
767                "hello",
768                {
769                    "type": "Text",
770                    "text": "world"
771                }
772            ]
773            "#,
774        )
775        .unwrap();
776        assert_eq!(
777            part,
778            vec![
779                ContentPart::Text {
780                    text: "hello".into()
781                },
782                ContentPart::Text {
783                    text: "world".into()
784                }
785            ]
786        );
787    }
788
789    #[test]
790    fn test_content_part_filedata_serde_optional() {
791        // mime_type = None -> 不序列化
792        let part = ContentPart::FileData {
793            file_uri: "gs://bucket/file".into(),
794            mime_type: None,
795        };
796        let v = serde_json::to_value(&part).unwrap();
797        assert_eq!(v.get("type").unwrap(), "FileData");
798        // 字段采用 camelCase
799        assert_eq!(v.get("fileUri").unwrap(), "gs://bucket/file");
800        assert!(v.get("mimeType").is_none());
801
802        // mime_type = Some -> 出现
803        let part2 = ContentPart::FileData {
804            file_uri: "gs://bucket/file2".into(),
805            mime_type: Some("image/png".into()),
806        };
807        let v2 = serde_json::to_value(&part2).unwrap();
808        assert_eq!(v2.get("type").unwrap(), "FileData");
809        assert_eq!(v2.get("fileUri").unwrap(), "gs://bucket/file2");
810        assert_eq!(v2.get("mimeType").unwrap(), "image/png");
811
812        // 反序列化校验
813        let back: ContentPart = serde_json::from_value(v2.clone()).unwrap();
814        assert_eq!(back, part2);
815        let back: ContentPart = v2.into();
816        assert_eq!(back, part2);
817    }
818
819    #[test]
820    fn test_content_part_inlinedata_serde() {
821        let part = ContentPart::InlineData {
822            mime_type: "text/plain".into(),
823            data: b"hello".to_vec().into(),
824        };
825        let v = serde_json::to_value(&part).unwrap();
826        assert_eq!(v.get("type").unwrap(), "InlineData");
827        assert_eq!(v.get("mimeType").unwrap(), "text/plain");
828        assert_eq!(v.get("data").unwrap(), "aGVsbG8=");
829
830        let back: ContentPart = serde_json::from_value(v.clone()).unwrap();
831        assert_eq!(back, part);
832        let back: ContentPart = v.into();
833        assert_eq!(back, part);
834    }
835
836    #[test]
837    fn test_content_part_any_serde() {
838        let v = json!({
839            "type": "text/plain",
840            "data": "aGVsbG8=",
841        });
842        let part: ContentPart = v.clone().into();
843        assert_eq!(part, ContentPart::Any(v));
844        let v2 = serde_json::to_value(&part).unwrap();
845        assert_eq!(v2.get("type").unwrap(), "text/plain");
846        assert_eq!(v2.get("data").unwrap(), "aGVsbG8=");
847
848        let part = ContentPart::Any(json!({
849            "data": "aGVsbG8=",
850        }));
851        let v2 = serde_json::to_value(&part).unwrap();
852        assert_eq!(v2.get("type").unwrap(), "Any");
853        assert_eq!(v2.get("data").unwrap(), "aGVsbG8=");
854    }
855
856    #[test]
857    fn test_content_part_toolcall_and_tooloutput_serde() {
858        let call = ContentPart::ToolCall {
859            name: "sum".into(),
860            args: serde_json::json!({"x":1, "y":2}),
861            call_id: None,
862        };
863        let v_call = serde_json::to_value(&call).unwrap();
864        assert_eq!(v_call.get("type").unwrap(), "ToolCall");
865        assert_eq!(v_call.get("name").unwrap(), "sum");
866        assert_eq!(
867            v_call.get("args").unwrap(),
868            &serde_json::json!({"x":1, "y":2})
869        );
870        // callId 省略
871        assert!(v_call.get("callId").is_none());
872        let back_call: ContentPart = serde_json::from_value(v_call.clone()).unwrap();
873        assert_eq!(back_call, call);
874        let back: ContentPart = v_call.into();
875        assert_eq!(back, call);
876
877        let out = ContentPart::ToolOutput {
878            name: "sum".into(),
879            output: serde_json::json!({"result":3}),
880            call_id: Some("c1".into()),
881            remote_id: None,
882        };
883        let v_out = serde_json::to_value(&out).unwrap();
884        assert_eq!(v_out.get("type").unwrap(), "ToolOutput");
885        assert_eq!(v_out.get("name").unwrap(), "sum");
886        assert_eq!(
887            v_out.get("output").unwrap(),
888            &serde_json::json!({"result":3})
889        );
890        // callId 存在
891        assert_eq!(v_out.get("callId").unwrap(), "c1");
892        let back_out: ContentPart = serde_json::from_value(v_out.clone()).unwrap();
893        assert_eq!(back_out, out);
894        let back: ContentPart = v_out.into();
895        assert_eq!(back, out);
896    }
897
898    #[test]
899    fn test_message_tool_calls_extract_from_content_parts() {
900        let parts = vec![
901            ContentPart::Text {
902                text: "hello".into(),
903            },
904            ContentPart::ToolCall {
905                name: "sum".into(),
906                args: serde_json::json!({"x":1, "y": 2}),
907                call_id: Some("abc".into()),
908            },
909            ContentPart::ToolCall {
910                name: "echo".into(),
911                args: serde_json::json!({"text":"hi"}),
912                call_id: None,
913            },
914        ];
915        let msg = Message {
916            role: "assistant".into(),
917            content: parts,
918            ..Default::default()
919        };
920        println!("{:#?}", json!(msg));
921
922        let calls = msg.tool_calls();
923        assert_eq!(calls.len(), 2);
924        assert_eq!(calls[0].name, "sum");
925        assert_eq!(calls[0].args, serde_json::json!({"x":1, "y":2}));
926        assert_eq!(calls[1].name, "echo");
927        assert_eq!(calls[1].args, serde_json::json!({"text":"hi"}));
928    }
929
930    #[test]
931    fn test_message_content_deserialize_from_string() {
932        // content as a plain string
933        let msg: Message = serde_json::from_value(serde_json::json!({
934            "role": "user",
935            "content": "hello world"
936        }))
937        .unwrap();
938        assert_eq!(msg.role, "user");
939        assert_eq!(msg.content.len(), 1);
940        assert_eq!(
941            msg.content[0],
942            ContentPart::Text {
943                text: "hello world".into()
944            }
945        );
946
947        // content as an array still works
948        let msg2: Message = serde_json::from_value(serde_json::json!({
949            "role": "assistant",
950            "content": [{"type": "Text", "text": "hi"}]
951        }))
952        .unwrap();
953        assert_eq!(msg2.content.len(), 1);
954        assert_eq!(msg2.content[0], ContentPart::Text { text: "hi".into() });
955
956        // missing content defaults to empty vec
957        let msg3: Message = serde_json::from_value(serde_json::json!({
958            "role": "system"
959        }))
960        .unwrap();
961        assert!(msg3.content.is_empty());
962    }
963
964    #[test]
965    fn test_request_meta_extra_flatten_serde() {
966        // empty extra should not serialize
967        let meta = RequestMeta {
968            engine: None,
969            user: None,
970            extra: Map::new(),
971        };
972        let v = serde_json::to_value(&meta).unwrap();
973        assert_eq!(v, serde_json::json!({}));
974
975        // extra should be flattened into the top-level object
976        let mut extra = Map::new();
977        extra.insert("foo".into(), serde_json::json!("bar"));
978        extra.insert("n".into(), serde_json::json!(1));
979        extra.insert("obj".into(), serde_json::json!({"x": true}));
980
981        let meta2 = RequestMeta {
982            engine: Some(Principal::from_text("aaaaa-aa").unwrap()),
983            user: Some("alice".into()),
984            extra,
985        };
986
987        let v2 = serde_json::to_value(&meta2).unwrap();
988        assert_eq!(v2.get("engine").unwrap(), "aaaaa-aa");
989        assert_eq!(v2.get("user").unwrap(), "alice");
990        assert_eq!(v2.get("foo").unwrap(), "bar");
991        assert_eq!(v2.get("n").unwrap(), 1);
992        assert_eq!(v2.get("obj").unwrap(), &serde_json::json!({"x": true}));
993        assert!(v2.get("extra").is_none());
994
995        // deserialization: unknown fields go into extra
996        let input = serde_json::json!({
997            "engine": "aaaaa-aa",
998            "user": "bob",
999            "k1": "v1",
1000            "k2": 2,
1001            "nested": {"a": 1}
1002        });
1003        let back: RequestMeta = serde_json::from_value(input).unwrap();
1004        assert_eq!(back.engine.unwrap().to_text(), "aaaaa-aa");
1005        assert_eq!(back.user.as_deref(), Some("bob"));
1006        assert_eq!(back.extra.get("k1").unwrap(), "v1");
1007        assert_eq!(back.extra.get("k2").unwrap(), 2);
1008        assert_eq!(
1009            back.extra.get("nested").unwrap(),
1010            &serde_json::json!({"a": 1})
1011        );
1012
1013        // round-trip (field-by-field)
1014        let back2: RequestMeta = serde_json::from_value(v2).unwrap();
1015        assert_eq!(back2.engine.unwrap().to_text(), "aaaaa-aa");
1016        assert_eq!(back2.user.as_deref(), Some("alice"));
1017        assert_eq!(back2.extra.get("foo").unwrap(), "bar");
1018        assert_eq!(back2.extra.get("n").unwrap(), 1);
1019        assert_eq!(
1020            back2.extra.get("obj").unwrap(),
1021            &serde_json::json!({"x": true})
1022        );
1023    }
1024}
anda_core/model.rs

anda_core/
model.rs
