From e45755ffea0e14626fd49b62661979d8548d1537 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9=20Cipriani=20Bandarra?= <andreban@gmail.com>
Date: Thu, 2 May 2024 18:36:37 +0100
Subject: [PATCH] Small refactor to close stream as soon as response is
 finished

---
 src/client.rs | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index aace80f..58c05ef 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -85,11 +85,24 @@ impl<T: TokenProvider + Clone> GeminiClient<T> {
                         if let Event::Message(event) = event {
                             let response: serde_json::error::Result<GenerateContentResponse> =
                                 serde_json::from_str(&event.data);
-                            if let Ok(response) = response {
-                                cloned_queue.push(Some(response.into_result()));
-                            } else {
-                                tracing::error!("Error parsing message: {}", event.data);
-                            };
+
+                            match response {
+                                Ok(response) => {
+                                    let result = response.into_result();
+                                    let finished = match &result {
+                                        Ok(result) => result.candidates[0].finish_reason.is_some(),
+                                        Err(_) => true,
+                                    };
+                                    cloned_queue.push(Some(result));
+                                    if finished {
+                                        break;
+                                    }
+                                }
+                                Err(_) => {
+                                    tracing::error!("Error parsing message: {}", event.data);
+                                    break;
+                                }
+                            }
                         }
                     }
                     Err(e) => {