Skip to main content

objectiveai_sdk/functions/alpha_scalar/
function.rs

1use crate::functions;
2use serde::{Deserialize, Serialize};
3use schemars::JsonSchema;
4
5#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, arbitrary::Arbitrary)]
6#[serde(tag = "type")]
7#[schemars(rename = "functions.alpha_scalar.RemoteFunction")]
8pub enum RemoteFunction {
9    #[schemars(title = "Branch")]
10    #[serde(rename = "alpha.scalar.branch.function")]
11    Branch {
12        description: String,
13        input_schema: super::expression::ScalarFunctionInputSchema,
14        tasks: Vec<super::BranchTaskExpression>,
15    },
16    #[schemars(title = "Leaf")]
17    #[serde(rename = "alpha.scalar.leaf.function")]
18    Leaf {
19        description: String,
20        input_schema: super::expression::ScalarFunctionInputSchema,
21        tasks: Vec<super::LeafTaskExpression>,
22    },
23}
24
25impl RemoteFunction {
26    pub fn tasks(&self) -> &[super::BranchTaskExpression] {
27        match self {
28            RemoteFunction::Branch { tasks, .. } => tasks,
29            RemoteFunction::Leaf { .. } => &[],
30        }
31    }
32
33    pub fn remotes(&self) -> impl Iterator<Item = &crate::RemotePath> {
34        self.tasks().iter().filter_map(|task| match task {
35            super::BranchTaskExpression::ScalarFunction(t) => Some(&t.path),
36            _ => None,
37        })
38    }
39
40    pub fn transpile(self) -> functions::RemoteFunction {
41        match self {
42            RemoteFunction::Branch {
43                description,
44                input_schema,
45                tasks,
46            } => functions::RemoteFunction::Scalar {
47                description,
48                input_schema:
49                    super::expression::scalar_function_input_schema::transpile(
50                        input_schema,
51                    ),
52                tasks: tasks
53                    .into_iter()
54                    .map(super::BranchTaskExpression::transpile)
55                    .collect(),
56            },
57            RemoteFunction::Leaf {
58                description,
59                input_schema,
60                tasks,
61            } => functions::RemoteFunction::Scalar {
62                description,
63                input_schema:
64                    super::expression::scalar_function_input_schema::transpile(
65                        input_schema,
66                    ),
67                tasks: tasks
68                    .into_iter()
69                    .map(super::LeafTaskExpression::transpile)
70                    .collect(),
71            },
72        }
73    }
74}
75
76#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, arbitrary::Arbitrary)]
77#[serde(tag = "type")]
78#[schemars(rename = "functions.alpha_scalar.InlineFunction")]
79pub enum InlineFunction {
80    #[schemars(title = "Branch")]
81    #[serde(rename = "alpha.scalar.branch.function")]
82    Branch {
83        tasks: Vec<super::BranchTaskExpression>,
84    },
85    #[schemars(title = "Leaf")]
86    #[serde(rename = "alpha.scalar.leaf.function")]
87    Leaf {
88        tasks: Vec<super::LeafTaskExpression>,
89    },
90}
91
92impl InlineFunction {
93    pub fn tasks(&self) -> &[super::BranchTaskExpression] {
94        match self {
95            InlineFunction::Branch { tasks, .. } => tasks,
96            InlineFunction::Leaf { .. } => &[],
97        }
98    }
99
100    pub fn remotes(&self) -> impl Iterator<Item = &crate::RemotePath> {
101        self.tasks().iter().filter_map(|task| match task {
102            super::BranchTaskExpression::ScalarFunction(t) => Some(&t.path),
103            _ => None,
104        })
105    }
106
107    pub fn transpile(self) -> functions::InlineFunction {
108        match self {
109            InlineFunction::Branch { tasks } => {
110                functions::InlineFunction::Scalar {
111                    tasks: tasks
112                        .into_iter()
113                        .map(super::BranchTaskExpression::transpile)
114                        .collect(),
115                }
116            }
117            InlineFunction::Leaf { tasks } => {
118                functions::InlineFunction::Scalar {
119                    tasks: tasks
120                        .into_iter()
121                        .map(super::LeafTaskExpression::transpile)
122                        .collect(),
123                }
124            }
125        }
126    }
127}