cleanup: remove expects for LLM completion calls

devflowinc · Dec 9, 2024 · 5281e45 · 5281e45
1 parent be27319
commit 5281e45
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 16 deletions.
diff --git 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
@@ -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