vv-agent 0.6.2

VectorVein agent runtime, SDK, CLI, tools, and workspace backends
Documentation
use serde_json::Value;

use crate::types::{CycleRecord, TaskTokenUsage, TokenUsage};

pub fn normalize_token_usage(raw_usage: &Value) -> TokenUsage {
    let Some(raw) = raw_usage.as_object() else {
        return TokenUsage::default();
    };

    let prompt_tokens = read_int(raw.get("prompt_tokens"));
    let completion_tokens = read_int(raw.get("completion_tokens"));
    let input_tokens = read_int(raw.get("input_tokens")).or(prompt_tokens);
    let output_tokens = read_int(raw.get("output_tokens")).or(completion_tokens);
    let total_tokens = read_int(raw.get("total_tokens")).unwrap_or_else(|| {
        prompt_tokens.or(input_tokens).unwrap_or_default()
            + completion_tokens.or(output_tokens).unwrap_or_default()
    });
    let cached_tokens = read_nested_int(
        raw_usage,
        &[
            &["prompt_tokens_details", "cached_tokens"],
            &["input_tokens_details", "cached_tokens"],
            &["cache_read_input_tokens"],
            &["cache_read_tokens"],
        ],
    )
    .unwrap_or_default();
    let reasoning_tokens = read_nested_int(
        raw_usage,
        &[
            &["completion_tokens_details", "reasoning_tokens"],
            &["output_tokens_details", "reasoning_tokens"],
            &["reasoning_tokens"],
        ],
    )
    .unwrap_or_default();
    let cache_creation_tokens = read_nested_int(
        raw_usage,
        &[
            &["input_tokens_details", "cache_creation_tokens"],
            &["prompt_tokens_details", "cache_creation_tokens"],
            &["cache_creation_input_tokens"],
            &["cache_creation_tokens"],
        ],
    )
    .unwrap_or_default();

    TokenUsage {
        prompt_tokens: prompt_tokens.or(input_tokens).unwrap_or_default(),
        completion_tokens: completion_tokens.or(output_tokens).unwrap_or_default(),
        total_tokens,
        cached_tokens,
        reasoning_tokens,
        input_tokens: input_tokens.unwrap_or_default(),
        output_tokens: output_tokens.unwrap_or_default(),
        cache_creation_tokens,
        raw: raw_usage.clone(),
    }
}

pub fn summarize_task_token_usage(cycles: &[CycleRecord]) -> TaskTokenUsage {
    let mut summary = TaskTokenUsage::default();
    for cycle in cycles {
        summary.add_cycle(cycle.index, cycle.token_usage.clone());
    }
    summary
}

fn read_nested_int(source: &Value, path_options: &[&[&str]]) -> Option<u64> {
    for path in path_options {
        let mut current = source;
        let mut matched = true;
        for key in *path {
            let Some(next) = current.as_object().and_then(|object| object.get(*key)) else {
                matched = false;
                break;
            };
            current = next;
        }
        if matched {
            if let Some(value) = read_int(Some(current)) {
                return Some(value);
            }
        }
    }
    None
}

fn read_int(value: Option<&Value>) -> Option<u64> {
    match value? {
        Value::Bool(_) => None,
        Value::Number(number) => number
            .as_u64()
            .or_else(|| {
                number
                    .as_i64()
                    .and_then(|value| (value >= 0).then_some(value as u64))
            })
            .or_else(|| number.as_f64().and_then(float_to_u64)),
        Value::String(value) => value.trim().parse::<u64>().ok(),
        _ => None,
    }
}

fn float_to_u64(value: f64) -> Option<u64> {
    value
        .is_finite()
        .then_some(value.trunc())
        .and_then(|value| {
            if value >= 0.0 {
                Some(value as u64)
            } else {
                None
            }
        })
}