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

1use crate::agent::completions::response::streaming::AgentCompletionIds;
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(untagged)]
15#[schemars(rename = "functions.executions.response.streaming.TaskChunk")]
16pub enum TaskChunk {
17    #[schemars(title = "FunctionExecution")]
18    FunctionExecution(super::FunctionExecutionTaskChunk),
19    #[schemars(title = "VectorCompletion")]
20    VectorCompletion(super::VectorCompletionTaskChunk),
21}
22
23impl AgentCompletionIds for TaskChunk {
24    fn agent_completion_ids(&self) -> impl Iterator<Item = &str> + Send {
25        // Enum dispatch: each variant's own impl returns its own concrete
26        // iterator type. We type-erase via Box<dyn ...> to unify them.
27        let iter: Box<dyn Iterator<Item = &str> + Send + '_> = match self {
28            TaskChunk::FunctionExecution(c) => {
29                Box::new(c.agent_completion_ids())
30            }
31            TaskChunk::VectorCompletion(c) => {
32                Box::new(c.agent_completion_ids())
33            }
34        };
35        iter
36    }
37}
38
39impl TaskChunk {
40    pub fn vector_completion_tasks(
41        &self,
42    ) -> impl Iterator<Item = &super::VectorCompletionTaskChunk> {
43        enum Iter<'a> {
44            FunctionExecution(
45                Box<
46                    dyn Iterator<Item = &'a super::VectorCompletionTaskChunk>
47                        + 'a,
48                >,
49            ),
50            VectorCompletion(Option<&'a super::VectorCompletionTaskChunk>),
51        }
52        impl<'a> Iterator for Iter<'a> {
53            type Item = &'a super::VectorCompletionTaskChunk;
54            fn next(&mut self) -> Option<Self::Item> {
55                match self {
56                    Iter::FunctionExecution(iter) => iter.next(),
57                    Iter::VectorCompletion(opt) => opt.take(),
58                }
59            }
60        }
61        match self {
62            TaskChunk::FunctionExecution(function) => Iter::FunctionExecution(
63                Box::new(function.inner.vector_completion_tasks()),
64            ),
65            TaskChunk::VectorCompletion(vector_completion) => {
66                Iter::VectorCompletion(Some(&vector_completion))
67            }
68        }
69    }
70
71    pub fn index(&self) -> u64 {
72        match self {
73            TaskChunk::FunctionExecution(chunk) => chunk.index,
74            TaskChunk::VectorCompletion(chunk) => chunk.index,
75        }
76    }
77
78    pub fn push(&mut self, other: &TaskChunk) {
79        match (self, other) {
80            (
81                TaskChunk::FunctionExecution(self_chunk),
82                TaskChunk::FunctionExecution(other_chunk),
83            ) => {
84                self_chunk.push(other_chunk);
85            }
86            (
87                TaskChunk::VectorCompletion(self_chunk),
88                TaskChunk::VectorCompletion(other_chunk),
89            ) => {
90                self_chunk.push(other_chunk);
91            }
92            _ => {}
93        }
94    }
95
96}
objectiveai_sdk/functions/executions/response/streaming/task_chunk.rs

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