This commit is contained in:
Paul Gauthier 2024-06-28 15:27:17 -07:00
parent d6467a8e30
commit e7aa10a89b

View file

@ -812,7 +812,7 @@ class Coder:
utils.show_messages(messages, functions=self.functions) utils.show_messages(messages, functions=self.functions)
self.multi_response_content = "" self.multi_response_content = ""
if self.show_pretty(): if self.show_pretty() and self.stream:
mdargs = dict(style=self.assistant_output_color, code_theme=self.code_theme) mdargs = dict(style=self.assistant_output_color, code_theme=self.code_theme)
self.mdstream = MarkdownStream(mdargs=mdargs) self.mdstream = MarkdownStream(mdargs=mdargs)
else: else:
@ -851,6 +851,9 @@ class Coder:
traceback.print_exc() traceback.print_exc()
return return
if self.mdstream:
self.live_incremental_response(True)
if self.multi_response_content: if self.multi_response_content:
self.multi_response_content += self.partial_response_content self.multi_response_content += self.partial_response_content
self.partial_response_content = self.multi_response_content self.partial_response_content = self.multi_response_content
@ -1158,48 +1161,40 @@ class Coder:
raise FinishReasonLength() raise FinishReasonLength()
def show_send_output_stream(self, completion): def show_send_output_stream(self, completion):
finish_reason_length = False for chunk in completion:
if len(chunk.choices) == 0:
continue
try: if (
for chunk in completion: hasattr(chunk.choices[0], "finish_reason")
if len(chunk.choices) == 0: and chunk.choices[0].finish_reason == "length"
continue ):
raise FinishReasonLength()
if ( try:
hasattr(chunk.choices[0], "finish_reason") func = chunk.choices[0].delta.function_call
and chunk.choices[0].finish_reason == "length" # dump(func)
): for k, v in func.items():
if self.main_model.can_prefill: if k in self.partial_response_function_call:
finish_reason_length = True self.partial_response_function_call[k] += v
raise FinishReasonLength() else:
self.partial_response_function_call[k] = v
except AttributeError:
pass
try: try:
func = chunk.choices[0].delta.function_call text = chunk.choices[0].delta.content
# dump(func) if text:
for k, v in func.items(): self.partial_response_content += text
if k in self.partial_response_function_call: except AttributeError:
self.partial_response_function_call[k] += v text = None
else:
self.partial_response_function_call[k] = v
except AttributeError:
pass
try: if self.show_pretty():
text = chunk.choices[0].delta.content self.live_incremental_response(False)
if text: elif text:
self.partial_response_content += text sys.stdout.write(text)
except AttributeError: sys.stdout.flush()
text = None yield text
if self.show_pretty():
self.live_incremental_response(False)
elif text:
sys.stdout.write(text)
sys.stdout.flush()
yield text
finally:
if self.show_pretty() and not finish_reason_length:
self.live_incremental_response(True)
def live_incremental_response(self, final): def live_incremental_response(self, final):
show_resp = self.render_incremental_response(final) show_resp = self.render_incremental_response(final)