Skip to main content

objectiveai_sdk/functions/alpha_scalar/
function.rs

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