Skip to main content

objectiveai_sdk/agent/completions/message/
tool_message.rs

1//! Tool message types.
2
3use super::rich_content::{RichContent, RichContentExpression};
4use crate::functions;
5use functions::expression::{
6    ExpressionError, FromStarlarkValue, WithExpression,
7};
8use schemars::JsonSchema;
9use serde::{Deserialize, Serialize};
10use starlark::values::dict::DictRef as StarlarkDictRef;
11use starlark::values::{UnpackValue, Value as StarlarkValue};
12
13/// Vendor-extension metadata attached to a tool response. The
14/// `objectiveai-mcp-proxy` populates known keys (currently
15/// `notifications`); the SDK lossy-decodes the MCP `_meta` bag into
16/// this typed shape. Unknown keys are dropped.
17#[derive(
18    Debug,
19    Clone,
20    PartialEq,
21    Default,
22    Serialize,
23    Deserialize,
24    JsonSchema,
25    arbitrary::Arbitrary,
26)]
27#[schemars(rename = "agent.completions.message.ToolResponseMetadata")]
28pub struct ToolResponseMetadata {
29    /// Count of pending notifications the proxy drained and prepended
30    /// to the tool response's `content` before returning. Only set
31    /// when at least one notification was drained.
32    #[serde(skip_serializing_if = "Option::is_none")]
33    #[schemars(extend("omitempty" = true))]
34    #[arbitrary(with = crate::arbitrary_util::arbitrary_option_u64)]
35    pub notifications: Option<u64>,
36}
37
38/// A tool message containing the result of a tool call.
39#[derive(
40    Debug,
41    Clone,
42    PartialEq,
43    Serialize,
44    Deserialize,
45    JsonSchema,
46    arbitrary::Arbitrary,
47)]
48#[schemars(rename = "agent.completions.message.ToolMessage")]
49pub struct ToolMessage {
50    /// The content of the tool response.
51    pub content: RichContent,
52    /// The ID of the tool call this message responds to.
53    pub tool_call_id: String,
54    /// Optional vendor-extension metadata, populated by
55    /// `objectiveai-mcp-proxy` via MCP's `_meta` extension bag.
56    #[serde(skip_serializing_if = "Option::is_none")]
57    #[schemars(extend("omitempty" = true))]
58    pub metadata: Option<ToolResponseMetadata>,
59}
60
61impl ToolMessage {
62    /// Prepares the message by normalizing its content.
63    pub fn prepare(&mut self) {
64        self.content.prepare();
65    }
66
67}
68
69impl FromStarlarkValue for ToolMessage {
70    fn from_starlark_value(
71        value: &StarlarkValue,
72    ) -> Result<Self, ExpressionError> {
73        let dict = StarlarkDictRef::from_value(*value).ok_or_else(|| {
74            ExpressionError::StarlarkConversionError(
75                "ToolMessage: expected dict".into(),
76            )
77        })?;
78        let mut content = None;
79        let mut tool_call_id = None;
80        for (k, v) in dict.iter() {
81            let key = <&str as UnpackValue>::unpack_value(k)
82                .map_err(|e| {
83                    ExpressionError::StarlarkConversionError(e.to_string())
84                })?
85                .ok_or_else(|| {
86                    ExpressionError::StarlarkConversionError(
87                        "ToolMessage: expected string key".into(),
88                    )
89                })?;
90            match key {
91                "content" => {
92                    content = Some(RichContent::from_starlark_value(&v)?)
93                }
94                "tool_call_id" => {
95                    tool_call_id = Some(String::from_starlark_value(&v)?)
96                }
97                _ => {}
98            }
99            if content.is_some() && tool_call_id.is_some() {
100                break;
101            }
102        }
103        Ok(ToolMessage {
104            content: content.ok_or_else(|| {
105                ExpressionError::StarlarkConversionError(
106                    "ToolMessage: missing content".into(),
107                )
108            })?,
109            tool_call_id: tool_call_id.ok_or_else(|| {
110                ExpressionError::StarlarkConversionError(
111                    "ToolMessage: missing tool_call_id".into(),
112                )
113            })?,
114            metadata: None,
115        })
116    }
117}
118
119/// Expression variant of [`ToolMessage`] for dynamic content.
120#[derive(
121    Debug,
122    Clone,
123    PartialEq,
124    Serialize,
125    Deserialize,
126    JsonSchema,
127    arbitrary::Arbitrary,
128)]
129#[schemars(rename = "agent.completions.message.ToolMessageExpression")]
130pub struct ToolMessageExpression {
131    /// The content expression.
132    pub content: functions::expression::WithExpression<RichContentExpression>,
133    /// The tool call ID expression.
134    pub tool_call_id: functions::expression::WithExpression<String>,
135}
136
137impl ToolMessageExpression {
138    /// Compiles the expression into a concrete [`ToolMessage`].
139    pub fn compile(
140        self,
141        params: &functions::expression::Params,
142    ) -> Result<ToolMessage, functions::expression::ExpressionError> {
143        let content = self.content.compile_one(params)?.compile(params)?;
144        let tool_call_id = self.tool_call_id.compile_one(params)?;
145        Ok(ToolMessage {
146            content,
147            tool_call_id,
148            metadata: None,
149        })
150    }
151}
152
153impl FromStarlarkValue for ToolMessageExpression {
154    fn from_starlark_value(
155        value: &StarlarkValue,
156    ) -> Result<Self, ExpressionError> {
157        let dict = StarlarkDictRef::from_value(*value).ok_or_else(|| {
158            ExpressionError::StarlarkConversionError(
159                "ToolMessageExpression: expected dict".into(),
160            )
161        })?;
162        let mut content = None;
163        let mut tool_call_id = None;
164        for (k, v) in dict.iter() {
165            let key = <&str as UnpackValue>::unpack_value(k)
166                .map_err(|e| {
167                    ExpressionError::StarlarkConversionError(e.to_string())
168                })?
169                .ok_or_else(|| {
170                    ExpressionError::StarlarkConversionError(
171                        "ToolMessageExpression: expected string key".into(),
172                    )
173                })?;
174            match key {
175                "content" => {
176                    content = Some(WithExpression::Value(
177                        RichContentExpression::from_starlark_value(&v)?,
178                    ))
179                }
180                "tool_call_id" => {
181                    tool_call_id = Some(WithExpression::Value(
182                        String::from_starlark_value(&v)?,
183                    ))
184                }
185                _ => {}
186            }
187            if content.is_some() && tool_call_id.is_some() {
188                break;
189            }
190        }
191        Ok(ToolMessageExpression {
192            content: content.ok_or_else(|| {
193                ExpressionError::StarlarkConversionError(
194                    "ToolMessageExpression: missing content".into(),
195                )
196            })?,
197            tool_call_id: tool_call_id.ok_or_else(|| {
198                ExpressionError::StarlarkConversionError(
199                    "ToolMessageExpression: missing tool_call_id".into(),
200                )
201            })?,
202        })
203    }
204}