summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
-rw-r--r--phind/__init__.py47
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