objectiveai_sdk/functions/executions/response/streaming/
task_chunk.rs1use 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 #[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}