Skip to main content

objectiveai_sdk/functions/executions/response/streaming/
task_chunk.rs

1use serde::{Deserialize, Serialize};
2use schemars::JsonSchema;
3
4#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema, arbitrary::Arbitrary)]
5#[serde(untagged)]
6#[schemars(rename = "functions.executions.response.streaming.TaskChunk")]
7pub enum TaskChunk {
8    #[schemars(title = "FunctionExecution")]
9    FunctionExecution(super::FunctionExecutionTaskChunk),
10    #[schemars(title = "VectorCompletion")]
11    VectorCompletion(super::VectorCompletionTaskChunk),
12}
13
14impl TaskChunk {
15    pub fn vector_completion_tasks(
16        &self,
17    ) -> impl Iterator<Item = &super::VectorCompletionTaskChunk> {
18        enum Iter<'a> {
19            FunctionExecution(
20                Box<
21                    dyn Iterator<Item = &'a super::VectorCompletionTaskChunk>
22                        + 'a,
23                >,
24            ),
25            VectorCompletion(Option<&'a super::VectorCompletionTaskChunk>),
26        }
27        impl<'a> Iterator for Iter<'a> {
28            type Item = &'a super::VectorCompletionTaskChunk;
29            fn next(&mut self) -> Option<Self::Item> {
30                match self {
31                    Iter::FunctionExecution(iter) => iter.next(),
32                    Iter::VectorCompletion(opt) => opt.take(),
33                }
34            }
35        }
36        match self {
37            TaskChunk::FunctionExecution(function) => Iter::FunctionExecution(
38                Box::new(function.inner.vector_completion_tasks()),
39            ),
40            TaskChunk::VectorCompletion(vector_completion) => {
41                Iter::VectorCompletion(Some(&vector_completion))
42            }
43        }
44    }
45
46    pub fn index(&self) -> u64 {
47        match self {
48            TaskChunk::FunctionExecution(chunk) => chunk.index,
49            TaskChunk::VectorCompletion(chunk) => chunk.index,
50        }
51    }
52
53    pub fn push(&mut self, other: &TaskChunk) {
54        match (self, other) {
55            (
56                TaskChunk::FunctionExecution(self_chunk),
57                TaskChunk::FunctionExecution(other_chunk),
58            ) => {
59                self_chunk.push(other_chunk);
60            }
61            (
62                TaskChunk::VectorCompletion(self_chunk),
63                TaskChunk::VectorCompletion(other_chunk),
64            ) => {
65                self_chunk.push(other_chunk);
66            }
67            _ => {}
68        }
69    }
70
71    /// Produces log files for this task.
72    ///
73    /// Returns `(reference, files)`.
74    #[cfg(feature = "filesystem")]
75    pub fn produce_files(&self) -> (serde_json::Value, Vec<crate::filesystem::logs::LogFile>) {
76        match self {
77            TaskChunk::FunctionExecution(chunk) => chunk.produce_files(),
78            TaskChunk::VectorCompletion(chunk) => chunk.produce_files(),
79        }
80    }
81}