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(rename = "functions.alpha_scalar.LeafTaskExpression")]
65pub enum LeafTaskExpression {
66    #[serde(rename = "vector.completion")]
67    VectorCompletion(VectorCompletionTaskExpression),
68}
69
70impl LeafTaskExpression {
71    pub fn transpile(self) -> functions::TaskExpression {
72        match self {
73            LeafTaskExpression::VectorCompletion(task) => {
74                functions::TaskExpression::VectorCompletion(task.transpile())
75            }
76        }
77    }
78}
79
80#[derive(
81    Debug,
82    Clone,
83    PartialEq,
84    Serialize,
85    Deserialize,
86    JsonSchema,
87    arbitrary::Arbitrary,
88)]
89#[schemars(rename = "functions.alpha_scalar.ScalarFunctionTaskExpression")]
90pub struct ScalarFunctionTaskExpression {
91    #[serde(flatten)]
92    #[schemars(schema_with = "crate::flatten_schema::<crate::RemotePath>")]
93    pub path: crate::RemotePath,
94    #[serde(skip_serializing_if = "Option::is_none")]
95    #[schemars(extend("omitempty" = true))]
96    pub skip: Option<functions::expression::Expression>,
97    pub input: super::expression::ScalarFunctionInputValueExpression,
98}
99
100impl ScalarFunctionTaskExpression {
101    pub fn url(&self) -> String {
102        self.path.url()
103    }
104
105    pub fn transpile(self) -> functions::ScalarFunctionTaskExpression {
106        functions::ScalarFunctionTaskExpression {
107            path: self.path,
108            skip: self.skip,
109            map: None,
110            input:
111                super::expression::scalar_function_input_value_expression::transpile(
112                    self.input,
113                ),
114            output: functions::expression::Expression::Special(
115                functions::expression::Special::Output,
116            ),
117        }
118    }
119}
120
121#[derive(
122    Debug,
123    Clone,
124    PartialEq,
125    Serialize,
126    Deserialize,
127    JsonSchema,
128    arbitrary::Arbitrary,
129)]
130#[schemars(
131    rename = "functions.alpha_scalar.PlaceholderScalarFunctionTaskExpression"
132)]
133pub struct PlaceholderScalarFunctionTaskExpression {
134    pub input_schema: super::expression::ScalarFunctionInputSchema,
135    #[serde(skip_serializing_if = "Option::is_none")]
136    #[schemars(extend("omitempty" = true))]
137    pub skip: Option<functions::expression::Expression>,
138    pub input: super::expression::ScalarFunctionInputValueExpression,
139}
140
141impl PlaceholderScalarFunctionTaskExpression {
142    pub fn transpile(
143        self,
144    ) -> functions::PlaceholderScalarFunctionTaskExpression {
145        functions::PlaceholderScalarFunctionTaskExpression {
146            input_schema:
147                super::expression::scalar_function_input_schema::transpile(
148                    self.input_schema,
149                ),
150            skip: self.skip,
151            map: None,
152            input:
153                super::expression::scalar_function_input_value_expression::transpile(
154                    self.input,
155                ),
156            output: functions::expression::Expression::Special(
157                functions::expression::Special::Output,
158            ),
159        }
160    }
161
162    pub fn replace(
163        self,
164        path: &crate::RemotePath,
165    ) -> ScalarFunctionTaskExpression {
166        ScalarFunctionTaskExpression {
167            path: path.clone(),
168            skip: self.skip,
169            input: self.input,
170        }
171    }
172}
173
174#[derive(
175    Debug,
176    Clone,
177    PartialEq,
178    Serialize,
179    Deserialize,
180    JsonSchema,
181    arbitrary::Arbitrary,
182)]
183#[schemars(rename = "functions.alpha_scalar.VectorCompletionTaskExpression")]
184pub struct VectorCompletionTaskExpression {
185    #[serde(skip_serializing_if = "Option::is_none")]
186    #[schemars(extend("omitempty" = true))]
187    pub skip: Option<functions::expression::Expression>,
188    pub messages: functions::expression::Expression,
189    pub responses: Vec<agent::completions::message::RichContent>,
190}
191
192impl VectorCompletionTaskExpression {
193    pub fn transpile(self) -> functions::VectorCompletionTaskExpression {
194        functions::VectorCompletionTaskExpression {
195            skip: self.skip,
196            map: None,
197            messages: functions::expression::WithExpression::Expression(
198                self.messages,
199            ),
200            responses: functions::expression::WithExpression::Value(
201                self.responses
202                    .into_iter()
203                    .map(agent::completions::message::RichContentExpression::from)
204                    .map(functions::expression::WithExpression::Value)
205                    .collect(),
206            ),
207            output: functions::expression::Expression::Special(
208                functions::expression::Special::TaskOutputWeightedSum,
209            ),
210        }
211    }
212}