38 lines
1.1 KiB
Python
38 lines
1.1 KiB
Python
"""Amortize LangChain's partial-JSON parsing without buffering live UI output."""
|
|
|
|
from langchain_core.messages import AIMessageChunk
|
|
from langchain_core.messages.ai import add_ai_message_chunks
|
|
|
|
CHUNK_BATCH_SIZE = 128
|
|
|
|
|
|
class ChunkAccumulator:
|
|
def __init__(self):
|
|
self.response = None
|
|
self.pending = []
|
|
|
|
@property
|
|
def has_output(self):
|
|
return self.response is not None or bool(self.pending)
|
|
|
|
@property
|
|
def ready(self):
|
|
return len(self.pending) >= CHUNK_BATCH_SIZE
|
|
|
|
def append(self, chunk):
|
|
self.pending.append(chunk)
|
|
|
|
def flush(self):
|
|
chunks = ([self.response] if self.response is not None else []) + self.pending
|
|
if not chunks:
|
|
return None
|
|
if all(isinstance(chunk, AIMessageChunk) for chunk in chunks):
|
|
self.response = add_ai_message_chunks(chunks[0], *chunks[1:])
|
|
else:
|
|
result = chunks[0]
|
|
for chunk in chunks[1:]:
|
|
result = result + chunk
|
|
self.response = result
|
|
self.pending.clear()
|
|
return self.response
|