From 865be2bf75bb62b6aeee059f684400b4b9938a15 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Mon, 29 Apr 2024 06:51:03 +0800
Subject: feat: non-streaming returns completion stats (#456)

---
 src/client/claude.rs | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

(limited to 'src/client/claude.rs')

diff --git a/src/client/claude.rs b/src/client/claude.rs
index e4081e7..8bd87ee 100644
--- a/src/client/claude.rs
+++ b/src/client/claude.rs
@@ -1,6 +1,6 @@
 use super::{
-    catch_error, extract_system_message, ClaudeClient, ExtraConfig, ImageUrl, MessageContent,
-    MessageContentPart, Model, ModelConfig, PromptType, ReplyHandler, SendData,
+    catch_error, extract_system_message, ClaudeClient, CompletionStats, ExtraConfig, ImageUrl,
+    MessageContent, MessageContentPart, Model, ModelConfig, PromptType, ReplyHandler, SendData,
 };
 
 use crate::utils::PromptKind;
@@ -54,19 +54,14 @@ impl_client_trait!(
     claude_send_message_streaming
 );
 
-pub async fn claude_send_message(builder: RequestBuilder) -> Result<String> {
+pub async fn claude_send_message(builder: RequestBuilder) -> Result<(String, CompletionStats)> {
     let res = builder.send().await?;
     let status = res.status();
     let data: Value = res.json().await?;
     if status != 200 {
         catch_error(&data, status.as_u16())?;
     }
-
-    let output = data["content"][0]["text"]
-        .as_str()
-        .ok_or_else(|| anyhow!("Invalid response data: {data}"))?;
-
-    Ok(output.to_string())
+    claude_extract_completion(&data)
 }
 
 pub async fn claude_send_message_streaming(
@@ -195,3 +190,16 @@ pub fn claude_build_body(data: SendData, model: &Model) -> Result<Value> {
     }
     Ok(body)
 }
+
+pub fn claude_extract_completion(data: &Value) -> Result<(String, CompletionStats)> {
+    let text = data["content"][0]["text"]
+        .as_str()
+        .ok_or_else(|| anyhow!("Invalid response data: {data}"))?;
+
+    let stats = CompletionStats {
+        id: data["id"].as_str().map(|v| v.to_string()),
+        input_tokens: data["usage"]["input_tokens"].as_u64(),
+        output_tokens: data["usage"]["output_tokens"].as_u64(),
+    };
+    Ok((text.to_string(), stats))
+}
-- 
cgit v1.2.3