update async text releaser to handle non-string

This commit is contained in:
2026-01-30 13:49:47 +08:00
parent 6c72948f5b
commit ece43eec89

View File

@@ -405,7 +405,7 @@ class AsyncTextReleaser:
return "".join(result_parts) if result_parts else None return "".join(result_parts) if result_parts else None
async def release(self, text_iterator: AsyncIterator[str]) -> AsyncIterator[str]: async def release(self, text_iterator: AsyncIterator[Any]) -> AsyncIterator[Any]:
""" """
Async version of release that works with async generators. Async version of release that works with async generators.
@@ -426,10 +426,13 @@ class AsyncTextReleaser:
nonlocal producer_done nonlocal producer_done
async for chunk in text_iterator: async for chunk in text_iterator:
start_pos = len(self._accumulated_text) if isinstance(chunk, str):
self._accumulated_text += chunk start_pos = len(self._accumulated_text)
end_pos = len(self._accumulated_text) self._accumulated_text += chunk
buffer.append((chunk, start_pos, end_pos)) end_pos = len(self._accumulated_text)
buffer.append((chunk, start_pos, end_pos))
else:
buffer.append((chunk, None, None))
# Process available chunks # Process available chunks
self._search_for_keys(state, len(self._accumulated_text)) self._search_for_keys(state, len(self._accumulated_text))
@@ -439,6 +442,12 @@ class AsyncTextReleaser:
chunk_at_yield = buffer[state.yield_idx] chunk_at_yield = buffer[state.yield_idx]
y_chunk, y_start, y_end = chunk_at_yield y_chunk, y_start, y_end = chunk_at_yield
# If it is not string; return the thing
if y_start is None: # Non-string item - yield immediately
state.yield_idx += 1
yield y_chunk
continue
if y_end > safe_end_pos: if y_end > safe_end_pos:
break break
@@ -471,6 +480,12 @@ class AsyncTextReleaser:
chunk_at_yield = buffer[state.yield_idx] chunk_at_yield = buffer[state.yield_idx]
y_chunk, y_start, y_end = chunk_at_yield y_chunk, y_start, y_end = chunk_at_yield
# If it is not string; return the thing
if y_start is None: # Non-string item - yield immediately
state.yield_idx += 1
yield y_chunk
continue
self._update_delay_mode(state, y_start, y_end) self._update_delay_mode(state, y_start, y_end)
if self._should_skip_to_end_key(state, y_end): if self._should_skip_to_end_key(state, y_end):