Skip to content

fix(providers): normalize OpenAI-compatible cache usage #536

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Raymond8196 wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

fix(providers): normalize OpenAI-compatible cache usage #536

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 23 additions & 13 deletions src-tauri/crates/agent-core/src/core/providers/openai_compat/client.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,14 +6,13 @@

use reqwest::Client;
use serde_json::Value;
use std::collections::HashMap;
use tracing::{debug, warn};

use super::types::{ApiErrorResponse, StreamChunk, ToolCallResponse};
use crate::providers::openai_policy::{resolve_openai_chat_wire_policy, OpenAiChatWirePolicy};
use crate::providers::registry::{provider_id, ProviderSpec};
use crate::providers::traits::{
finish_reason as finish, usage_key, LLMResponse, ProviderConfig, ProviderError, ToolCallRequest,
finish_reason as finish, LLMResponse, ProviderConfig, ProviderError, ToolCallRequest,
};
use crate::utils::build_http_client;

Expand Down Expand Up @@ -211,7 +210,7 @@ impl OpenAICompatClient {
/// into a single `LLMResponse`.
pub(super) fn reassemble_sse_to_response(body: &str) -> Result<LLMResponse, ProviderError> {
let mut content = String::new();
let mut usage: HashMap<String, i64> = HashMap::new();
let mut usage = std::collections::HashMap::new();

for line in body.lines() {
let line = line.trim();
Expand Down Expand Up @@ -240,15 +239,7 @@ impl OpenAICompatClient {
}
}
if let Some(ref api_usage) = chunk.usage {
usage.insert(
usage_key::PROMPT_TOKENS.to_string(),
api_usage.prompt_tokens,
);
usage.insert(
usage_key::COMPLETION_TOKENS.to_string(),
api_usage.completion_tokens,
);
usage.insert(usage_key::TOTAL_TOKENS.to_string(), api_usage.total_tokens);
usage.extend(api_usage.to_usage_map());
}
}

Expand All @@ -274,7 +265,7 @@ impl OpenAICompatClient {
#[cfg(test)]
mod tests {
use super::OpenAICompatClient;
use crate::providers::traits::ProviderError;
use crate::providers::traits::{usage_key, ProviderError};

#[test]
fn usage_limit_http_429_is_typed_and_non_transient() {
Expand Down Expand Up @@ -307,4 +298,23 @@ mod tests {
} if message == "Slow down"
));
}

#[test]
fn sse_reassembly_normalizes_standard_cached_tokens() {
let body = concat!(
"data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":null}]}\n\n",
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":1200,\"completion_tokens\":300,\"total_tokens\":1500,\"prompt_tokens_details\":{\"cached_tokens\":800}}}\n\n",
"data: [DONE]\n\n"
);

let response = OpenAICompatClient::reassemble_sse_to_response(body)
.expect("OpenAI-compatible SSE body should reassemble");

assert_eq!(response.content.as_deref(), Some("ok"));
assert_eq!(response.usage[usage_key::PROMPT_TOKENS], 400);
assert_eq!(response.usage[usage_key::COMPLETION_TOKENS], 300);
assert_eq!(response.usage[usage_key::TOTAL_TOKENS], 1500);
assert_eq!(response.usage[usage_key::CACHE_READ_TOKENS], 800);
assert!(!response.usage.contains_key(usage_key::CACHE_WRITE_TOKENS));
}
}
73 changes: 66 additions & 7 deletions src-tauri/crates/agent-core/src/core/providers/openai_compat/streaming/chat.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
//! Non-streaming `chat()` implementation for OpenAI-compatible providers.

use serde_json::Value;
use std::collections::HashMap;
use tracing::{info, warn};

use super::super::client::OpenAICompatClient;
Expand Down Expand Up @@ -173,12 +172,10 @@ pub(super) async fn run_chat(
.next()
.ok_or_else(|| ProviderError::ParseError("No choices in response".to_string()))?;

let mut usage = HashMap::new();
if let Some(api_usage) = parsed.usage {
usage.insert("prompt_tokens".to_string(), api_usage.prompt_tokens);
usage.insert("completion_tokens".to_string(), api_usage.completion_tokens);
usage.insert("total_tokens".to_string(), api_usage.total_tokens);
}
let usage = parsed
.usage
.map(|api_usage| api_usage.to_usage_map())
.unwrap_or_default();

let tool_calls = choice
.message
Expand Down Expand Up @@ -262,3 +259,65 @@ fn split_inline_thinking(

(content_out, merged_reasoning)
}

#[cfg(test)]
mod tests {
use super::*;
use crate::providers::registry::{find_by_name, provider_id};
use crate::providers::traits::{usage_key, LLMProvider, ProviderConfig};
use std::collections::HashMap;
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};

#[tokio::test]
async fn non_streaming_standard_usage_normalizes_cached_tokens() {
crate::test_support::install_crypto_provider_for_tests();

let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/chat/completions"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"choices": [{
"message": {"content": "ok"},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 1200,
"completion_tokens": 300,
"total_tokens": 1500,
"prompt_tokens_details": {"cached_tokens": 800}
}
})))
.mount(&server)
.await;

let spec = find_by_name(provider_id::OPENAI).expect("OpenAI provider registered");
let client = OpenAICompatClient::new(
ProviderConfig {
api_key: "test-key".to_string(),
api_base: Some(server.uri()),
extra_headers: HashMap::new(),
is_azure: false,
},
spec,
"gpt-4.1".to_string(),
);

let response = client
.chat(
&[serde_json::json!({"role": "user", "content": "hello"})],
None,
"gpt-4.1",
1024,
0.0,
)
.await
.expect("OpenAI-compatible response should parse");

assert_eq!(response.usage[usage_key::PROMPT_TOKENS], 400);
assert_eq!(response.usage[usage_key::COMPLETION_TOKENS], 300);
assert_eq!(response.usage[usage_key::TOTAL_TOKENS], 1500);
assert_eq!(response.usage[usage_key::CACHE_READ_TOKENS], 800);
assert!(!response.usage.contains_key(usage_key::CACHE_WRITE_TOKENS));
}
}
75 changes: 68 additions & 7 deletions src-tauri/crates/agent-core/src/core/providers/openai_compat/streaming/sse_stream.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ use crate::providers::openai_policy::ChatTokenLimitField;
use crate::providers::registry::provider_id;
use crate::providers::safe_truncate::safe_truncate_utf8;
use crate::providers::traits::{
finish_reason as finish, LLMResponse, ProviderError, StreamDelta, StreamErrorKind,
finish_reason as finish, usage_key, LLMResponse, ProviderError, StreamDelta, StreamErrorKind,
ToolCallDelta, ToolCallRequest,
};
use crate::providers::wire_sanitize::{
Expand Down Expand Up @@ -514,13 +514,22 @@ pub(super) async fn run_chat_streaming(

// Usage (usually on final chunk when stream_options.include_usage=true)
if let Some(ref usage) = chunk.usage {
let normalized_usage = usage.to_usage_map();
debug!(
"[streaming-usage] OpenAI chunk usage: prompt={}, completion={}, total={}",
usage.prompt_tokens, usage.completion_tokens, usage.total_tokens
"[streaming-usage] OpenAI chunk usage: prompt={}, completion={}, total={}, cache_read={}, cache_write={}",
usage.prompt_tokens,
usage.completion_tokens,
usage.total_tokens,
normalized_usage
.get(usage_key::CACHE_READ_TOKENS)
.copied()
.unwrap_or(0),
normalized_usage
.get(usage_key::CACHE_WRITE_TOKENS)
.copied()
.unwrap_or(0),
);
final_usage.insert("prompt_tokens".to_string(), usage.prompt_tokens);
final_usage.insert("completion_tokens".to_string(), usage.completion_tokens);
final_usage.insert("total_tokens".to_string(), usage.total_tokens);
final_usage.extend(normalized_usage);
}
}
if stream_done {
Expand Down Expand Up @@ -667,7 +676,7 @@ pub(super) async fn run_chat_streaming(
mod tests {
use super::*;
use crate::providers::registry::{find_by_name, provider_id};
use crate::providers::traits::{LLMProvider, ProviderConfig};
use crate::providers::traits::{usage_key, LLMProvider, ProviderConfig};
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};

Expand Down Expand Up @@ -757,4 +766,56 @@ mod tests {
assert!(response.tool_calls.is_empty());
assert_eq!(response.content, None);
}

#[tokio::test]
async fn deepseek_stream_normalizes_cache_hit_and_miss_tokens() {
crate::test_support::install_crypto_provider_for_tests();

let server = MockServer::start().await;
let body = concat!(
"data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":null}]}\n\n",
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":1200,\"completion_tokens\":300,\"total_tokens\":1500,\"prompt_cache_hit_tokens\":800,\"prompt_cache_miss_tokens\":400}}\n\n",
"data: [DONE]\n\n"
);
Mock::given(method("POST"))
.and(path("/chat/completions"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(body),
)
.mount(&server)
.await;

let spec = find_by_name(provider_id::DEEPSEEK).expect("DeepSeek provider registered");
let client = OpenAICompatClient::new(
ProviderConfig {
api_key: "test-key".to_string(),
api_base: Some(server.uri()),
extra_headers: HashMap::new(),
is_azure: false,
},
spec,
"deepseek-chat".to_string(),
);

let response = client
.chat_streaming(
&[serde_json::json!({"role": "user", "content": "hello"})],
None,
"deepseek-chat",
1024,
0.0,
&|_| {},
None,
)
.await
.expect("DeepSeek stream should parse");

assert_eq!(response.usage[usage_key::PROMPT_TOKENS], 400);
assert_eq!(response.usage[usage_key::COMPLETION_TOKENS], 300);
assert_eq!(response.usage[usage_key::TOTAL_TOKENS], 1500);
assert_eq!(response.usage[usage_key::CACHE_READ_TOKENS], 800);
assert!(!response.usage.contains_key(usage_key::CACHE_WRITE_TOKENS));
}
}
Loading
Loading