From c4d759dfba8d72298ad405ad9cc79083b306d782 Mon Sep 17 00:00:00 2001 From: takatost Date: Thu, 17 Aug 2023 13:40:00 +0800 Subject: [PATCH] fix: wenxin error not raise when stream mode (#884) --- api/core/third_party/langchain/llms/wenxin.py | 26 ++++++++++++++----- 1 file changed, 20 insertions(+), 6 deletions(-) diff --git a/api/core/third_party/langchain/llms/wenxin.py b/api/core/third_party/langchain/llms/wenxin.py index a10fb82b71..7b31b67624 100644 --- a/api/core/third_party/langchain/llms/wenxin.py +++ b/api/core/third_party/langchain/llms/wenxin.py @@ -3,6 +3,7 @@ from __future__ import annotations import json import logging +from json import JSONDecodeError from typing import ( Any, Dict, @@ -223,11 +224,24 @@ class Wenxin(LLM): for token in self._client.post(request).iter_lines(): if token: token = token.decode("utf-8") - completion = json.loads(token[5:]) - yield GenerationChunk(text=completion['result']) - if run_manager: - run_manager.on_llm_new_token(completion['result']) + if token.startswith('data:'): + completion = json.loads(token[5:]) - if completion['is_end']: - break + yield GenerationChunk(text=completion['result']) + if run_manager: + run_manager.on_llm_new_token(completion['result']) + + if completion['is_end']: + break + else: + try: + json_response = json.loads(token) + except JSONDecodeError: + raise ValueError(f"Wenxin Response Error {token}") + + raise ValueError( + f"Wenxin API {json_response['error_code']}" + f" error: {json_response['error_msg']}, " + f"please confirm if the model you have chosen is already paid for." + )