Skip to main content

gproxy_protocol/adapt/generate/stream/fanout/
event.rs

1use super::*;
2use crate::adapt::generate::fanout::output::Client;
3use crate::wire::{
4    gemini as g,
5    openai::chat::{self as h, stream as hs},
6};
7
8/// Sealed by NativeEvent. Rewrites only the aggregate envelope of concrete client events.
9pub trait FanoutEvent: NativeEvent {
10    fn project(
11        &mut self,
12        index: usize,
13        id: &str,
14        created: &mut Option<i64>,
15    ) -> Result<(), TransformError>;
16    fn aggregate(
17        values: Vec<Self::Full>,
18        id: String,
19        report: &mut Report,
20    ) -> Result<Self::Full, TransformError>;
21    fn tail(value: &Self::Full, usage: bool) -> Result<Option<Self>, TransformError>;
22    fn equivalent(a: &Self::Full, b: &Self::Full) -> bool;
23    fn visible(&self) -> bool;
24}
25
26impl FanoutEvent for hs::ChatCompletionChunk {
27    fn project(
28        &mut self,
29        index: usize,
30        id: &str,
31        created: &mut Option<i64>,
32    ) -> Result<(), TransformError> {
33        self.id = id.into();
34        self.created = *created.get_or_insert(self.created);
35        self.usage = None;
36        self.service_tier = None;
37        self.system_fingerprint = None;
38        if self.choices.len() > 1 || self.choices.iter().any(|c| c.index != 0) {
39            return Err(invalid("child Chat choice is not singleton"));
40        }
41        for c in &mut self.choices {
42            c.index = i64::try_from(index).map_err(|_| limit("choice index overflow"))?;
43        }
44        Ok(())
45    }
46    fn aggregate(
47        values: Vec<Self::Full>,
48        id: String,
49        report: &mut Report,
50    ) -> Result<Self::Full, TransformError> {
51        let mut value = <h::GenerateContentResponseBody as Client>::aggregate(values, id, report)?;
52        if value.service_tier.is_some()
53            || value.system_fingerprint.is_some()
54            || value.moderation.is_some()
55        {
56            report.omitted(
57                "service_tier,system_fingerprint,moderation",
58                "aggregate metadata cannot be established before incremental exposure",
59            );
60        }
61        value.service_tier = None;
62        value.system_fingerprint = None;
63        value.moderation = None;
64        Ok(value)
65    }
66    fn tail(value: &Self::Full, usage: bool) -> Result<Option<Self>, TransformError> {
67        if !usage || value.usage.is_none() {
68            return Ok(None);
69        }
70        let mut v = Self::builder(
71            value.id.clone(),
72            Vec::new(),
73            value.created,
74            value.model.clone(),
75            hs::ChunkObject::ChatCompletionChunk,
76        )
77        .build();
78        v.usage = Some(
79            value
80                .usage
81                .clone()
82                .map(crate::transform::generate::stream::chat::usage::synthesize)
83                .transpose()?,
84        );
85        Ok(Some(v))
86    }
87    fn equivalent(a: &Self::Full, b: &Self::Full) -> bool {
88        a == b
89    }
90    fn visible(&self) -> bool {
91        !self.choices.is_empty()
92    }
93}
94
95impl FanoutEvent for g::GenerateContentResponseBody {
96    fn project(
97        &mut self,
98        index: usize,
99        id: &str,
100        _: &mut Option<i64>,
101    ) -> Result<(), TransformError> {
102        self.response_id = Some(id.into());
103        self.usage_metadata = None;
104        self.model_version = None;
105        self.model_status = None;
106        self.prompt_feedback = None;
107        if self
108            .candidates
109            .as_ref()
110            .is_some_and(|v| v.len() > 1 || v.iter().any(|c| c.index.is_some_and(|i| i != 0)))
111        {
112            return Err(invalid("child Gemini candidate is not singleton"));
113        }
114        for c in self.candidates.iter_mut().flatten() {
115            c.index = Some(i64::try_from(index).map_err(|_| limit("candidate index overflow"))?);
116        }
117        Ok(())
118    }
119    fn aggregate(
120        values: Vec<Self::Full>,
121        id: String,
122        report: &mut Report,
123    ) -> Result<Self::Full, TransformError> {
124        let mut value = <Self as Client>::aggregate(values, id, report)?;
125        if value.model_version.is_some()
126            || value.model_status.is_some()
127            || value.prompt_feedback.is_some()
128        {
129            report.omitted(
130                "modelVersion,modelStatus,promptFeedback",
131                "aggregate metadata cannot be established before incremental exposure",
132            );
133        }
134        value.model_version = None;
135        value.model_status = None;
136        value.prompt_feedback = None;
137        Ok(value)
138    }
139    fn tail(value: &Self::Full, _: bool) -> Result<Option<Self>, TransformError> {
140        let mut v = Self::builder().build();
141        v.response_id = value.response_id.clone();
142        v.usage_metadata = value.usage_metadata.clone();
143        Ok(Some(v))
144    }
145    fn equivalent(a: &Self::Full, b: &Self::Full) -> bool {
146        a == b
147    }
148    fn visible(&self) -> bool {
149        self.candidates.as_ref().is_some_and(|v| !v.is_empty())
150    }
151}