84 lines
3.0 KiB
Python
84 lines
3.0 KiB
Python
import itertools
|
|
import json
|
|
|
|
from langchain_core.messages import AIMessageChunk
|
|
import langchain_core.messages.ai as ai
|
|
import pytest
|
|
|
|
from artemis.llm.chunk_accumulator import ChunkAccumulator
|
|
from artemis.llm.response_guard import ResponseIntegrityError, ensure_response_complete
|
|
|
|
|
|
def test_batching_preserves_interleaved_tools_unicode_metadata_and_usage():
|
|
arguments = [json.dumps({"target": "中文🙂" * 80}), json.dumps({"target": 3})]
|
|
chunks = [AIMessageChunk(content="Live text", id="resp_one")]
|
|
for index in range(2):
|
|
chunks.append(
|
|
AIMessageChunk(
|
|
content=[],
|
|
tool_call_chunks=[
|
|
{"name": "click", "args": "", "id": f"call_{index}", "index": index}
|
|
],
|
|
)
|
|
)
|
|
for pair in itertools.zip_longest(*arguments):
|
|
for index, value in enumerate(pair):
|
|
if value is not None:
|
|
chunks.append(
|
|
AIMessageChunk(content=[], tool_call_chunks=[{"args": value, "index": index}])
|
|
)
|
|
chunks.append(
|
|
AIMessageChunk(
|
|
content=[],
|
|
response_metadata={"id": "resp_one", "status": "completed"},
|
|
usage_metadata={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
|
|
chunk_position="last",
|
|
)
|
|
)
|
|
expected = chunks[0]
|
|
accumulator = ChunkAccumulator()
|
|
max_pending = 0
|
|
for index, chunk in enumerate(chunks):
|
|
if index:
|
|
expected = expected + chunk
|
|
accumulator.append(chunk)
|
|
if accumulator.ready:
|
|
accumulator.flush()
|
|
max_pending = max(max_pending, len(accumulator.pending))
|
|
result = accumulator.flush()
|
|
assert result.model_dump() == expected.model_dump()
|
|
assert [t["args"] for t in result.tool_calls] == [json.loads(a) for a in arguments]
|
|
assert result.usage_metadata["total_tokens"] == 15
|
|
assert max_pending < 256
|
|
assert not accumulator.pending
|
|
|
|
|
|
def test_partial_json_is_not_reparsed_for_every_single_character(monkeypatch):
|
|
chunks = [
|
|
AIMessageChunk(
|
|
content=[], tool_call_chunks=[{"name": "click", "args": '{"target":x', "index": 0}]
|
|
)
|
|
]
|
|
chunks.extend(
|
|
AIMessageChunk(content=[], tool_call_chunks=[{"args": "x", "index": 0}]) for _ in range(800)
|
|
)
|
|
original = ai.parse_partial_json
|
|
parsed_lengths = []
|
|
|
|
def counted(arguments, *args, **kwargs):
|
|
parsed_lengths.append(len(arguments))
|
|
return original(arguments, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(ai, "parse_partial_json", counted)
|
|
accumulator = ChunkAccumulator()
|
|
for chunk in chunks:
|
|
accumulator.append(chunk)
|
|
if accumulator.ready:
|
|
accumulator.flush()
|
|
result = accumulator.flush()
|
|
assert result.tool_call_chunks[0]["args"] == '{"target":x' + "x" * 800
|
|
assert len(parsed_lengths) < 20
|
|
result.response_metadata["status"] = "completed"
|
|
with pytest.raises(ResponseIntegrityError):
|
|
ensure_response_complete(result)
|