Merge pull request #16 from atomiechen/dev

Dev

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "HandyLLM"
-version = "0.5.4"
+version = "0.5.5"
 authors = [
   { name="Atomie CHEN", email="[email protected]" },
 ]

src/handyllm/openai_api.py

@@ -84,18 +84,15 @@ def get_model_engine_map(cls, model_engine_map=None):
     def stream_chat_with_role(response):
         role = ''
         for data in response:
-            if 'choices' not in data:
-                continue
-            if not data['choices']:
-                continue
-            if 'delta' not in data['choices'][0]:
-                continue
-            message = data['choices'][0]['delta']
-            if 'role' in message:
-                role = message['role']
-            if 'content' in message:
-                text = message['content']
-                yield role, text
+            try:
+                message = data['choices'][0]['delta']
+                if 'role' in message:
+                    role = message['role']
+                if 'content' in message:
+                    text = message['content']
+                    yield role, text
+            except (KeyError, IndexError):
+                pass
 
     @staticmethod
     def stream_chat(response):
@@ -105,13 +102,10 @@ def stream_chat(response):
     @staticmethod
     def stream_completions(response):
         for data in response:
-            if 'choices' not in data:
-                continue
-            if not data['choices']:
-                continue
-            if 'text' not in data['choices'][0]:
-                continue
-            yield data['choices'][0]['text']
+            try:
+                yield data['choices'][0]['text']
+            except (KeyError, IndexError):
+                pass
 
     @classmethod
     def api_request_endpoint(
@@ -223,18 +217,24 @@ def wrapper(response):
                         text = ''
                         role = ''
                         for data in response:
-                            message = data['choices'][0]['delta']
-                            if 'role' in message:
-                                role = message['role']
-                            if 'content' in message:
-                                text += message['content']
+                            try:
+                                message = data['choices'][0]['delta']
+                                if 'role' in message:
+                                    role = message['role']
+                                if 'content' in message:
+                                    text += message['content']
+                            except (KeyError, IndexError):
+                                pass
                             yield data
                         log_strs.append(cls.converter.chat2raw([{'role': role, 'content': text}]))
                         log_strs.append(" OUTPUT END ".center(50, '-')+"\n")
                         logger.info('\n'.join(log_strs))
                     response = wrapper(response)
                 else:
-                    log_strs.append(cls.converter.chat2raw([response['choices'][0]['message']]))
+                    try:
+                        log_strs.append(cls.converter.chat2raw([response['choices'][0]['message']]))
+                    except (KeyError, IndexError):
+                        log_strs.append("Wrong response format, no message found")
                     log_strs.append(" OUTPUT END ".center(50, '-')+"\n")
                     logger.info('\n'.join(log_strs))
         except Exception as e:
@@ -294,14 +294,20 @@ def completions(cls, prompt, logger=None, log_marks=[], **kwargs):
                     def wrapper(response):
                         text = ''
                         for data in response:
-                            text += data['choices'][0]['text']
+                            try:
+                                text += data['choices'][0]['text']
+                            except (KeyError, IndexError):
+                                pass
                             yield data
                         log_strs.append(text)
                         log_strs.append(" OUTPUT END ".center(50, '-')+"\n")
                         logger.info('\n'.join(log_strs))
                     response = wrapper(response)
                 else:
-                    log_strs.append(response['choices'][0]['text'])
+                    try:
+                        log_strs.append(response['choices'][0]['text'])
+                    except (KeyError, IndexError):
+                        log_strs.append("Wrong response format, no text found")
                     log_strs.append(" OUTPUT END ".center(50, '-')+"\n")
                     logger.info('\n'.join(log_strs))
         except Exception as e: