From 8c22fa7560520441d5e836e6e726414ae219cd4c Mon Sep 17 00:00:00 2001
From: erhant <erhany96@gmail.com>
Date: Thu, 9 May 2024 10:05:54 +0300
Subject: [PATCH] flatten final data

---
 examples/basic_chatbot.rs        |  4 ++--
 src/generation/completion/mod.rs | 23 ++++++++---------------
 2 files changed, 10 insertions(+), 17 deletions(-)

diff --git a/examples/basic_chatbot.rs b/examples/basic_chatbot.rs
index 2637b94..2749e43 100644
--- a/examples/basic_chatbot.rs
+++ b/examples/basic_chatbot.rs
@@ -38,8 +38,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 stdout.write_all(ele.response.as_bytes()).await?;
                 stdout.flush().await?;
 
-                if let Some(final_data) = ele.final_data {
-                    context = Some(final_data.context);
+                if ele.context.is_some() {
+                    context = ele.context;
                 }
             }
         }
diff --git a/src/generation/completion/mod.rs b/src/generation/completion/mod.rs
index b81081a..a7ba8cd 100644
--- a/src/generation/completion/mod.rs
+++ b/src/generation/completion/mod.rs
@@ -103,23 +103,16 @@ pub struct GenerationResponse {
     pub response: String,
     /// Whether the completion is done. If the completion is streaming, this will be false until the last response.
     pub done: bool,
-    #[serde(flatten)]
-    /// The final data of the completion. This is only present if the completion is done.
-    pub final_data: Option<GenerationFinalResponseData>,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct GenerationFinalResponseData {
     /// An encoding of the conversation used in this response, this can be sent in the next request to keep a conversational memory
-    pub context: GenerationContext,
+    pub context: Option<GenerationContext>,
     /// Time spent generating the response
-    pub total_duration: u64,
+    pub total_duration: Option<u64>,
     /// Number of tokens in the prompt
-    pub prompt_eval_count: u16,
+    pub prompt_eval_count: Option<u16>,
     /// Time spent in nanoseconds evaluating the prompt
-    pub prompt_eval_duration: u64,
-    /// Number of tokens the response
-    pub eval_count: u16,
-    /// Time in nanoseconds spent generating the response
-    pub eval_duration: u64,
+    pub prompt_eval_duration: Option<u64>,
+    /// Number of tokens in the response
+    pub eval_count: Option<u16>,
+    /// Time spent in nanoseconds generating the response
+    pub eval_duration: Option<u64>,
 }