Skip to content

Commit

Permalink
Cleaner stream handling in Answer class (#2314)
Browse files Browse the repository at this point in the history
* add cleaner stream

* add cleaner stream handling
  • Loading branch information
pablonyx authored Sep 3, 2024
1 parent af66650 commit fb95398
Showing 1 changed file with 14 additions and 8 deletions.
22 changes: 14 additions & 8 deletions backend/danswer/llm/answering/answer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit fb95398

Please sign in to comment.