feat: non-streaming returns completion stats (#456)

1 files changed, 16 insertions, 5 deletions

diff --git a/src/client/common.rs b/src/client/common.rs
index 695245a..a58c962 100644
--- a/src/client/common.rs
+++ b/src/client/common.rs
@@ -126,7 +126,7 @@ macro_rules! register_client {
                     client.set_model(model);
                 } else {
                     anyhow::bail!(
-                        "The current model lacks the corresponding capability."
+                        "The current model is incapable of doing that."
                     );
                 }
             }
@@ -260,7 +260,7 @@ macro_rules! impl_client_trait {
                 &self,
                 client: &reqwest::Client,
                 data: $crate::client::SendData,
-            ) -> anyhow::Result<String> {
+            ) -> anyhow::Result<(String, $crate::client::CompletionStats)> {
                 let builder = self.request_builder(client, data)?;
                 $send_message(builder).await
             }
@@ -330,11 +330,11 @@ pub trait Client: Sync + Send {
         Ok(client)
     }
 
-    async fn send_message(&self, input: Input) -> Result<String> {
+    async fn send_message(&self, input: Input) -> Result<(String, CompletionStats)> {
         let global_config = self.config().0;
         if global_config.read().dry_run {
             let content = global_config.read().echo_messages(&input);
-            return Ok(content);
+            return Ok((content, CompletionStats::default()));
         }
         let client = self.build_client()?;
         let data = global_config.read().prepare_send_data(&input, false)?;
@@ -384,7 +384,11 @@ pub trait Client: Sync + Send {
         }
     }
 
-    async fn send_message_inner(&self, client: &ReqwestClient, data: SendData) -> Result<String>;
+    async fn send_message_inner(
+        &self,
+        client: &ReqwestClient,
+        data: SendData,
+    ) -> Result<(String, CompletionStats)>;
 
     async fn send_message_streaming_inner(
         &self,
@@ -414,6 +418,13 @@ pub struct SendData {
     pub stream: bool,
 }
 
+#[derive(Debug, Clone, Default)]
+pub struct CompletionStats {
+    pub id: Option<String>,
+    pub input_tokens: Option<u64>,
+    pub output_tokens: Option<u64>,
+}
+
 pub type PromptType<'a> = (&'a str, &'a str, bool, PromptKind);
 
 pub fn create_config(list: &[PromptType], client: &str) -> Result<(String, Value)> {