149 lines
4.5 KiB
Python
149 lines
4.5 KiB
Python
"""U3 / G8 delta_flush_interval 调速测试。
|
|
|
|
覆盖 R11-R12, R14:
|
|
- R11 chunk 按 flush_interval_ms 间隔批量 yield
|
|
- R12 配置化(flush_interval_ms=0 退化为逐 chunk yield)
|
|
- R14 自检:合并 content 等于原始 chunks 拼接(不丢字符)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
|
|
|
|
from agentkit.core.react import ReActEngine
|
|
from agentkit.llm.protocol import StreamChunk
|
|
|
|
|
|
class _StubGateway:
|
|
"""模拟 LLMGateway,yield 一串 StreamChunk 后结束。"""
|
|
|
|
def __init__(self, chunks: list[str]):
|
|
self._chunks = chunks
|
|
|
|
def get_provider_name_for_model(self, model: str) -> str | None:
|
|
return None
|
|
|
|
async def chat_stream(self, **kwargs):
|
|
for c in self._chunks:
|
|
yield StreamChunk(content=c, model="test")
|
|
|
|
|
|
def _collect_token_events(events) -> list[str]:
|
|
return [e.data["content"] for e in events if e.event_type == "token"]
|
|
|
|
|
|
# ---- R12 Config: flush_interval_ms=0 → 逐 chunk yield(向后兼容) ----
|
|
|
|
|
|
async def test_flush_interval_zero_yields_per_chunk():
|
|
chunks = ["H", "e", "l", "l", "o"]
|
|
gw = _StubGateway(chunks)
|
|
engine = ReActEngine(llm_gateway=gw, flush_interval_ms=0)
|
|
events = []
|
|
async for ev in engine.execute_stream(
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
tools=[],
|
|
model="test",
|
|
):
|
|
events.append(ev)
|
|
tokens = _collect_token_events(events)
|
|
# 5 chunks → 5 token events(每个内容 = 单 chunk)
|
|
assert tokens == ["H", "e", "l", "l", "o"]
|
|
|
|
|
|
# ---- R11 Happy path: flush_interval_ms > 0 → 批量合并 ----
|
|
|
|
|
|
async def test_flush_interval_batches_chunks_by_interval():
|
|
chunks = ["a", "b", "c", "d", "e", "f"]
|
|
gw = _StubGateway(chunks)
|
|
# 间隔设很大(10s),所有 chunks 在第一个 interval 内累积,流结束后最终 flush
|
|
engine = ReActEngine(llm_gateway=gw, flush_interval_ms=10000)
|
|
events = []
|
|
async for ev in engine.execute_stream(
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
tools=[],
|
|
model="test",
|
|
):
|
|
events.append(ev)
|
|
tokens = _collect_token_events(events)
|
|
# 所有 chunk 累积到流结束,最终 flush 一次 → 1 个 token event,content = 全拼接
|
|
assert len(tokens) == 1
|
|
assert tokens[0] == "abcdef"
|
|
|
|
|
|
# ---- R14 Self-check: 合并 content 等于原始 chunks 拼接(不丢字符) ----
|
|
|
|
|
|
async def test_no_character_loss_after_merge():
|
|
chunks = ["Hello", " ", "World", "!", "你好", "世界"]
|
|
gw = _StubGateway(chunks)
|
|
engine = ReActEngine(llm_gateway=gw, flush_interval_ms=10000)
|
|
events = []
|
|
async for ev in engine.execute_stream(
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
tools=[],
|
|
model="test",
|
|
):
|
|
events.append(ev)
|
|
tokens = _collect_token_events(events)
|
|
# 合并所有 token events 的 content 等于原始 chunks 拼接
|
|
merged = "".join(tokens)
|
|
assert merged == "".join(chunks) == "Hello World!你好世界"
|
|
|
|
|
|
# ---- Edge: 流结束 mid-interval → 最终 flush 剩余 buffer ----
|
|
|
|
|
|
async def test_final_flush_on_stream_end():
|
|
chunks = ["x", "y", "z"]
|
|
gw = _StubGateway(chunks)
|
|
engine = ReActEngine(llm_gateway=gw, flush_interval_ms=10000)
|
|
events = []
|
|
async for ev in engine.execute_stream(
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
tools=[],
|
|
model="test",
|
|
):
|
|
events.append(ev)
|
|
tokens = _collect_token_events(events)
|
|
# mid-interval 累积 → 流结束最终 flush 一次
|
|
assert tokens == ["xyz"]
|
|
|
|
|
|
# ---- Edge: 单个 chunk 后流结束 → 立即 flush ----
|
|
|
|
|
|
async def test_single_chunk_immediate_flush():
|
|
chunks = ["only"]
|
|
gw = _StubGateway(chunks)
|
|
engine = ReActEngine(llm_gateway=gw, flush_interval_ms=10000)
|
|
events = []
|
|
async for ev in engine.execute_stream(
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
tools=[],
|
|
model="test",
|
|
):
|
|
events.append(ev)
|
|
tokens = _collect_token_events(events)
|
|
assert tokens == ["only"]
|
|
|
|
|
|
# ---- Edge: chunks 含空 content(usage-only chunk)不进 buffer ----
|
|
|
|
|
|
async def test_empty_content_chunk_not_buffered():
|
|
chunks = ["a", "", "b"] # 中间 chunk 空
|
|
gw = _StubGateway(chunks)
|
|
engine = ReActEngine(llm_gateway=gw, flush_interval_ms=10000)
|
|
events = []
|
|
async for ev in engine.execute_stream(
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
tools=[],
|
|
model="test",
|
|
):
|
|
events.append(ev)
|
|
tokens = _collect_token_events(events)
|
|
# 空 chunk 跳过 buffer,最终 flush "ab"
|
|
assert tokens == ["ab"]
|