systemprompt-models 0.47.0

//! Parses a buffered Gemini reply into a [`CanonicalResponse`].
//!
//! # Usage counts
//!
//! Gemini reports `thoughtsTokenCount` beside `candidatesTokenCount`, so a
//! turn that spent its whole budget thinking arrives with `candidates: 0` and
//! looks free; `CanonicalUsage` requires `output_tokens` to include reasoning,
//! so the two are summed and thoughts kept on their own for reporting.
//! `cachedContentTokenCount` is a *subset* of `promptTokenCount`, while
//! `CanonicalUsage::input_tokens` is exclusive of cache reads, so it is
//! subtracted -- left in, the cached slice bills at the input rate as well as
//! the cache-read rate. The wire total already counts it once and is kept.
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

use serde_json::Value;
use uuid::Uuid;

use super::wire::{GeminiCandidate, GeminiPart, GeminiResponse, GeminiUsageMetadata};
use crate::wire::canonical::{
    CanonicalContent, CanonicalResponse, CanonicalStopReason, CanonicalUsage, CodeExecutionOutput,
    GroundedSource, Grounding,
};
use crate::wire::defect::{BodyDefect, buffered_body_defect};
use crate::wire::error::WireParseError;

const GEMINI_GROUNDING_RELEVANCE: f32 = 0.85;

#[must_use]
pub fn stop_reason(finish: &str) -> CanonicalStopReason {
    match finish {
        "STOP" => CanonicalStopReason::EndTurn,
        "MAX_TOKENS" => CanonicalStopReason::MaxTokens,
        _ => CanonicalStopReason::Other,
    }
}

pub fn parse_response(
    value: &Value,
    fallback_model: &str,
) -> Result<CanonicalResponse, WireParseError> {
    let parsed: GeminiResponse =
        serde_json::from_value(value.clone()).map_err(WireParseError::Gemini)?;

    let id = parsed
        .response_id
        .unwrap_or_else(|| format!("msg_{}", Uuid::new_v4().simple()));
    let model = parsed
        .model_version
        .unwrap_or_else(|| fallback_model.to_owned());

    let usage = usage(parsed.usage_metadata);
    let candidate = parsed.candidates.into_iter().next();
    let raw_finish_reason = candidate.as_ref().and_then(|c| c.finish_reason.clone());
    let stop_reason = raw_finish_reason.as_deref().map(stop_reason);
    let grounding = candidate.as_ref().and_then(grounding_from_candidate);
    let parts = candidate.and_then(|c| c.content).map(|c| c.parts);
    let (content, code_execution) = parts.map_or_else(
        || (Vec::new(), None),
        |parts| (parts_to_content(&parts), code_execution(&parts)),
    );

    // Why: Gemini reports finishReason STOP even when the candidate it just
    // returned is a functionCall, so the wire's reason cannot tell "finished
    // talking" from "wants a tool run". Left as EndTurn it renders as
    // `finish_reason: "stop"`, the client ends the turn, and the call rides
    // along in the payload unexecuted. The content is the only reliable
    // signal; MAX_TOKENS still wins, since a call truncated mid-args is not
    // one the client can run.
    let has_tool_use = content
        .iter()
        .any(|c| matches!(c, CanonicalContent::ToolUse { .. }));
    let stop_reason = stop_reason.map(|r| r.with_tool_use(has_tool_use));

    Ok(CanonicalResponse {
        id,
        model,
        content,
        stop_reason,
        usage,
        grounding,
        code_execution,
        raw_finish_reason,
        ..Default::default()
    })
}

// Why: thoughts are summed into output and cached is subtracted from prompt;
// the module head explains both conventions and why billing depends on them.
fn usage(meta: Option<GeminiUsageMetadata>) -> CanonicalUsage {
    meta.map_or_else(CanonicalUsage::default, |u| CanonicalUsage {
        input_tokens: u.prompt.saturating_sub(u.cached),
        output_tokens: u.candidates + u.thoughts,
        cache_read_tokens: u.cached,
        cache_creation_tokens: 0,
        reasoning_tokens: u.thoughts,
        total_tokens: if u.total > 0 {
            u.total
        } else {
            u.prompt + u.candidates + u.thoughts
        },
    })
}

fn grounding_from_candidate(candidate: &GeminiCandidate) -> Option<Grounding> {
    let meta = candidate.grounding_metadata.as_ref()?;
    let sources: Vec<GroundedSource> = meta
        .grounding_chunks
        .iter()
        .filter_map(|c| c.web.as_ref())
        .filter(|w| !w.uri.is_empty())
        .map(|w| GroundedSource {
            uri: w.uri.clone(),
            title: w.title.clone(),
            relevance: Some(GEMINI_GROUNDING_RELEVANCE),
            ..GroundedSource::default()
        })
        .collect();
    if sources.is_empty() && meta.web_search_queries.is_empty() {
        return None;
    }
    Some(Grounding {
        sources,
        queries: meta.web_search_queries.clone(),
    })
}

fn code_execution(parts: &[GeminiPart]) -> Option<CodeExecutionOutput> {
    let mut output = CodeExecutionOutput::default();
    let mut seen = false;
    for part in parts {
        match part {
            GeminiPart::ExecutableCode { executable_code } => {
                seen = true;
                output.language.clone_from(&executable_code.language);
                output.code.clone_from(&executable_code.code);
            },
            GeminiPart::CodeExecutionResult {
                code_execution_result,
            } => {
                seen = true;
                output.result.clone_from(&code_execution_result.output);
                output.outcome.clone_from(&code_execution_result.outcome);
            },
            _ => {},
        }
    }
    seen.then_some(output)
}

pub(super) fn parts_to_content(parts: &[GeminiPart]) -> Vec<CanonicalContent> {
    parts
        .iter()
        .filter_map(|part| match part {
            GeminiPart::Text {
                text,
                thought: Some(true),
                thought_signature,
            } => Some(CanonicalContent::Thinking {
                text: text.clone(),
                signature: thought_signature.clone(),
                id: None,
                encrypted_content: None,
            }),
            GeminiPart::Text { text, .. } if !text.is_empty() => {
                Some(CanonicalContent::Text(text.clone()))
            },
            GeminiPart::FunctionCall {
                function_call,
                thought_signature,
            } => Some(CanonicalContent::ToolUse {
                id: format!("call_{}", Uuid::new_v4().simple()),
                name: function_call.name.clone(),
                input: function_call.args.clone(),
                signature: thought_signature.clone(),
            }),
            _ => None,
        })
        .collect()
}

// Why: runs before `parse_response`, which is total and would turn a body
// carrying nothing into a well-formed empty turn. `None` means the body is
// worth parsing.
#[must_use]
pub fn buffered_defect(value: &Value) -> Option<BodyDefect> {
    buffered_body_defect(value, "candidates", "usageMetadata")
}