Update networks/huggingface_streamer.py

networks/huggingface_streamer.py

@@ -203,7 +203,19 @@ class HuggingfaceStreamer:
             output = self.message_outputer.output(
                 content=content, content_type=content_type
             )
-            yield "{" + output + "}"
+            
+            output = json.loads(output.split('data: ')[1])
+            print("OUTPUT OLLAMA STREAMER", output)
+            output['repsonse'] = output['choices'][-1]['delta'].get('content')
+            output['done'] = output['choices'][-1].get('finish_reason')
+            output['created_at'] = output.get('created')
+            yield json.dumps(output)
 
         if not is_finished:
-            yield "{" + self.message_outputer.output(content="", content_type="Finished") + "}"
+            final_output = self.message_outputer.output(content="", content_type="Finished").split('data: ')[1]
+            final_output = json.loads(final_output)
+            final_output = final_output['choices'][-1]['delta'].get('content')
+            final_output['done'] = output['choices'][-1].get('finish_reason')
+            final_output['created_at'] = final_output.get('created')
+            yield json.dumps(final_output)
+            # yield self.message_outputer.output(content="", content_type="Finished").split('data: ')[1]