Fix JSON unmarshaling and error handling in chatStreamGemini function

swuecho · Sep 15, 2024 · 22635cc · 22635cc
1 parent b778147
commit 22635cc
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/api/chat_main_handler.go b/api/chat_main_handler.go
@@ -1483,11 +1483,19 @@ func (h *ChatHandler) chatStreamGemini(w http.ResponseWriter, chatSession sqlc_q
 			log.Printf("%s", line)
 		}
 		if err != nil {
+			// Create an instance of ErrorResponse
+			var errorResponse gemini.ErrorResponse
+
+			// Unmarshal the JSON string into the ErrorResponse struct
+			err := json.Unmarshal(line, &errorResponse)
+			if err != nil {
+				fmt.Println("Error unmarshaling JSON:", err)
+			}
 			if errors.Is(err, io.EOF) {
-				log.Printf("End of stream reached: %+v", err)
+				log.Printf("End of stream reached: %+v, %+v", err, errorResponse)
 				return answer, answer_id, false
 			} else {
-				log.Printf("Error while reading response: %+v", err)
+				log.Printf("Error while reading response: %+v, %+v", err, errorResponse)
 				return "", "", true
 			}
 		}

diff --git a/api/llm/gemini/gemini.go b/api/llm/gemini/gemini.go
@@ -182,3 +182,11 @@ func GenGemminPayload(chat_compeletion_messages []models.Message, chatFiles []sq
 	}
 	return payloadBytes, nil
 }
+
+type ErrorResponse struct {
+	Error struct {
+		Code    int    `json:"code"`
+		Message string `json:"message"`
+		Status  string `json:"status"`
+	} `json:"error"`
+}