1use crate::agent;
29use serde::{Deserialize, Serialize};
30use schemars::JsonSchema;
31
32#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, arbitrary::Arbitrary)]
38#[serde(tag = "type")]
39#[schemars(rename = "functions.TaskExpression")]
40pub enum TaskExpression {
41 #[schemars(title = "ScalarFunction")]
42 #[serde(rename = "scalar.function")]
43 ScalarFunction(ScalarFunctionTaskExpression),
44 #[schemars(title = "VectorFunction")]
45 #[serde(rename = "vector.function")]
46 VectorFunction(VectorFunctionTaskExpression),
47 #[schemars(title = "VectorCompletion")]
48 #[serde(rename = "vector.completion")]
49 VectorCompletion(VectorCompletionTaskExpression),
50 #[schemars(title = "PlaceholderScalarFunction")]
51 #[serde(rename = "placeholder.scalar.function")]
52 PlaceholderScalarFunction(PlaceholderScalarFunctionTaskExpression),
53 #[schemars(title = "PlaceholderVectorFunction")]
54 #[serde(rename = "placeholder.vector.function")]
55 PlaceholderVectorFunction(PlaceholderVectorFunctionTaskExpression),
56}
57
58impl TaskExpression {
59 pub fn url(&self) -> Option<String> {
60 match self {
61 TaskExpression::ScalarFunction(task) => Some(task.url()),
62 TaskExpression::VectorFunction(task) => Some(task.url()),
63 TaskExpression::VectorCompletion(_) => None,
64 TaskExpression::PlaceholderScalarFunction(_) => None,
65 TaskExpression::PlaceholderVectorFunction(_) => None,
66 }
67 }
68
69 pub fn take_skip(&mut self) -> Option<super::expression::Expression> {
71 match self {
72 TaskExpression::ScalarFunction(task) => task.skip.take(),
73 TaskExpression::VectorFunction(task) => task.skip.take(),
74 TaskExpression::VectorCompletion(task) => task.skip.take(),
75 TaskExpression::PlaceholderScalarFunction(task) => task.skip.take(),
76 TaskExpression::PlaceholderVectorFunction(task) => task.skip.take(),
77 }
78 }
79
80 pub fn map(&self) -> Option<&super::expression::Expression> {
82 match self {
83 TaskExpression::ScalarFunction(task) => task.map.as_ref(),
84 TaskExpression::VectorFunction(task) => task.map.as_ref(),
85 TaskExpression::VectorCompletion(task) => task.map.as_ref(),
86 TaskExpression::PlaceholderScalarFunction(task) => {
87 task.map.as_ref()
88 }
89 TaskExpression::PlaceholderVectorFunction(task) => {
90 task.map.as_ref()
91 }
92 }
93 }
94
95 pub fn compile(
97 self,
98 params: &super::expression::Params,
99 ) -> Result<Task, super::expression::ExpressionError> {
100 match self {
101 TaskExpression::ScalarFunction(task) => {
102 task.compile(params).map(Task::ScalarFunction)
103 }
104 TaskExpression::VectorFunction(task) => {
105 task.compile(params).map(Task::VectorFunction)
106 }
107 TaskExpression::VectorCompletion(task) => {
108 task.compile(params).map(Task::VectorCompletion)
109 }
110 TaskExpression::PlaceholderScalarFunction(task) => {
111 task.compile(params).map(Task::PlaceholderScalarFunction)
112 }
113 TaskExpression::PlaceholderVectorFunction(task) => {
114 task.compile(params).map(Task::PlaceholderVectorFunction)
115 }
116 }
117 }
118}
119
120#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
125#[serde(tag = "type")]
126#[schemars(rename = "functions.Task")]
127pub enum Task {
128 #[schemars(title = "ScalarFunction")]
130 #[serde(rename = "scalar.function")]
131 ScalarFunction(ScalarFunctionTask),
132 #[schemars(title = "VectorFunction")]
134 #[serde(rename = "vector.function")]
135 VectorFunction(VectorFunctionTask),
136 #[schemars(title = "VectorCompletion")]
138 #[serde(rename = "vector.completion")]
139 VectorCompletion(VectorCompletionTask),
140 #[schemars(title = "PlaceholderScalarFunction")]
142 #[serde(rename = "placeholder.scalar.function")]
143 PlaceholderScalarFunction(PlaceholderScalarFunctionTask),
144 #[schemars(title = "PlaceholderVectorFunction")]
146 #[serde(rename = "placeholder.vector.function")]
147 PlaceholderVectorFunction(PlaceholderVectorFunctionTask),
148}
149
150impl Task {
151 pub fn compile_output(
152 &self,
153 input: &super::expression::InputValue,
154 raw_output: super::expression::TaskOutput,
155 ) -> Result<
156 super::expression::TaskOutputOwned,
157 super::expression::ExpressionError,
158 > {
159 match self {
160 Task::ScalarFunction(task) => {
161 task.compile_output(input, raw_output)
162 }
163 Task::VectorFunction(task) => {
164 task.compile_output(input, raw_output)
165 }
166 Task::VectorCompletion(task) => {
167 task.compile_output(input, raw_output)
168 }
169 Task::PlaceholderScalarFunction(task) => {
170 task.compile_output(input, raw_output)
171 }
172 Task::PlaceholderVectorFunction(task) => {
173 task.compile_output(input, raw_output)
174 }
175 }
176 }
177}
178
179#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, arbitrary::Arbitrary)]
181#[schemars(rename = "functions.ScalarFunctionTaskExpression")]
182pub struct ScalarFunctionTaskExpression {
183 #[serde(flatten)]
184 #[schemars(schema_with = "crate::flatten_schema::<crate::RemotePath>")]
185 pub path: crate::RemotePath,
186
187 #[serde(skip_serializing_if = "Option::is_none")]
189 #[schemars(extend("omitempty" = true))]
190 pub skip: Option<super::expression::Expression>,
191
192 #[serde(skip_serializing_if = "Option::is_none")]
195 #[schemars(extend("omitempty" = true))]
196 pub map: Option<super::expression::Expression>,
197
198 pub input:
201 super::expression::WithExpression<super::expression::InputValueExpression>,
202
203 pub output: super::expression::Expression,
219}
220
221impl ScalarFunctionTaskExpression {
222 pub fn url(&self) -> String {
223 self.path.url()
224 }
225
226 pub fn compile(
228 self,
229 params: &super::expression::Params,
230 ) -> Result<ScalarFunctionTask, super::expression::ExpressionError> {
231 let input = self.input.compile_one(params)?.compile(params)?;
232 Ok(ScalarFunctionTask {
233 path: self.path,
234 input,
235 output: self.output,
236 })
237 }
238}
239
240#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
242#[schemars(rename = "functions.ScalarFunctionTask")]
243pub struct ScalarFunctionTask {
244 #[serde(flatten)]
245 #[schemars(schema_with = "crate::flatten_schema::<crate::RemotePath>")]
246 pub path: crate::RemotePath,
247 pub input: super::expression::InputValue,
249 pub output: super::expression::Expression,
255}
256
257impl ScalarFunctionTask {
258 pub fn url(&self) -> String {
259 self.path.url()
260 }
261
262 pub fn compile_output(
263 &self,
264 input: &super::expression::InputValue,
265 raw_output: super::expression::TaskOutput,
266 ) -> Result<
267 super::expression::TaskOutputOwned,
268 super::expression::ExpressionError,
269 > {
270 let params =
271 super::expression::Params::Ref(super::expression::ParamsRef {
272 input,
273 output: Some(raw_output),
274 map: None,
275 tasks_min: None,
276 tasks_max: None,
277 depth: None,
278 name: None,
279 spec: None,
280 });
281 let compiled_output = self.output.compile_one(¶ms)?;
282 Ok(compiled_output)
283 }
284}
285
286#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, arbitrary::Arbitrary)]
288#[schemars(rename = "functions.VectorFunctionTaskExpression")]
289pub struct VectorFunctionTaskExpression {
290 #[serde(flatten)]
291 #[schemars(schema_with = "crate::flatten_schema::<crate::RemotePath>")]
292 pub path: crate::RemotePath,
293
294 #[serde(skip_serializing_if = "Option::is_none")]
296 #[schemars(extend("omitempty" = true))]
297 pub skip: Option<super::expression::Expression>,
298
299 #[serde(skip_serializing_if = "Option::is_none")]
302 #[schemars(extend("omitempty" = true))]
303 pub map: Option<super::expression::Expression>,
304
305 pub input:
308 super::expression::WithExpression<super::expression::InputValueExpression>,
309
310 pub output: super::expression::Expression,
326}
327
328impl VectorFunctionTaskExpression {
329 pub fn url(&self) -> String {
330 self.path.url()
331 }
332
333 pub fn compile(
335 self,
336 params: &super::expression::Params,
337 ) -> Result<VectorFunctionTask, super::expression::ExpressionError> {
338 let input = self.input.compile_one(params)?.compile(params)?;
339 Ok(VectorFunctionTask {
340 path: self.path,
341 input,
342 output: self.output,
343 })
344 }
345}
346
347#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
349#[schemars(rename = "functions.VectorFunctionTask")]
350pub struct VectorFunctionTask {
351 #[serde(flatten)]
352 #[schemars(schema_with = "crate::flatten_schema::<crate::RemotePath>")]
353 pub path: crate::RemotePath,
354 pub input: super::expression::InputValue,
356 pub output: super::expression::Expression,
362}
363
364impl VectorFunctionTask {
365 pub fn url(&self) -> String {
366 self.path.url()
367 }
368
369 pub fn compile_output(
370 &self,
371 input: &super::expression::InputValue,
372 raw_output: super::expression::TaskOutput,
373 ) -> Result<
374 super::expression::TaskOutputOwned,
375 super::expression::ExpressionError,
376 > {
377 let params =
378 super::expression::Params::Ref(super::expression::ParamsRef {
379 input,
380 output: Some(raw_output),
381 map: None,
382 tasks_min: None,
383 tasks_max: None,
384 depth: None,
385 name: None,
386 spec: None,
387 });
388 let compiled_output = self.output.compile_one(¶ms)?;
389 Ok(compiled_output)
390 }
391}
392
393#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, arbitrary::Arbitrary)]
395#[schemars(rename = "functions.VectorCompletionTaskExpression")]
396pub struct VectorCompletionTaskExpression {
397 #[serde(skip_serializing_if = "Option::is_none")]
399 #[schemars(extend("omitempty" = true))]
400 pub skip: Option<super::expression::Expression>,
401
402 #[serde(skip_serializing_if = "Option::is_none")]
405 #[schemars(extend("omitempty" = true))]
406 pub map: Option<super::expression::Expression>,
407
408 pub messages: super::expression::WithExpression<
411 Vec<
412 super::expression::WithExpression<
413 agent::completions::message::MessageExpression,
414 >,
415 >,
416 >,
417 pub responses: super::expression::WithExpression<
420 Vec<
421 super::expression::WithExpression<
422 agent::completions::message::RichContentExpression,
423 >,
424 >,
425 >,
426
427 pub output: super::expression::Expression,
439}
440
441impl VectorCompletionTaskExpression {
442 pub fn compile(
444 self,
445 params: &super::expression::Params,
446 ) -> Result<VectorCompletionTask, super::expression::ExpressionError> {
447 let messages = self.messages.compile_one(params)?;
449 let mut compiled_messages = Vec::with_capacity(messages.len());
450 for message in messages {
451 match message.compile_one_or_many(params)? {
452 super::expression::OneOrMany::One(one_message) => {
453 compiled_messages.push(one_message.compile(params)?);
454 }
455 super::expression::OneOrMany::Many(many_messages) => {
456 for message in many_messages {
457 compiled_messages.push(message.compile(params)?);
458 }
459 }
460 }
461 }
462
463 let responses = self.responses.compile_one(params)?;
465 let mut compiled_responses = Vec::with_capacity(responses.len());
466 for response in responses {
467 match response.compile_one_or_many(params)? {
468 super::expression::OneOrMany::One(one_response) => {
469 compiled_responses.push(one_response.compile(params)?);
470 }
471 super::expression::OneOrMany::Many(many_responses) => {
472 for response in many_responses {
473 compiled_responses.push(response.compile(params)?);
474 }
475 }
476 }
477 }
478
479 Ok(VectorCompletionTask {
480 messages: compiled_messages,
481 responses: compiled_responses,
482 output: self.output,
483 })
484 }
485}
486
487#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
489#[schemars(rename = "functions.VectorCompletionTask")]
490pub struct VectorCompletionTask {
491 pub messages: Vec<agent::completions::message::Message>,
493 pub responses: Vec<agent::completions::message::RichContent>,
495 pub output: super::expression::Expression,
501}
502
503impl VectorCompletionTask {
504 pub fn compile_output(
505 &self,
506 input: &super::expression::InputValue,
507 raw_output: super::expression::TaskOutput,
508 ) -> Result<
509 super::expression::TaskOutputOwned,
510 super::expression::ExpressionError,
511 > {
512 let params =
513 super::expression::Params::Ref(super::expression::ParamsRef {
514 input,
515 output: Some(raw_output),
516 map: None,
517 tasks_min: None,
518 tasks_max: None,
519 depth: None,
520 name: None,
521 spec: None,
522 });
523 let compiled_output = self.output.compile_one(¶ms)?;
524 Ok(compiled_output)
525 }
526}
527
528#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, arbitrary::Arbitrary)]
533#[schemars(rename = "functions.PlaceholderScalarFunctionTaskExpression")]
534pub struct PlaceholderScalarFunctionTaskExpression {
535 pub input_schema: super::expression::InputSchema,
537
538 #[serde(skip_serializing_if = "Option::is_none")]
540 #[schemars(extend("omitempty" = true))]
541 pub skip: Option<super::expression::Expression>,
542
543 #[serde(skip_serializing_if = "Option::is_none")]
546 #[schemars(extend("omitempty" = true))]
547 pub map: Option<super::expression::Expression>,
548
549 pub input:
552 super::expression::WithExpression<super::expression::InputValueExpression>,
553
554 pub output: super::expression::Expression,
557}
558
559impl PlaceholderScalarFunctionTaskExpression {
560 pub fn compile(
561 self,
562 params: &super::expression::Params,
563 ) -> Result<PlaceholderScalarFunctionTask, super::expression::ExpressionError>
564 {
565 let input = self.input.compile_one(params)?.compile(params)?;
566 Ok(PlaceholderScalarFunctionTask {
567 input_schema: self.input_schema,
568 input,
569 output: self.output,
570 })
571 }
572}
573
574#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
579#[schemars(rename = "functions.PlaceholderScalarFunctionTask")]
580pub struct PlaceholderScalarFunctionTask {
581 pub input_schema: super::expression::InputSchema,
583 pub input: super::expression::InputValue,
585 pub output: super::expression::Expression,
587}
588
589impl PlaceholderScalarFunctionTask {
590 pub fn compile_output(
591 &self,
592 input: &super::expression::InputValue,
593 raw_output: super::expression::TaskOutput,
594 ) -> Result<
595 super::expression::TaskOutputOwned,
596 super::expression::ExpressionError,
597 > {
598 let params =
599 super::expression::Params::Ref(super::expression::ParamsRef {
600 input,
601 output: Some(raw_output),
602 map: None,
603 tasks_min: None,
604 tasks_max: None,
605 depth: None,
606 name: None,
607 spec: None,
608 });
609 let compiled_output = self.output.compile_one(¶ms)?;
610 Ok(compiled_output)
611 }
612}
613
614#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, arbitrary::Arbitrary)]
619#[schemars(rename = "functions.PlaceholderVectorFunctionTaskExpression")]
620pub struct PlaceholderVectorFunctionTaskExpression {
621 pub input_schema: super::expression::InputSchema,
623
624 pub output_length: super::expression::Expression,
627
628 pub input_split: super::expression::Expression,
631
632 pub input_merge: super::expression::Expression,
635
636 #[serde(skip_serializing_if = "Option::is_none")]
638 #[schemars(extend("omitempty" = true))]
639 pub skip: Option<super::expression::Expression>,
640
641 #[serde(skip_serializing_if = "Option::is_none")]
644 #[schemars(extend("omitempty" = true))]
645 pub map: Option<super::expression::Expression>,
646
647 pub input:
650 super::expression::WithExpression<super::expression::InputValueExpression>,
651
652 pub output: super::expression::Expression,
655}
656
657impl PlaceholderVectorFunctionTaskExpression {
658 pub fn compile(
659 self,
660 params: &super::expression::Params,
661 ) -> Result<PlaceholderVectorFunctionTask, super::expression::ExpressionError>
662 {
663 let input = self.input.compile_one(params)?.compile(params)?;
664 Ok(PlaceholderVectorFunctionTask {
665 input_schema: self.input_schema,
666 output_length: self.output_length,
667 input_split: self.input_split,
668 input_merge: self.input_merge,
669 input,
670 output: self.output,
671 })
672 }
673}
674
675#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
680#[schemars(rename = "functions.PlaceholderVectorFunctionTask")]
681pub struct PlaceholderVectorFunctionTask {
682 pub input_schema: super::expression::InputSchema,
684 pub output_length: super::expression::Expression,
686 pub input_split: super::expression::Expression,
688 pub input_merge: super::expression::Expression,
690 pub input: super::expression::InputValue,
692 pub output: super::expression::Expression,
694}
695
696impl PlaceholderVectorFunctionTask {
697 pub fn compile_output(
698 &self,
699 input: &super::expression::InputValue,
700 raw_output: super::expression::TaskOutput,
701 ) -> Result<
702 super::expression::TaskOutputOwned,
703 super::expression::ExpressionError,
704 > {
705 let params =
706 super::expression::Params::Ref(super::expression::ParamsRef {
707 input,
708 output: Some(raw_output),
709 map: None,
710 tasks_min: None,
711 tasks_max: None,
712 depth: None,
713 name: None,
714 spec: None,
715 });
716 let compiled_output = self.output.compile_one(¶ms)?;
717 Ok(compiled_output)
718 }
719}
720
721#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
727#[serde(untagged)]
728#[schemars(rename = "functions.CompiledTask")]
729pub enum CompiledTask {
730 #[schemars(title = "One")]
732 One(Task),
733 #[schemars(title = "Many")]
735 Many(Vec<Task>),
736}