fix: openai SSE parsing

jiacai2050 · Jul 28, 2024 · 82a349b · 82a349b
1 parent e244f4b
commit 82a349b
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 10 deletions.
diff --git a/shellgpt/api/llm.py b/shellgpt/api/llm.py
@@ -57,32 +57,47 @@ def chat_openai(self, prompt, stream, add_system_message):
 
         answer = ''
         current = b''
+        # https://github.com/openai/openai-python#streaming-responses
+        # The response is SSE, so we need to parse the response line by line.
         for item in r.iter_content(chunk_size=None):
-            debug_print(f'item: {item}')
+            debug_print(f'\nitem: {item}\ncurrent: {current}')
             for msg in item.split(b'\n\n'):
-                msg = msg.strip(b'data: ')
+                msg = msg.removeprefix(b'data: ')
                 if len(msg) == 0:
                     continue
 
                 current += msg
 
-                # when current end with '}', it's the end of the message
+                # when current end with '}', it maybe the end of the message
                 if current[-1] == 125:
-                    msg = current
+                    # msg is a complete JSON message
+                    # `data:` may appear in the middle of the message, so we need to remove it again.
+                    msg = current.removeprefix(b'data: ')
                     current = b''
                 else:
                     continue
 
                 s = msg.decode('utf-8')
+                debug_print(f'\nitem to decode: {s}')
                 if s == '[DONE]':
                     self.messages.append({'role': 'assistant', 'content': answer})
                     return
                 else:
-                    resp = json.loads(s)
-                    for item in resp['choices']:
-                        msg = item['delta']['content']
-                        answer += msg
-                        yield msg
+                    try:
+                        resp = json.loads(s)
+                        for item in resp['choices']:
+                            if 'content' not in item['delta']:
+                                continue
+
+                            msg = item['delta']['content']
+                            answer += msg
+                            yield msg
+                    except json.JSONDecodeError as e:
+                        debug_print(f'Error when decode JSON: {s}, err:{e}')
+                        # this means the message is not a JSON message, so we need to continue searching next }.
+                        current = msg
+                        continue
+
 
     def make_messages(self, prompt, support_image, add_system_message):
         model = self.model

diff --git a/shellgpt/app.py b/shellgpt/app.py
@@ -15,6 +15,7 @@
     IS_TTY,
 )
 from .utils.common import (
+    is_verbose,
     load_contents_from_config,
     execute_cmd,
     copy_text,
@@ -26,7 +27,7 @@
 from .tui.app import ShellGPTApp
 from .history import History
 
-__version__ = '0.5.1'
+__version__ = '0.5.2'
 
 
 def init_app():
@@ -207,6 +208,8 @@ def infer(self, prompt):
                 print()
         except Exception as e:
             print(f'Error when infer: ${e}')
+            if is_verbose():
+                raise e
 
     def explain_cmd(self, cmd):
         self.last_answer = ''

diff --git a/shellgpt/utils/common.py b/shellgpt/utils/common.py
@@ -24,6 +24,8 @@ def set_verbose(v):
     global IS_VERBOSE
     IS_VERBOSE = v
 
+def is_verbose():
+    return IS_VERBOSE
 
 def debug_print(msg):
     if IS_VERBOSE: