use super::{
common::{Budget, StreamEnd, chat_finish, id},
*,
};
use crate::transform::generate::stream::{
chat,
gemini::{GeminiStreamCollector, GeminiStreamLimits},
};
#[derive(Debug, Clone)]
pub struct GeminiToChatContext {
pub created: i64,
pub model: Option<String>,
}
pub struct GeminiToChatStream {
source: Option<GeminiStreamCollector>,
policy: TargetIdPolicy,
flow: IdentityFlow,
budget: Budget,
calls: super::super::content::Calls,
signatures: super::super::content::Signatures,
tool_indexes: BTreeMap<i64, i64>,
context: GeminiToChatContext,
response_id: Option<String>,
model: Option<String>,
pending: Vec<g::GenerateContentResponseBody>,
emitted: bool,
failed: bool,
report: Report,
}
impl GeminiToChatStream {
pub fn new(
context: GeminiToChatContext,
flow: IdentityFlow,
limits: StreamLimits,
) -> Result<Self, TransformError> {
Self::new_with_policy(
context,
flow,
limits,
TargetIdPolicy::new(crate::Dialect::OpenAiChat),
)
}
pub fn new_with_policy(
context: GeminiToChatContext,
flow: IdentityFlow,
limits: StreamLimits,
policy: TargetIdPolicy,
) -> Result<Self, TransformError> {
if policy.dialect != crate::Dialect::OpenAiChat {
return Err(TransformError::shape(
"stream.policy",
"target policy dialect mismatch",
));
}
if context.created < 0 || context.model.as_ref().is_some_and(|v| v.is_empty()) {
return Err(invalid("invalid factual creation/model context"));
}
if let Some(model) = &context.model {
super::common::bounded(model, limits.max_bytes)?;
}
Ok(Self {
source: Some(GeminiStreamCollector::new(GeminiStreamLimits {
max_bytes: limits.max_bytes,
})),
policy,
flow,
budget: Budget::new(limits),
calls: Default::default(),
signatures: Default::default(),
tool_indexes: BTreeMap::new(),
model: context.model.clone(),
context,
response_id: None,
pending: Vec::new(),
emitted: false,
failed: false,
report: Default::default(),
})
}
pub fn identities(&self) -> &IdentityFlow {
&self.flow
}
pub fn push(
&mut self,
chunk: g::GenerateContentResponseBody,
) -> Result<Converted<Vec<s::ChatCompletionChunk>>, TransformError> {
if self.failed {
return Err(invalid("stream already failed"));
}
let result = self.push_inner(chunk.into_declared());
if result.is_err() {
self.failed = true;
}
result
}
fn push_inner(
&mut self,
chunk: g::GenerateContentResponseBody,
) -> Result<Converted<Vec<s::ChatCompletionChunk>>, TransformError> {
self.budget.input(&chunk)?;
self.source
.as_mut()
.ok_or_else(|| invalid("source consumed"))?
.push(chunk.clone())?;
if let Some(model) = &chunk.model_version {
if self.emitted && self.model.as_ref().is_some_and(|v| v != model) {
return Err(invalid(
"source model contradicts already emitted model fact",
));
}
self.model = Some(model.clone());
}
if let Some(value) = &chunk.response_id {
self.response_id = Some(value.clone());
}
if self.model.is_none() {
self.pending.push(chunk);
return Ok(Converted {
value: Vec::new(),
report: Report::default(),
});
}
let mut output = Vec::new();
for pending in std::mem::take(&mut self.pending) {
self.content(pending, &mut output)?;
}
self.content(chunk, &mut output)?;
self.emitted |= !output.is_empty();
Ok(Converted {
value: output,
report: std::mem::take(&mut self.report),
})
}
fn content(
&mut self,
chunk: g::GenerateContentResponseBody,
out: &mut Vec<s::ChatCompletionChunk>,
) -> Result<(), TransformError> {
let policy = self.policy.clone();
let response = id(
&mut self.flow,
&policy,
IdentityRole::Response,
crate::Dialect::Gemini,
self.response_id.clone(),
0,
)?;
for candidate in chunk.candidates.into_iter().flatten() {
let index = candidate.index.unwrap_or(0);
if let Some(reason) = candidate.finish_reason
&& reason != g::FinishReason::Unspecified
{
chat_finish(
reason,
self.tool_indexes.get(&index).is_some_and(|n| *n > 0),
)?;
}
let Some(mut content) = candidate.content else {
continue;
};
if content.role.is_none() {
content.role = Some("model".into());
}
if content.role.as_deref() != Some("model") {
return Err(invalid("generated candidate must have model role"));
}
for part in content.parts.iter().flatten() {
if part
.function_call
.as_ref()
.is_some_and(|call| call.name.is_empty())
{
return Err(TransformError::missing_metadata("function_call.name"));
}
}
let messages = super::super::content::gemini_content_to_chat(
content,
&mut self.report,
&mut self.flow,
&policy,
&mut self.calls,
Some(&mut self.signatures),
)?;
for message in messages {
let c::ChatMessage::Assistant(message) = message else {
return Err(invalid(
"model candidate produced nonassistant Chat message",
));
};
let mut delta = s::Delta::builder().role(s::DeltaRole::Assistant).build();
let mut has_content = false;
if let Some(details) = message.reasoning_details.clone().flatten() {
delta.reasoning_details = Some(Some(details));
has_content = true;
}
if let Some(text) = c::visible_reasoning(
&message.reasoning_content,
&message.reasoning,
&message.reasoning_details,
) {
delta.reasoning_content = Some(Some(text));
has_content = true;
}
if let Some(content) = message.content.flatten() {
let mut text = String::new();
match content {
c::AssistantContent::Text(v) => text = v,
c::AssistantContent::Parts(parts) => {
for part in parts {
match part {
c::AssistantContentPart::Text(v) => text.push_str(&v.text),
c::AssistantContentPart::Refusal(_) => {
return Err(invalid(
"unexpected refusal in model text mapping",
));
}
}
}
}
}
delta.content = Some(Some(text));
has_content = true;
}
let mut tools = Vec::new();
for tool in message.tool_calls.into_iter().flatten() {
let c::MessageToolCall::Function(tool) = tool else {
return Err(invalid("Gemini function produced nonfunction call"));
};
let next = self.tool_indexes.entry(index).or_default();
let mut call = s::DeltaToolCall::builder(*next).build();
*next = next.checked_add(1).ok_or_else(limit)?;
call.id = Some(Some(tool.id));
call.type_ = Some(Some(s::DeltaToolCallType::Function));
call.function = Some(Some(
s::DeltaFunctionCall::builder()
.name(tool.function.name)
.arguments(tool.function.arguments)
.build(),
));
tools.push(call);
}
if !tools.is_empty() {
delta.tool_calls = Some(Some(tools));
has_content = true;
}
if has_content {
let chunk = self.chunk(
response.clone(),
vec![s::StreamChoice::builder(index, delta).build()],
);
self.emit(chunk, out)?;
}
}
}
Ok(())
}
fn emit(
&mut self,
chunk: s::ChatCompletionChunk,
out: &mut Vec<s::ChatCompletionChunk>,
) -> Result<(), TransformError> {
self.budget.output(std::slice::from_ref(&chunk))?;
out.push(chunk);
Ok(())
}
fn chunk(&self, response: String, choices: Vec<s::StreamChoice>) -> s::ChatCompletionChunk {
s::ChatCompletionChunk::builder(
response,
choices,
self.context.created,
self.model.clone().expect("model checked before emission"),
s::ChunkObject::ChatCompletionChunk,
)
.build()
}
pub fn finish(mut self) -> Result<StreamEnd<s::ChatCompletionChunk>, TransformError> {
if self.failed {
return Err(invalid("stream failed before finish"));
}
let source = self
.source
.take()
.ok_or_else(|| invalid("source consumed"))?
.finish()?
.value;
if self.model.is_none() {
return Err(TransformError::missing_metadata(
"response.model_version or factual context.model",
));
}
let mut chunks = Vec::new();
for pending in std::mem::take(&mut self.pending) {
self.content(pending, &mut chunks)?;
}
let policy = self.policy.clone();
let response = id(
&mut self.flow,
&policy,
IdentityRole::Response,
crate::Dialect::Gemini,
source.response_id,
0,
)?;
let candidates = source
.candidates
.filter(|v| !v.is_empty())
.ok_or_else(|| invalid("prompt blocked without generated candidates"))?;
let mut choices = Vec::new();
for (expected, candidate) in candidates.into_iter().enumerate() {
let index = candidate.index.unwrap_or(0);
if index != expected as i64 {
return Err(invalid(
"final candidate indexes cannot form a complete native Chat choice set",
));
}
let finish = chat_finish(
candidate
.finish_reason
.ok_or_else(|| TransformError::missing_metadata("candidate.finish_reason"))?,
self.tool_indexes.get(&index).is_some_and(|n| *n > 0),
)?;
let logs = candidate
.logprobs_result
.map(|v| super::super::logs::to_chat(v, &mut self.report))
.map(crate::transform::optional)
.transpose()?
.flatten()
.map(chat::logs::synthesize)
.map(crate::transform::optional)
.transpose()?
.flatten();
let mut choice = s::StreamChoice::builder(
index,
s::Delta::builder().role(s::DeltaRole::Assistant).build(),
)
.finish_reason(finish)
.build();
choice.logprobs = logs.map(Some);
choices.push(choice);
for (present, field) in [
(candidate.safety_ratings.is_some(), "safety_ratings"),
(candidate.citation_metadata.is_some(), "citation_metadata"),
(
candidate.grounding_attributions.is_some(),
"grounding_attributions",
),
(candidate.grounding_metadata.is_some(), "grounding_metadata"),
(candidate.avg_logprobs.is_some(), "avg_logprobs"),
(
candidate.url_context_metadata.is_some(),
"url_context_metadata",
),
(candidate.token_count.is_some(), "token_count"),
(candidate.finish_message.is_some(), "finish_message"),
] {
if present {
self.report.omitted(
format!("candidate.{field}"),
"Chat stream has no corresponding declared metadata field",
);
}
}
}
let terminal = self.chunk(response.clone(), choices);
self.emit(terminal, &mut chunks)?;
if let Some(usage) = source.usage_metadata {
let usage = super::super::usage::to_chat(&usage, &mut self.report)?;
let mut tail = self.chunk(response, Vec::new());
tail.usage = Some(Some(chat::usage::synthesize(usage)?));
self.emit(tail, &mut chunks)?;
}
if source.prompt_feedback.is_some() {
self.report.omitted(
"prompt_feedback",
"Chat has no equivalent Gemini prompt feedback",
);
}
if source.model_status.is_some() {
self.report
.omitted("model_status", "Chat has no Gemini model-status field");
}
Ok(StreamEnd {
chunks,
identities: self.flow,
report: self.report,
})
}
}