fix: OpenAI finish response parsing

dustinblackman · Jan 30, 2024 · ec3b054 · ec3b054
1 parent 9b950a5
commit ec3b054
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 5 deletions.
diff --git a/src/infrastructure/backends/openai.rs b/src/infrastructure/backends/openai.rs
@@ -53,12 +53,13 @@ struct CompletionRequest {
 
 #[derive(Default, Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 struct CompletionDeltaResponse {
-    content: String,
+    content: Option<String>,
 }
 
 #[derive(Default, Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 struct CompletionChoiceResponse {
     delta: CompletionDeltaResponse,
+    finish_reason: Option<String>,
 }
 
 #[derive(Default, Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -202,7 +203,19 @@ impl Backend for OpenAI {
             let ores: CompletionResponse = serde_json::from_str(&cleaned_line).unwrap();
             tracing::debug!(body = ?ores, "Completion response");
 
-            let text = ores.choices[0].delta.content.to_string();
+            let choice = &ores.choices[0];
+            if choice.finish_reason.is_some() {
+                break;
+            }
+            if choice.delta.content.is_none() {
+                continue;
+            }
+
+            let text = choice.delta.content.clone().unwrap().to_string();
+            if text.is_empty() {
+                continue;
+            }
+
             last_message += &text;
             let msg = BackendResponse {
                 author: Author::Model,

diff --git a/src/infrastructure/backends/openai_test.rs b/src/infrastructure/backends/openai_test.rs
@@ -102,20 +102,29 @@ async fn it_gets_completions() -> Result<()> {
     let first_line = serde_json::to_string(&CompletionResponse {
         choices: vec![CompletionChoiceResponse {
             delta: CompletionDeltaResponse {
-                content: "Hello ".to_string(),
+                content: Some("Hello ".to_string()),
             },
+            finish_reason: None,
         }],
     })?;
 
     let second_line = serde_json::to_string(&CompletionResponse {
         choices: vec![CompletionChoiceResponse {
             delta: CompletionDeltaResponse {
-                content: "World".to_string(),
+                content: Some("World".to_string()),
             },
+            finish_reason: None,
         }],
     })?;
 
-    let body = [first_line, second_line].join("\n");
+    let third_line = serde_json::to_string(&CompletionResponse {
+        choices: vec![CompletionChoiceResponse {
+            delta: CompletionDeltaResponse { content: None },
+            finish_reason: Some("stop".to_string()),
+        }],
+    })?;
+
+    let body = [first_line, second_line, third_line].join("\n");
     let prompt = BackendPrompt {
         text: "Say hi to the world".to_string(),
         backend_context: serde_json::to_string(&vec![MessageRequest {