Skip to main content

systemprompt_models/wire/gemini/
response.rs

1//! Parses a buffered Gemini reply into a [`CanonicalResponse`].
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use serde_json::Value;
7use uuid::Uuid;
8
9use super::wire::{GeminiCandidate, GeminiPart, GeminiResponse, GeminiUsageMetadata};
10use crate::wire::canonical::{
11    CanonicalContent, CanonicalResponse, CanonicalStopReason, CanonicalUsage, CodeExecutionOutput,
12    GroundedSource, Grounding,
13};
14
15const GEMINI_GROUNDING_RELEVANCE: f32 = 0.85;
16
17#[must_use]
18pub fn stop_reason(finish: &str) -> CanonicalStopReason {
19    match finish {
20        "STOP" => CanonicalStopReason::EndTurn,
21        "MAX_TOKENS" => CanonicalStopReason::MaxTokens,
22        _ => CanonicalStopReason::Other,
23    }
24}
25
26#[must_use]
27pub fn parse_response(value: &Value, fallback_model: &str) -> CanonicalResponse {
28    let parsed: GeminiResponse = serde_json::from_value(value.clone()).unwrap_or(GeminiResponse {
29        candidates: Vec::new(),
30        usage_metadata: None,
31        response_id: None,
32        model_version: None,
33    });
34
35    let id = parsed
36        .response_id
37        .unwrap_or_else(|| format!("msg_{}", Uuid::new_v4().simple()));
38    let model = parsed
39        .model_version
40        .unwrap_or_else(|| fallback_model.to_owned());
41
42    let usage = usage(parsed.usage_metadata);
43    let candidate = parsed.candidates.into_iter().next();
44    let raw_finish_reason = candidate.as_ref().and_then(|c| c.finish_reason.clone());
45    let stop_reason = raw_finish_reason.as_deref().map(stop_reason);
46    let grounding = candidate.as_ref().and_then(grounding_from_candidate);
47    let parts = candidate.and_then(|c| c.content).map(|c| c.parts);
48    let (content, code_execution) = parts.map_or_else(
49        || (Vec::new(), None),
50        |parts| (parts_to_content(&parts), code_execution(&parts)),
51    );
52
53    CanonicalResponse {
54        id,
55        model,
56        content,
57        stop_reason,
58        usage,
59        grounding,
60        code_execution,
61        raw_finish_reason,
62    }
63}
64
65fn usage(meta: Option<GeminiUsageMetadata>) -> CanonicalUsage {
66    meta.map_or_else(CanonicalUsage::default, |u| CanonicalUsage {
67        input_tokens: u.prompt,
68        output_tokens: u.candidates,
69        cache_read_tokens: u.cached,
70        cache_creation_tokens: 0,
71        total_tokens: if u.total > 0 {
72            u.total
73        } else {
74            u.prompt + u.candidates
75        },
76    })
77}
78
79fn grounding_from_candidate(candidate: &GeminiCandidate) -> Option<Grounding> {
80    let meta = candidate.grounding_metadata.as_ref()?;
81    let sources: Vec<GroundedSource> = meta
82        .grounding_chunks
83        .iter()
84        .filter_map(|c| c.web.as_ref())
85        .filter(|w| !w.uri.is_empty())
86        .map(|w| GroundedSource {
87            uri: w.uri.clone(),
88            title: w.title.clone(),
89            relevance: Some(GEMINI_GROUNDING_RELEVANCE),
90            ..GroundedSource::default()
91        })
92        .collect();
93    if sources.is_empty() && meta.web_search_queries.is_empty() {
94        return None;
95    }
96    Some(Grounding {
97        sources,
98        queries: meta.web_search_queries.clone(),
99    })
100}
101
102fn code_execution(parts: &[GeminiPart]) -> Option<CodeExecutionOutput> {
103    let mut output = CodeExecutionOutput::default();
104    let mut seen = false;
105    for part in parts {
106        match part {
107            GeminiPart::ExecutableCode { executable_code } => {
108                seen = true;
109                output.language.clone_from(&executable_code.language);
110                output.code.clone_from(&executable_code.code);
111            },
112            GeminiPart::CodeExecutionResult {
113                code_execution_result,
114            } => {
115                seen = true;
116                output.result.clone_from(&code_execution_result.output);
117                output.outcome.clone_from(&code_execution_result.outcome);
118            },
119            _ => {},
120        }
121    }
122    seen.then_some(output)
123}
124
125pub(super) fn parts_to_content(parts: &[GeminiPart]) -> Vec<CanonicalContent> {
126    parts
127        .iter()
128        .filter_map(|part| match part {
129            GeminiPart::Text { text } if !text.is_empty() => {
130                Some(CanonicalContent::Text(text.clone()))
131            },
132            GeminiPart::FunctionCall {
133                function_call,
134                thought_signature,
135            } => Some(CanonicalContent::ToolUse {
136                id: format!("call_{}", Uuid::new_v4().simple()),
137                name: function_call.name.clone(),
138                input: function_call.args.clone(),
139                signature: thought_signature.clone(),
140            }),
141            _ => None,
142        })
143        .collect()
144}