From fb95398e5b0e491825a131b87a64146102bd3a4b Mon Sep 17 00:00:00 2001 From: pablodanswer Date: Tue, 3 Sep 2024 11:36:01 -0700 Subject: [PATCH] Cleaner stream handling in Answer class (#2314) * add cleaner stream * add cleaner stream handling --- backend/danswer/llm/answering/answer.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/backend/danswer/llm/answering/answer.py b/backend/danswer/llm/answering/answer.py index 7e7733cc0fe..0a0a1c52afa 100644 --- a/backend/danswer/llm/answering/answer.py +++ b/backend/danswer/llm/answering/answer.py @@ -550,15 +550,21 @@ def _process_stream( answer_style_configs=self.answer_style_config, ) - def _stream() -> Iterator[str | StreamStopInfo]: - yield cast(str | StreamStopInfo, message) - yield from (cast(str | StreamStopInfo, item) for item in stream) + stream_stop_info = None + + def _stream() -> Iterator[str]: + nonlocal stream_stop_info + yield cast(str, message) + for item in stream: + if isinstance(item, StreamStopInfo): + stream_stop_info = item + return + yield cast(str, item) - for item in _stream(): - if isinstance(item, StreamStopInfo): - yield item - else: - yield from process_answer_stream_fn(iter([item])) + yield from process_answer_stream_fn(_stream()) + + if stream_stop_info: + yield stream_stop_info processed_stream = [] for processed_packet in _process_stream(output_generator):