Skip to main content

objectiveai_sdk/functions/alpha_scalar/
task.rs

1use crate::{agent, functions};
2use schemars::JsonSchema;
3use serde::{Deserialize, Serialize};
4
5#[derive(
6    Debug,
7    Clone,
8    PartialEq,
9    Serialize,
10    Deserialize,
11    JsonSchema,
12    arbitrary::Arbitrary,
13)]
14#[serde(tag = "type")]
15#[schemars(rename = "functions.alpha_scalar.BranchTaskExpression")]
16pub enum BranchTaskExpression {
17    #[schemars(title = "ScalarFunction")]
18    #[serde(rename = "alpha.scalar.function")]
19    ScalarFunction(ScalarFunctionTaskExpression),
20    #[schemars(title = "PlaceholderScalarFunction")]
21    #[serde(rename = "placeholder.alpha.scalar.function")]
22    PlaceholderScalarFunction(PlaceholderScalarFunctionTaskExpression),
23}
24
25impl BranchTaskExpression {
26    pub fn url(&self) -> Option<String> {
27        match self {
28            BranchTaskExpression::ScalarFunction(task) => Some(task.url()),
29            BranchTaskExpression::PlaceholderScalarFunction(_) => None,
30        }
31    }
32
33    pub fn transpile(self) -> functions::TaskExpression {
34        match self {
35            BranchTaskExpression::ScalarFunction(task) => {
36                functions::TaskExpression::ScalarFunction(task.transpile())
37            }
38            BranchTaskExpression::PlaceholderScalarFunction(task) => {
39                functions::TaskExpression::PlaceholderScalarFunction(
40                    task.transpile(),
41                )
42            }
43        }
44    }
45
46    pub fn is_placeholder(&self) -> bool {
47        match self {
48            BranchTaskExpression::ScalarFunction(_) => false,
49            BranchTaskExpression::PlaceholderScalarFunction(_) => true,
50        }
51    }
52}
53
54#[derive(
55    Debug,
56    Clone,
57    PartialEq,
58    Serialize,
59    Deserialize,
60    JsonSchema,
61    arbitrary::Arbitrary,
62)]
63#[serde(tag = "type")]
64#[schemars(
65    rename = "functions.alpha_scalar.PartialPlaceholderBranchTaskExpression"
66)]
67pub enum PartialPlaceholderBranchTaskExpression {
68    #[serde(rename = "placeholder.alpha.scalar.function")]
69    PlaceholderScalarFunction(PartialPlaceholderScalarFunctionTaskExpression),
70}
71
72impl PartialPlaceholderBranchTaskExpression {
73    pub fn complete(
74        self,
75        name: String,
76        depth: u64,
77        min_branch_width: u64,
78        max_branch_width: u64,
79        min_leaf_width: u64,
80        max_leaf_width: u64,
81    ) -> BranchTaskExpression {
82        match self {
83            PartialPlaceholderBranchTaskExpression::PlaceholderScalarFunction(
84                task,
85            ) => BranchTaskExpression::PlaceholderScalarFunction(
86                task.complete(
87                    name,
88                    depth,
89                    min_branch_width,
90                    max_branch_width,
91                    min_leaf_width,
92                    max_leaf_width,
93                ),
94            ),
95        }
96    }
97}
98
99#[derive(
100    Debug,
101    Clone,
102    PartialEq,
103    Serialize,
104    Deserialize,
105    JsonSchema,
106    arbitrary::Arbitrary,
107)]
108#[serde(tag = "type")]
109#[schemars(rename = "functions.alpha_scalar.LeafTaskExpression")]
110pub enum LeafTaskExpression {
111    #[serde(rename = "vector.completion")]
112    VectorCompletion(VectorCompletionTaskExpression),
113}
114
115impl LeafTaskExpression {
116    pub fn transpile(self) -> functions::TaskExpression {
117        match self {
118            LeafTaskExpression::VectorCompletion(task) => {
119                functions::TaskExpression::VectorCompletion(task.transpile())
120            }
121        }
122    }
123}
124
125#[derive(
126    Debug,
127    Clone,
128    PartialEq,
129    Serialize,
130    Deserialize,
131    JsonSchema,
132    arbitrary::Arbitrary,
133)]
134#[schemars(rename = "functions.alpha_scalar.ScalarFunctionTaskExpression")]
135pub struct ScalarFunctionTaskExpression {
136    #[serde(flatten)]
137    #[schemars(schema_with = "crate::flatten_schema::<crate::RemotePath>")]
138    pub path: crate::RemotePath,
139    #[serde(skip_serializing_if = "Option::is_none")]
140    #[schemars(extend("omitempty" = true))]
141    pub skip: Option<functions::expression::Expression>,
142    pub input: super::expression::ScalarFunctionInputValueExpression,
143}
144
145impl ScalarFunctionTaskExpression {
146    pub fn url(&self) -> String {
147        self.path.url()
148    }
149
150    pub fn transpile(self) -> functions::ScalarFunctionTaskExpression {
151        functions::ScalarFunctionTaskExpression {
152            path: self.path,
153            skip: self.skip,
154            map: None,
155            input:
156                super::expression::scalar_function_input_value_expression::transpile(
157                    self.input,
158                ),
159            output: functions::expression::Expression::Special(
160                functions::expression::Special::Output,
161            ),
162        }
163    }
164}
165
166#[derive(
167    Debug,
168    Clone,
169    PartialEq,
170    Serialize,
171    Deserialize,
172    JsonSchema,
173    arbitrary::Arbitrary,
174)]
175#[schemars(
176    rename = "functions.alpha_scalar.PlaceholderScalarFunctionTaskExpression"
177)]
178pub struct PlaceholderScalarFunctionTaskExpression {
179    #[serde(flatten)]
180    pub params: functions::inventions::Params,
181    pub input_schema: super::expression::ScalarFunctionInputSchema,
182    #[serde(skip_serializing_if = "Option::is_none")]
183    #[schemars(extend("omitempty" = true))]
184    pub skip: Option<functions::expression::Expression>,
185    pub input: super::expression::ScalarFunctionInputValueExpression,
186}
187
188impl PlaceholderScalarFunctionTaskExpression {
189    pub fn transpile(
190        self,
191    ) -> functions::PlaceholderScalarFunctionTaskExpression {
192        functions::PlaceholderScalarFunctionTaskExpression {
193            input_schema:
194                super::expression::scalar_function_input_schema::transpile(
195                    self.input_schema,
196                ),
197            skip: self.skip,
198            map: None,
199            input:
200                super::expression::scalar_function_input_value_expression::transpile(
201                    self.input,
202                ),
203            output: functions::expression::Expression::Special(
204                functions::expression::Special::Output,
205            ),
206        }
207    }
208
209    pub fn replace(
210        self,
211        path: &crate::RemotePath,
212    ) -> ScalarFunctionTaskExpression {
213        ScalarFunctionTaskExpression {
214            path: path.clone(),
215            skip: self.skip,
216            input: self.input,
217        }
218    }
219}
220
221#[derive(
222    Debug,
223    Clone,
224    PartialEq,
225    Serialize,
226    Deserialize,
227    JsonSchema,
228    arbitrary::Arbitrary,
229)]
230#[schemars(
231    rename = "functions.alpha_scalar.PartialPlaceholderScalarFunctionTaskExpression"
232)]
233pub struct PartialPlaceholderScalarFunctionTaskExpression {
234    pub spec: String,
235    pub input_schema: super::expression::ScalarFunctionInputSchema,
236    #[serde(skip_serializing_if = "Option::is_none")]
237    #[schemars(extend("omitempty" = true))]
238    pub skip: Option<functions::expression::Expression>,
239    pub input: super::expression::ScalarFunctionInputValueExpression,
240}
241
242impl PartialPlaceholderScalarFunctionTaskExpression {
243    pub fn complete(
244        self,
245        name: String,
246        depth: u64,
247        min_branch_width: u64,
248        max_branch_width: u64,
249        min_leaf_width: u64,
250        max_leaf_width: u64,
251    ) -> PlaceholderScalarFunctionTaskExpression {
252        PlaceholderScalarFunctionTaskExpression {
253            params: functions::inventions::Params {
254                depth,
255                min_branch_width,
256                max_branch_width,
257                min_leaf_width,
258                max_leaf_width,
259                name,
260                spec: self.spec,
261            },
262            input_schema: self.input_schema,
263            skip: self.skip,
264            input: self.input,
265        }
266    }
267}
268
269#[derive(
270    Debug,
271    Clone,
272    PartialEq,
273    Serialize,
274    Deserialize,
275    JsonSchema,
276    arbitrary::Arbitrary,
277)]
278#[schemars(rename = "functions.alpha_scalar.VectorCompletionTaskExpression")]
279pub struct VectorCompletionTaskExpression {
280    #[serde(skip_serializing_if = "Option::is_none")]
281    #[schemars(extend("omitempty" = true))]
282    pub skip: Option<functions::expression::Expression>,
283    pub messages: functions::expression::Expression,
284    pub responses: Vec<agent::completions::message::RichContent>,
285}
286
287impl VectorCompletionTaskExpression {
288    pub fn transpile(self) -> functions::VectorCompletionTaskExpression {
289        functions::VectorCompletionTaskExpression {
290            skip: self.skip,
291            map: None,
292            messages: functions::expression::WithExpression::Expression(
293                self.messages,
294            ),
295            responses: functions::expression::WithExpression::Value(
296                self.responses
297                    .into_iter()
298                    .map(agent::completions::message::RichContentExpression::from)
299                    .map(functions::expression::WithExpression::Value)
300                    .collect(),
301            ),
302            output: functions::expression::Expression::Special(
303                functions::expression::Special::TaskOutputWeightedSum,
304            ),
305        }
306    }
307}