Skip to main content

zai_rs/
agent.rs

1//! # Agent v1 wire contracts
2//!
3//! Strongly typed request and response values for the three frozen Agent v1
4//! operations:
5//!
6//! - `POST /v1/agents`
7//! - `POST /v1/agents/async-result`
8//! - `POST /v1/agents/conversation`
9//!
10//! This module intentionally contains no network facade. Use these values with
11//! a transport that targets the corresponding Agent v1 routes.
12
13mod request;
14mod response;
15
16pub use request::*;
17pub use response::*;
18
19use std::marker::PhantomData;
20
21use serde::{Deserialize, Serialize};
22
23use crate::{ZaiResult, client::validation::require_non_blank};
24
25/// Open variables object accepted by the Agent invocation schema.
26///
27/// This is the only open JSON map in the Agent v1 request contracts.
28#[derive(Clone, Default, Serialize, Deserialize)]
29#[serde(transparent)]
30pub struct AgentCustomVariables(serde_json::Map<String, serde_json::Value>);
31
32impl std::fmt::Debug for AgentCustomVariables {
33    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34        formatter
35            .debug_struct("AgentCustomVariables")
36            .field("keys", &self.0.keys().collect::<Vec<_>>())
37            .field("values", &"[REDACTED]")
38            .finish()
39    }
40}
41
42impl AgentCustomVariables {
43    /// Construct an empty variables object.
44    pub fn new() -> Self {
45        Self(serde_json::Map::new())
46    }
47
48    /// Insert an open-schema variable value.
49    pub fn insert(
50        &mut self,
51        key: impl Into<String>,
52        value: serde_json::Value,
53    ) -> Option<serde_json::Value> {
54        self.0.insert(key.into(), value)
55    }
56
57    /// Return whether no variables are present.
58    pub fn is_empty(&self) -> bool {
59        self.0.is_empty()
60    }
61
62    /// Borrow the underlying open map.
63    pub const fn as_map(&self) -> &serde_json::Map<String, serde_json::Value> {
64        &self.0
65    }
66}
67
68impl From<serde_json::Map<String, serde_json::Value>> for AgentCustomVariables {
69    fn from(value: serde_json::Map<String, serde_json::Value>) -> Self {
70        Self(value)
71    }
72}
73
74/// Type-state marker for a non-streaming Agent invocation.
75#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)]
76pub struct NonStreaming;
77
78/// Type-state marker for a streaming Agent invocation.
79#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)]
80pub struct Streaming;
81
82/// Sealed mapping from invocation type-state to its wire-level `stream` value.
83pub trait StreamMode: private::Sealed {
84    /// Boolean serialized into the request body.
85    const WIRE: bool;
86}
87
88impl StreamMode for NonStreaming {
89    const WIRE: bool = false;
90}
91
92impl StreamMode for Streaming {
93    const WIRE: bool = true;
94}
95
96mod private {
97    pub trait Sealed {}
98    impl Sealed for super::NonStreaming {}
99    impl Sealed for super::Streaming {}
100}
101
102/// Builder for an [`AgentInvokeRequest`].
103pub struct AgentInvokeRequestBuilder<N: StreamMode> {
104    agent_id: AgentId,
105    messages: Vec<AgentMessage>,
106    custom_variables: AgentCustomVariables,
107    marker: PhantomData<N>,
108}
109
110impl<N: StreamMode> AgentInvokeRequestBuilder<N> {
111    /// Append one message.
112    pub fn message(mut self, message: AgentMessage) -> Self {
113        self.messages.push(message);
114        self
115    }
116
117    /// Append multiple messages in order.
118    pub fn messages(mut self, messages: impl IntoIterator<Item = AgentMessage>) -> Self {
119        self.messages.extend(messages);
120        self
121    }
122
123    /// Replace the open variables object.
124    pub fn custom_variables(mut self, variables: impl Into<AgentCustomVariables>) -> Self {
125        self.custom_variables = variables.into();
126        self
127    }
128
129    /// Validate the required non-empty message list and build the request.
130    pub fn build(self) -> ZaiResult<AgentInvokeRequest<N>> {
131        if self.messages.is_empty() {
132            return Err(crate::client::validation::invalid(
133                "messages must contain at least one item",
134            ));
135        }
136        Ok(AgentInvokeRequest {
137            agent_id: self.agent_id,
138            stream: N::WIRE,
139            messages: self.messages,
140            custom_variables: self.custom_variables,
141            marker: PhantomData,
142        })
143    }
144
145    /// Switch this request builder to streaming output.
146    pub fn streaming(self) -> AgentInvokeRequestBuilder<Streaming> {
147        AgentInvokeRequestBuilder {
148            agent_id: self.agent_id,
149            messages: self.messages,
150            custom_variables: self.custom_variables,
151            marker: PhantomData,
152        }
153    }
154}
155
156/// Frozen request body for `POST /v1/agents`.
157#[derive(Clone, Serialize)]
158pub struct AgentInvokeRequest<N: StreamMode> {
159    agent_id: AgentId,
160    stream: bool,
161    messages: Vec<AgentMessage>,
162    #[serde(skip_serializing_if = "AgentCustomVariables::is_empty")]
163    custom_variables: AgentCustomVariables,
164    #[serde(skip)]
165    marker: PhantomData<N>,
166}
167
168impl<N: StreamMode> std::fmt::Debug for AgentInvokeRequest<N> {
169    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
170        let custom_variables = if self.custom_variables.is_empty() {
171            "none"
172        } else {
173            "[REDACTED]"
174        };
175        formatter
176            .debug_struct("AgentInvokeRequest")
177            .field("agent_id", &"[REDACTED]")
178            .field("stream", &self.stream)
179            .field("message_count", &self.messages.len())
180            .field("custom_variables", &custom_variables)
181            .finish()
182    }
183}
184
185impl<N: StreamMode> AgentInvokeRequest<N> {
186    /// Start an invocation request for a frozen Agent v1 identifier.
187    pub fn builder(agent_id: AgentId) -> AgentInvokeRequestBuilder<N> {
188        AgentInvokeRequestBuilder {
189            agent_id,
190            messages: Vec::new(),
191            custom_variables: AgentCustomVariables::new(),
192            marker: PhantomData,
193        }
194    }
195
196    /// Return the selected agent identifier.
197    pub const fn agent_id(&self) -> AgentId {
198        self.agent_id
199    }
200
201    /// Return the resolved streaming flag.
202    pub const fn stream(&self) -> bool {
203        self.stream
204    }
205
206    /// Borrow request messages.
207    pub fn messages(&self) -> &[AgentMessage] {
208        &self.messages
209    }
210
211    /// Borrow the open variables object.
212    pub const fn custom_variables(&self) -> &AgentCustomVariables {
213        &self.custom_variables
214    }
215}
216
217/// Frozen request body for `POST /v1/agents/async-result`.
218#[derive(Clone, Serialize)]
219pub struct AgentAsyncResultRequest {
220    async_id: String,
221    agent_id: String,
222}
223
224impl std::fmt::Debug for AgentAsyncResultRequest {
225    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
226        formatter
227            .debug_struct("AgentAsyncResultRequest")
228            .field("async_id", &"[REDACTED]")
229            .field("agent_id", &"[REDACTED]")
230            .finish()
231    }
232}
233
234impl AgentAsyncResultRequest {
235    /// Validate the two required identifiers and construct the request.
236    pub fn new(agent_id: impl Into<String>, async_id: impl Into<String>) -> ZaiResult<Self> {
237        let agent_id = agent_id.into();
238        let async_id = async_id.into();
239        require_non_blank(&agent_id, "agent_id")?;
240        require_non_blank(&async_id, "async_id")?;
241        Ok(Self { async_id, agent_id })
242    }
243
244    /// Borrow the agent identifier.
245    pub fn agent_id(&self) -> &str {
246        &self.agent_id
247    }
248
249    /// Borrow the asynchronous task identifier.
250    pub fn async_id(&self) -> &str {
251        &self.async_id
252    }
253}
254
255/// One slide page descriptor in conversation custom variables.
256#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
257#[serde(deny_unknown_fields)]
258pub struct AgentSlidePage {
259    #[serde(skip_serializing_if = "Option::is_none")]
260    position: Option<f64>,
261    #[serde(skip_serializing_if = "Option::is_none")]
262    width: Option<f64>,
263    #[serde(skip_serializing_if = "Option::is_none")]
264    height: Option<f64>,
265}
266
267impl AgentSlidePage {
268    /// Construct a fully specified slide page descriptor.
269    pub const fn new(position: f64, width: f64, height: f64) -> Self {
270        Self {
271            position: Some(position),
272            width: Some(width),
273            height: Some(height),
274        }
275    }
276
277    /// Return the slide position.
278    pub const fn position(&self) -> Option<f64> {
279        self.position
280    }
281
282    /// Return the slide width in centimetres.
283    pub const fn width(&self) -> Option<f64> {
284        self.width
285    }
286
287    /// Return the slide height in centimetres.
288    pub const fn height(&self) -> Option<f64> {
289        self.height
290    }
291}
292
293/// Closed custom variables accepted by the slide conversation endpoint.
294#[derive(Debug, Clone, Default, Serialize, Deserialize)]
295#[serde(deny_unknown_fields)]
296pub struct AgentConversationVariables {
297    #[serde(skip_serializing_if = "Option::is_none")]
298    include_pdf: Option<bool>,
299    #[serde(skip_serializing_if = "Option::is_none")]
300    pages: Option<Vec<AgentSlidePage>>,
301}
302
303impl AgentConversationVariables {
304    /// Construct empty optional conversation variables.
305    pub const fn new() -> Self {
306        Self {
307            include_pdf: None,
308            pages: None,
309        }
310    }
311
312    /// Configure whether a PDF export is included.
313    pub const fn with_include_pdf(mut self, include_pdf: bool) -> Self {
314        self.include_pdf = Some(include_pdf);
315        self
316    }
317
318    /// Configure slide page descriptors.
319    pub fn with_pages(mut self, pages: impl IntoIterator<Item = AgentSlidePage>) -> Self {
320        self.pages = Some(pages.into_iter().collect());
321        self
322    }
323
324    /// Return the configured PDF flag.
325    pub const fn include_pdf(&self) -> Option<bool> {
326        self.include_pdf
327    }
328
329    /// Borrow configured page descriptors.
330    pub fn pages(&self) -> Option<&[AgentSlidePage]> {
331        self.pages.as_deref()
332    }
333}
334
335/// Frozen request body for `POST /v1/agents/conversation`.
336#[derive(Clone, Serialize)]
337pub struct AgentConversationRequest {
338    agent_id: String,
339    conversation_id: String,
340    #[serde(skip_serializing_if = "Option::is_none")]
341    custom_variables: Option<AgentConversationVariables>,
342}
343
344impl std::fmt::Debug for AgentConversationRequest {
345    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
346        formatter
347            .debug_struct("AgentConversationRequest")
348            .field("agent_id", &"[REDACTED]")
349            .field("conversation_id", &"[REDACTED]")
350            .field(
351                "custom_variables",
352                &self.custom_variables.as_ref().map(|_| "[REDACTED]"),
353            )
354            .finish()
355    }
356}
357
358impl AgentConversationRequest {
359    /// Validate identifiers and construct a conversation-continuation request.
360    pub fn new(agent_id: impl Into<String>, conversation_id: impl Into<String>) -> ZaiResult<Self> {
361        let agent_id = agent_id.into();
362        let conversation_id = conversation_id.into();
363        require_non_blank(&agent_id, "agent_id")?;
364        require_non_blank(&conversation_id, "conversation_id")?;
365        Ok(Self {
366            agent_id,
367            conversation_id,
368            custom_variables: None,
369        })
370    }
371
372    /// Set the endpoint's closed custom-variables object.
373    pub fn with_custom_variables(mut self, variables: AgentConversationVariables) -> Self {
374        self.custom_variables = Some(variables);
375        self
376    }
377
378    /// Borrow the agent identifier.
379    pub fn agent_id(&self) -> &str {
380        &self.agent_id
381    }
382
383    /// Borrow the conversation identifier.
384    pub fn conversation_id(&self) -> &str {
385        &self.conversation_id
386    }
387
388    /// Borrow the optional custom variables.
389    pub const fn custom_variables(&self) -> Option<&AgentConversationVariables> {
390        self.custom_variables.as_ref()
391    }
392}
393
394#[cfg(test)]
395mod tests {
396    use super::*;
397
398    #[test]
399    fn invoke_request_serializes_the_exact_nonstreaming_schema() {
400        let request = AgentInvokeRequest::<NonStreaming>::builder(AgentId::GeneralTranslation)
401            .message(AgentMessage::user("hello"))
402            .build()
403            .unwrap();
404
405        assert_eq!(
406            serde_json::to_value(&request).unwrap(),
407            serde_json::json!({
408                "agent_id": "general_translation",
409                "stream": false,
410                "messages": [{"role": "user", "content": "hello"}]
411            })
412        );
413        assert!(!format!("{request:?}").contains("hello"));
414    }
415
416    #[test]
417    fn invoke_request_serializes_closed_multimodal_parts_and_open_variables() {
418        let mut variables = AgentCustomVariables::new();
419        variables.insert("target_lang", serde_json::json!("en"));
420        let request = AgentInvokeRequest::<NonStreaming>::builder(AgentId::GeneralTranslation)
421            .message(AgentMessage::user(vec![
422                AgentRequestContentPart::text("translate this"),
423                AgentRequestContentPart::file_id("file-1"),
424                AgentRequestContentPart::file_url("https://example.test/report.pdf"),
425                AgentRequestContentPart::image_url("https://example.test/image.png"),
426            ]))
427            .custom_variables(variables)
428            .streaming()
429            .build()
430            .unwrap();
431
432        assert_eq!(
433            serde_json::to_value(request).unwrap(),
434            serde_json::json!({
435                "agent_id": "general_translation",
436                "stream": true,
437                "messages": [{
438                    "role": "user",
439                    "content": [
440                        {"type": "text", "text": "translate this"},
441                        {"type": "file_id", "file_id": "file-1"},
442                        {"type": "file_url", "file_url": "https://example.test/report.pdf"},
443                        {"type": "image_url", "image_url": "https://example.test/image.png"}
444                    ]
445                }],
446                "custom_variables": {"target_lang": "en"}
447            })
448        );
449    }
450
451    #[test]
452    fn invoke_request_requires_at_least_one_message() {
453        assert!(
454            AgentInvokeRequest::<NonStreaming>::builder(AgentId::GeneralTranslation)
455                .build()
456                .is_err()
457        );
458    }
459
460    #[test]
461    fn async_result_request_is_exact_and_redacted() {
462        assert!(AgentAsyncResultRequest::new(" ", "task-1").is_err());
463        assert!(AgentAsyncResultRequest::new("agent-1", " ").is_err());
464        let request = AgentAsyncResultRequest::new("agent-1", "task-1").unwrap();
465        assert_eq!(
466            serde_json::to_value(&request).unwrap(),
467            serde_json::json!({"async_id": "task-1", "agent_id": "agent-1"})
468        );
469        let debug = format!("{request:?}");
470        assert!(!debug.contains("agent-1"));
471        assert!(!debug.contains("task-1"));
472    }
473
474    #[test]
475    fn conversation_request_uses_custom_variables_instead_of_messages() {
476        let variables = AgentConversationVariables::new()
477            .with_include_pdf(true)
478            .with_pages([AgentSlidePage::new(1.0, 25.4, 14.29)]);
479        let request = AgentConversationRequest::new("slides_glm_agent", "conversation-1")
480            .unwrap()
481            .with_custom_variables(variables);
482
483        assert_eq!(
484            serde_json::to_value(&request).unwrap(),
485            serde_json::json!({
486                "agent_id": "slides_glm_agent",
487                "conversation_id": "conversation-1",
488                "custom_variables": {
489                    "include_pdf": true,
490                    "pages": [{"position": 1.0, "width": 25.4, "height": 14.29}]
491                }
492            })
493        );
494        let debug = format!("{request:?}");
495        assert!(!debug.contains("slides_glm_agent"));
496        assert!(!debug.contains("conversation-1"));
497    }
498}