From ead3a240d09eac61adf9d8f5194aeb9b0f39d8be Mon Sep 17 00:00:00 2001
From: CloudBridge UY <admin@cloudbridge.com.uy>
Date: Tue, 15 Aug 2023 12:13:22 -0300
Subject: [PATCH] fix(commands): Stop spinner after response and other fixes

- anthropic, openai, and vertex commands now stop the spinner after receiving the response
- Added #[serde(rename = "claude-2")] attribute to Model::Claude2
- Citation fields are now optional
---
 crates/c/src/commands/anthropic.rs |  8 +++++---
 crates/c/src/commands/openai.rs    |  6 +++---
 crates/c/src/commands/vertex.rs    | 18 +++++++++---------
 3 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/crates/c/src/commands/anthropic.rs b/crates/c/src/commands/anthropic.rs
index 6a288be..f78e942 100644
--- a/crates/c/src/commands/anthropic.rs
+++ b/crates/c/src/commands/anthropic.rs
@@ -23,6 +23,7 @@ pub struct Chunk {
 #[serde(rename_all = "kebab-case")]
 pub enum Model {
     #[default]
+    #[serde(rename = "claude-2")]
     Claude2,
     ClaudeV1,
     ClaudeV1_100k,
@@ -335,6 +336,9 @@ pub async fn run(mut options: CommandOptions) -> Result<()> {
     } else {
         let response = complete(&session).await?;
 
+        // Stop the spinner.
+        spinner.stop();
+
         // Print the response output.
         print_output(&session.meta.format, &response)?;
 
@@ -349,9 +353,6 @@ pub async fn run(mut options: CommandOptions) -> Result<()> {
     // Save the session to a file.
     session.save()?;
 
-    // Stop the spinner.
-    spinner.stop();
-
     Ok(())
 }
 
@@ -362,6 +363,7 @@ pub fn complete_prompt(
     max_tokens_to_sample: u32,
 ) -> Result<String> {
     let max = max_supported_tokens - max_tokens_to_sample;
+
     messages.push(Message::new("".to_string(), Role::Assistant, false));
 
     tracing::event!(
diff --git a/crates/c/src/commands/openai.rs b/crates/c/src/commands/openai.rs
index e08aabb..b760889 100644
--- a/crates/c/src/commands/openai.rs
+++ b/crates/c/src/commands/openai.rs
@@ -509,6 +509,9 @@ pub async fn run(mut options: CommandOptions) -> Result<()> {
     } else {
         let response = complete(&session).await?;
 
+        // Stop the spinner.
+        spinner.stop();
+
         // Print the response output.
         print_output(&session.meta.format, &response)?;
 
@@ -523,9 +526,6 @@ pub async fn run(mut options: CommandOptions) -> Result<()> {
     // Save the session to a file.
     session.save()?;
 
-    // Stop the spinner.
-    spinner.stop();
-
     Ok(())
 }
 
diff --git a/crates/c/src/commands/vertex.rs b/crates/c/src/commands/vertex.rs
index 8e77a21..6f7faaa 100644
--- a/crates/c/src/commands/vertex.rs
+++ b/crates/c/src/commands/vertex.rs
@@ -159,12 +159,12 @@ pub struct Candidate {
 #[derive(Debug, Serialize, Deserialize, Default, Clone)]
 #[serde(rename_all = "camelCase")]
 pub struct Citation {
-    start_index: u32,
-    end_index: u32,
-    url: String,
-    title: String,
-    license: String,
-    publication_date: String,
+    start_index: Option<u32>,
+    end_index: Option<u32>,
+    url: Option<String>,
+    title: Option<String>,
+    license: Option<String>,
+    publication_date: Option<String>,
 }
 
 #[derive(Debug, Serialize, Deserialize, Default, Clone)]
@@ -401,6 +401,9 @@ pub async fn run(mut options: CommandOptions) -> Result<()> {
 
     let response = complete(&session).await?;
 
+    // Stop the spinner.
+    spinner.stop();
+
     // Print the response output.
     print_output(&session.meta.format, &response)?;
 
@@ -422,9 +425,6 @@ pub async fn run(mut options: CommandOptions) -> Result<()> {
     // Save the session to a file.
     session.save()?;
 
-    // Stop the spinner.
-    spinner.stop();
-
     Ok(())
 }