diff options
| -rw-r--r-- | phind/__init__.py | 47 |
1 files changed, 28 insertions, 19 deletions
diff --git a/phind/__init__.py b/phind/__init__.py index 469c09b7..69468703 100644 --- a/phind/__init__.py +++ b/phind/__init__.py @@ -204,25 +204,34 @@ class StreamingCompletion: while StreamingCompletion.stream_completed != True or not StreamingCompletion.message_queue.empty(): try: - message = StreamingCompletion.message_queue.get(timeout=0) - for token in findall(r'(?<=data: )(.+?)(?=\r\n\r\n)', message.decode()): - yield PhindResponse({ - 'id' : f'cmpl-1337-{int(time())}', - 'object' : 'text_completion', - 'created': int(time()), - 'model' : model, - 'choices': [{ - 'text' : token, - 'index' : 0, - 'logprobs' : None, - 'finish_reason' : 'stop' - }], - 'usage': { - 'prompt_tokens' : len(prompt), - 'completion_tokens' : len(token), - 'total_tokens' : len(prompt) + len(token) - } - }) + chunk = StreamingCompletion.message_queue.get(timeout=0) + + if chunk == b'data: \r\ndata: \r\ndata: \r\n\r\n': + chunk = b'data: \n\n\r\n\r\n' + + chunk = chunk.decode() + + chunk = chunk.replace('data: \r\n\r\ndata: ', 'data: \n') + chunk = chunk.replace('\r\ndata: \r\ndata: \r\n\r\n', '\n\n\r\n\r\n') + chunk = chunk.replace('data: ', '').replace('\r\n\r\n', '') + + yield PhindResponse({ + 'id' : f'cmpl-1337-{int(time())}', + 'object' : 'text_completion', + 'created': int(time()), + 'model' : model, + 'choices': [{ + 'text' : chunk, + 'index' : 0, + 'logprobs' : None, + 'finish_reason' : 'stop' + }], + 'usage': { + 'prompt_tokens' : len(prompt), + 'completion_tokens' : len(chunk), + 'total_tokens' : len(prompt) + len(chunk) + } + }) except Empty: pass |