From 5281e45e75f014817d9673aa410fc635dd25bf2f Mon Sep 17 00:00:00 2001
From: skeptrune <nick.k@trieve.ai>
Date: Mon, 9 Dec 2024 11:12:44 -0800
Subject: [PATCH] cleanup: remove expects for LLM completion calls

---
 server/src/handlers/message_handler.rs   | 31 +++++++++++++++---------
 server/src/operators/message_operator.rs | 18 ++++++++++----
 2 files changed, 33 insertions(+), 16 deletions(-)

diff --git a/server/src/handlers/message_handler.rs b/server/src/handlers/message_handler.rs
index 6ca865d52e..f4f44ee606 100644
--- a/server/src/handlers/message_handler.rs
+++ b/server/src/handlers/message_handler.rs
@@ -1003,17 +1003,26 @@ pub async fn get_suggested_queries(
     .collect();
 
     while queries.len() < 3 {
-        query = client
-            .chat()
-            .create(parameters.clone())
-            .await
-            .expect("No LLM Completion for topic");
-        queries = match &query
-            .choices
-            .first()
-            .expect("No response for LLM completion")
-            .message
-        {
+        query = match client.chat().create(parameters.clone()).await {
+            Ok(query) => query,
+            Err(err) => {
+                log::error!(
+                    "Error generating suggested queries when queries are less than 3: {}",
+                    err
+                );
+                return Err(ServiceError::BadRequest(err.to_string()));
+            }
+        };
+        let first_query = match query.choices.first() {
+            Some(first_query) => first_query,
+            None => {
+                log::error!("Error generating suggested queries when queries are less than 3: No first query in choices");
+                return Err(ServiceError::BadRequest(
+                    "No first query in choices on call to LLM".to_string(),
+                ));
+            }
+        };
+        queries = match &first_query.message {
             ChatMessage::User {
                 content: ChatMessageContent::Text(content),
                 ..
diff --git a/server/src/operators/message_operator.rs b/server/src/operators/message_operator.rs
index d827953b42..927f62ea9e 100644
--- a/server/src/operators/message_operator.rs
+++ b/server/src/operators/message_operator.rs
@@ -303,11 +303,19 @@ pub async fn get_rag_chunks_query(
             ..Default::default()
         };
 
-        let search_query_from_message_to_query_prompt = client
-            .chat()
-            .create(gen_inference_parameters)
-            .await
-            .expect("No LLM Completion for chunk search");
+        let search_query_from_message_to_query_prompt =
+            match client.chat().create(gen_inference_parameters).await {
+                Ok(query) => query,
+                Err(err) => {
+                    log::error!(
+                        "Error getting LLM completion for message to query prompt {:?}",
+                        err
+                    );
+                    return Err(actix_web::error::ErrorInternalServerError(
+                        "Error getting LLM completion for message to query prompt",
+                    ));
+                }
+            };
 
         query = match &search_query_from_message_to_query_prompt
             .choices