Skip to main content

gproxy_protocol/transform/claude/generate_content/gemini/
response.rs

1use crate::claude::count_tokens::types::{BetaThinkingBlockType, BetaToolUseBlockType};
2use crate::claude::create_message::response::ClaudeCreateMessageResponse;
3use crate::claude::create_message::types::{
4    BetaContentBlock, BetaMessage, BetaMessageRole, BetaMessageType, BetaServiceTier,
5    BetaStopReason, BetaTextBlock, BetaTextBlockType, Model,
6};
7use crate::claude::types::ClaudeResponseHeaders;
8use crate::gemini::generate_content::response::GeminiGenerateContentResponse;
9use crate::gemini::generate_content::types::{GeminiBlockReason, GeminiFinishReason};
10use crate::transform::claude::generate_content::utils::beta_usage_from_counts;
11use crate::transform::claude::utils::beta_error_response_from_status_message;
12use crate::transform::utils::TransformError;
13
14impl TryFrom<GeminiGenerateContentResponse> for ClaudeCreateMessageResponse {
15    type Error = TransformError;
16
17    fn try_from(value: GeminiGenerateContentResponse) -> Result<Self, TransformError> {
18        Ok(match value {
19            GeminiGenerateContentResponse::Success {
20                stats_code,
21                headers,
22                body,
23            } => {
24                let mut content = Vec::new();
25                let mut has_tool_use = false;
26
27                let candidate = body
28                    .candidates
29                    .clone()
30                    .and_then(|items| items.into_iter().next());
31                if let Some(candidate) = candidate {
32                    if let Some(candidate_content) = candidate.content {
33                        for (idx, part) in candidate_content.parts.into_iter().enumerate() {
34                            if part.thought.unwrap_or(false) {
35                                if let Some(text) = part.text {
36                                    content.push(BetaContentBlock::Thinking(
37                                        crate::claude::create_message::types::BetaThinkingBlock {
38                                            signature: part
39                                                .thought_signature
40                                                .unwrap_or_else(|| format!("thought_{idx}")),
41                                            thinking: text,
42                                            type_: BetaThinkingBlockType::Thinking,
43                                        },
44                                    ));
45                                }
46                            } else if let Some(text) = part.text {
47                                content.push(BetaContentBlock::Text(BetaTextBlock {
48                                    citations: None,
49                                    text,
50                                    type_: BetaTextBlockType::Text,
51                                }));
52                            }
53
54                            if let Some(function_call) = part.function_call {
55                                has_tool_use = true;
56                                content.push(BetaContentBlock::ToolUse(
57                                    crate::claude::create_message::types::BetaToolUseBlock {
58                                        id: function_call
59                                            .id
60                                            .unwrap_or_else(|| format!("tool_call_{idx}")),
61                                        input: function_call.args.unwrap_or_default(),
62                                        name: function_call.name,
63                                        type_: BetaToolUseBlockType::ToolUse,
64                                        cache_control: None,
65                                        caller: None,
66                                    },
67                                ));
68                            }
69
70                            if let Some(function_response) = part.function_response {
71                                let response_text =
72                                    serde_json::to_string(&function_response.response)
73                                        .unwrap_or_default();
74                                if !response_text.is_empty() {
75                                    content.push(BetaContentBlock::Text(BetaTextBlock {
76                                        citations: None,
77                                        text: response_text,
78                                        type_: BetaTextBlockType::Text,
79                                    }));
80                                }
81                            }
82
83                            if let Some(executable_code) = part.executable_code {
84                                content.push(BetaContentBlock::Text(BetaTextBlock {
85                                    citations: None,
86                                    text: executable_code.code,
87                                    type_: BetaTextBlockType::Text,
88                                }));
89                            }
90
91                            if let Some(code_execution_result) = part.code_execution_result
92                                && let Some(output) = code_execution_result.output
93                                && !output.is_empty()
94                            {
95                                content.push(BetaContentBlock::Text(BetaTextBlock {
96                                    citations: None,
97                                    text: output,
98                                    type_: BetaTextBlockType::Text,
99                                }));
100                            }
101
102                            if let Some(file_data) = part.file_data {
103                                content.push(BetaContentBlock::Text(BetaTextBlock {
104                                    citations: None,
105                                    text: file_data.file_uri,
106                                    type_: BetaTextBlockType::Text,
107                                }));
108                            }
109                        }
110                    }
111
112                    if content.is_empty() {
113                        content.push(BetaContentBlock::Text(BetaTextBlock {
114                            citations: None,
115                            text: candidate.finish_message.unwrap_or_default(),
116                            type_: BetaTextBlockType::Text,
117                        }));
118                    }
119
120                    let stop_reason = match candidate.finish_reason {
121                        Some(GeminiFinishReason::MaxTokens) => Some(BetaStopReason::MaxTokens),
122                        Some(GeminiFinishReason::MalformedFunctionCall)
123                        | Some(GeminiFinishReason::UnexpectedToolCall)
124                        | Some(GeminiFinishReason::TooManyToolCalls)
125                        | Some(GeminiFinishReason::MissingThoughtSignature) => {
126                            Some(BetaStopReason::ToolUse)
127                        }
128                        Some(GeminiFinishReason::Safety)
129                        | Some(GeminiFinishReason::Recitation)
130                        | Some(GeminiFinishReason::Blocklist)
131                        | Some(GeminiFinishReason::ProhibitedContent)
132                        | Some(GeminiFinishReason::Spii)
133                        | Some(GeminiFinishReason::ImageSafety)
134                        | Some(GeminiFinishReason::ImageProhibitedContent)
135                        | Some(GeminiFinishReason::ImageRecitation) => {
136                            Some(BetaStopReason::Refusal)
137                        }
138                        _ => {
139                            if has_tool_use {
140                                Some(BetaStopReason::ToolUse)
141                            } else {
142                                Some(BetaStopReason::EndTurn)
143                            }
144                        }
145                    };
146
147                    let usage_metadata = body.usage_metadata.unwrap_or_default();
148                    let prompt_input_tokens = usage_metadata
149                        .prompt_token_count
150                        .unwrap_or(0)
151                        .saturating_add(usage_metadata.tool_use_prompt_token_count.unwrap_or(0));
152                    let cached_tokens = usage_metadata.cached_content_token_count.unwrap_or(0);
153                    let output_tokens = usage_metadata
154                        .candidates_token_count
155                        .unwrap_or(0)
156                        .saturating_add(usage_metadata.thoughts_token_count.unwrap_or(0));
157                    let total_input_tokens = usage_metadata
158                        .total_token_count
159                        .map(|total| total.saturating_sub(output_tokens))
160                        .unwrap_or_else(|| prompt_input_tokens.saturating_add(cached_tokens));
161                    let input_tokens = total_input_tokens.saturating_sub(cached_tokens);
162                    let usage = beta_usage_from_counts(
163                        input_tokens,
164                        cached_tokens,
165                        output_tokens,
166                        BetaServiceTier::Standard,
167                    );
168
169                    ClaudeCreateMessageResponse::Success {
170                        stats_code,
171                        headers: ClaudeResponseHeaders {
172                            extra: headers.extra,
173                        },
174                        body: BetaMessage {
175                            id: body.response_id.unwrap_or_default(),
176                            container: None,
177                            content,
178                            context_management: None,
179                            model: Model::Custom(body.model_version.unwrap_or_default()),
180                            role: BetaMessageRole::Assistant,
181                            stop_reason,
182                            stop_sequence: None,
183                            type_: BetaMessageType::Message,
184                            usage,
185                        },
186                    }
187                } else {
188                    let block_reason = body
189                        .prompt_feedback
190                        .as_ref()
191                        .and_then(|feedback| feedback.block_reason.as_ref());
192                    let stop_reason = match block_reason {
193                        Some(GeminiBlockReason::Safety)
194                        | Some(GeminiBlockReason::Blocklist)
195                        | Some(GeminiBlockReason::ProhibitedContent)
196                        | Some(GeminiBlockReason::ImageSafety) => Some(BetaStopReason::Refusal),
197                        _ => Some(BetaStopReason::EndTurn),
198                    };
199                    let fallback_text = match block_reason {
200                        Some(GeminiBlockReason::Safety) => "blocked_by_safety".to_string(),
201                        Some(GeminiBlockReason::Other) => "blocked".to_string(),
202                        Some(GeminiBlockReason::Blocklist) => "blocked_by_blocklist".to_string(),
203                        Some(GeminiBlockReason::ProhibitedContent) => {
204                            "blocked_by_prohibited_content".to_string()
205                        }
206                        Some(GeminiBlockReason::ImageSafety) => {
207                            "blocked_by_image_safety".to_string()
208                        }
209                        Some(GeminiBlockReason::BlockReasonUnspecified) | None => String::new(),
210                    };
211                    let usage = beta_usage_from_counts(0, 0, 0, BetaServiceTier::Standard);
212                    ClaudeCreateMessageResponse::Success {
213                        stats_code,
214                        headers: ClaudeResponseHeaders {
215                            extra: headers.extra,
216                        },
217                        body: BetaMessage {
218                            id: body.response_id.unwrap_or_default(),
219                            container: None,
220                            content: vec![BetaContentBlock::Text(BetaTextBlock {
221                                citations: None,
222                                text: fallback_text,
223                                type_: BetaTextBlockType::Text,
224                            })],
225                            context_management: None,
226                            model: Model::Custom(body.model_version.unwrap_or_default()),
227                            role: BetaMessageRole::Assistant,
228                            stop_reason,
229                            stop_sequence: None,
230                            type_: BetaMessageType::Message,
231                            usage,
232                        },
233                    }
234                }
235            }
236            GeminiGenerateContentResponse::Error {
237                stats_code,
238                headers,
239                body,
240            } => ClaudeCreateMessageResponse::Error {
241                stats_code,
242                headers: ClaudeResponseHeaders {
243                    extra: headers.extra,
244                },
245                body: beta_error_response_from_status_message(stats_code, body.error.message),
246            },
247        })
248    }
249}