gproxy_protocol/transform/generate/claude_gemini/stream/

claude_to_gemini.rs

use super::super::ClaudeGeminiUsageFacts;
use super::{
    claude_blocks::Blocks,
    common::{Budget, StreamEnd, StreamLimits, bound, invalid, normalize_gemini, response_id},
    usage::ClaudeUsageProgress,
};
use crate::transform::generate::stream::{
    claude::{ClaudeStreamCollector, ClaudeStreamLimits},
    gemini::{GeminiStreamCollector, GeminiStreamLimits},
};
use crate::{
    transform::{
        Converted, Report, TransformError,
        identity::{IdentityFlow, TargetIdPolicy},
    },
    wire::{DeclaredFields, claude::stream as s, gemini as g},
};

#[derive(Debug, Default, Clone, Copy)]
pub struct ClaudeToGeminiContext {
    pub usage: ClaudeGeminiUsageFacts,
}

pub struct ClaudeToGeminiStream {
    source: Option<ClaudeStreamCollector>,
    target: Option<GeminiStreamCollector>,
    flow: IdentityFlow,
    policy: TargetIdPolicy,
    budget: Budget,
    limits: StreamLimits,
    blocks: Blocks,
    facts: ClaudeGeminiUsageFacts,
    usage: ClaudeUsageProgress,
    response_id: Option<String>,
    model: Option<String>,
    failed: bool,
    stopped: bool,
}

impl ClaudeToGeminiStream {
    pub fn new(
        context: ClaudeToGeminiContext,
        flow: IdentityFlow,
        limits: StreamLimits,
    ) -> Result<Self, TransformError> {
        Self::new_with_policy(
            context,
            flow,
            limits,
            TargetIdPolicy::new(crate::Dialect::Gemini),
        )
    }
    pub fn new_with_policy(
        context: ClaudeToGeminiContext,
        flow: IdentityFlow,
        limits: StreamLimits,
        policy: TargetIdPolicy,
    ) -> Result<Self, TransformError> {
        if policy.dialect != crate::Dialect::Gemini {
            return Err(TransformError::shape(
                "stream.policy",
                "target policy dialect mismatch",
            ));
        }

        Ok(Self {
            source: Some(ClaudeStreamCollector::new(ClaudeStreamLimits {
                max_text_bytes: limits.max_bytes,
                max_json_bytes: limits.max_bytes,
            })),
            target: Some(GeminiStreamCollector::new(GeminiStreamLimits {
                max_bytes: limits.max_bytes,
            })),
            flow,
            policy,
            budget: Budget::new(limits),
            limits,
            blocks: Blocks::new(limits),
            facts: context.usage,
            usage: Default::default(),
            response_id: None,
            model: None,
            failed: false,
            stopped: false,
        })
    }
    pub(crate) fn reserve_external_ids(
        &mut self,
        role: crate::transform::identity::IdentityRole,
        ids: &std::collections::BTreeSet<String>,
    ) -> Result<(), TransformError> {
        self.flow.reserve_external_ids(role, ids).map_err(|e| {
            TransformError::new(
                crate::transform::TransformErrorKind::Conflict,
                "fanout.ids",
                e.to_string(),
            )
        })
    }
    pub fn identities(&self) -> &IdentityFlow {
        &self.flow
    }
    pub fn push(
        &mut self,
        event: s::StreamEvent,
    ) -> Result<Converted<Vec<g::GenerateContentResponseBody>>, TransformError> {
        if self.failed || self.stopped {
            self.failed = true;
            return Err(invalid("stream", "event after message_stop or failure"));
        }
        let result = self.inner(event.into_declared());
        if result.is_err() {
            self.failed = true;
        }
        result
    }
    fn inner(
        &mut self,
        event: s::StreamEvent,
    ) -> Result<Converted<Vec<g::GenerateContentResponseBody>>, TransformError> {
        self.budget.input(&event)?;
        self.source
            .as_mut()
            .ok_or_else(|| invalid("stream", "source consumed"))?
            .push(event.clone())?;
        match event {
            s::StreamEvent::MessageStart(v) => {
                self.usage.start(&v.message.usage)?;
                self.response_id = Some(response_id(
                    &mut self.flow,
                    &self.policy,
                    crate::Dialect::Claude,
                    Some(v.message.id),
                )?);
                self.model = Some(v.message.model);
            }
            s::StreamEvent::ContentBlockStart(v) => {
                self.blocks.start(v, &mut self.flow, &self.policy)?
            }
            s::StreamEvent::ContentBlockDelta(v) => self.blocks.delta(v)?,
            s::StreamEvent::ContentBlockStop(v) => self.blocks.stop(v.index)?,
            s::StreamEvent::MessageDelta(v) => self.usage.delta(&v.usage)?,
            s::StreamEvent::MessageStop(_) => self.stopped = true,
            s::StreamEvent::Ping(_) => {}
            s::StreamEvent::Error(_) => {
                return Err(invalid("stream.error", "Claude upstream error"));
            }
        }
        let mut out = Vec::new();
        while let Some(part) = self.blocks.next()? {
            let body = g::GenerateContentResponseBody::builder()
                .candidates(vec![
                    g::Candidate::builder()
                        .index(0)
                        .content(
                            g::Content::builder()
                                .role("model")
                                .parts(vec![part])
                                .build(),
                        )
                        .build(),
                ])
                .build();
            self.emit(&mut out, body)?;
        }
        self.blocks.check_pending()?;
        Ok(Converted {
            value: out,
            report: Report::default(),
        })
    }
    fn emit(
        &mut self,
        out: &mut Vec<g::GenerateContentResponseBody>,
        mut body: g::GenerateContentResponseBody,
    ) -> Result<(), TransformError> {
        body.response_id = self.response_id.clone();
        body.model_version = self.model.clone();
        self.budget.output(&body)?;
        self.target
            .as_mut()
            .ok_or_else(|| invalid("stream", "target consumed"))?
            .push(body.clone())?;
        out.push(body);
        Ok(())
    }
    pub fn finish(mut self) -> Result<StreamEnd<g::GenerateContentResponseBody>, TransformError> {
        if self.failed || !self.stopped || !self.blocks.is_empty() {
            return Err(invalid(
                "stream",
                "failed stream, premature EOF or pending blocks",
            ));
        }
        let mut source = self
            .source
            .take()
            .ok_or_else(|| invalid("stream", "source consumed"))?
            .finish()?
            .value;
        self.usage.finalize(&mut source.usage, self.facts)?;
        let converted = super::super::claude_to_gemini_response(
            source,
            self.facts,
            &mut self.flow.clone(),
            &self.policy,
        )?;
        let mut expected = converted.value;
        bound(&expected, self.limits.max_bytes)?;
        let mut tail = expected.clone();
        for candidate in tail.candidates.iter_mut().flatten() {
            if let Some(content) = &mut candidate.content {
                content.parts = Some(Vec::new());
            }
        }
        let mut out = Vec::new();
        self.emit(&mut out, tail)?;
        let mut actual = self
            .target
            .take()
            .ok_or_else(|| invalid("stream", "target consumed"))?
            .finish()?
            .value;
        normalize_gemini(&mut actual);
        normalize_gemini(&mut expected);
        if actual != expected {
            return Err(invalid(
                "stream.parity",
                "native Gemini output differs from completed Claude conversion",
            ));
        }
        Ok(StreamEnd {
            chunks: out,
            identities: self.flow,
            report: converted.report,
        })
    }
}