Files
plyngent/tests/test_agent/test_stream_merge.py
T

26 lines
1.2 KiB
Python

from __future__ import annotations
from plyngent.lmproto.openai_compatible.client import merge_stream_tool_calls
def test_merge_stream_tool_calls_accumulates_arguments() -> None:
lines = [
b'{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"c1","type":"function","function":{"name":"add","arguments":""}}]}}]}',
b'{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\\"a\\":"}}]}}]}',
b'{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"1}"}}]},"finish_reason":"tool_calls"}]}',
]
calls = merge_stream_tool_calls(lines)
assert len(calls) == 1
assert calls[0].id == "c1"
assert calls[0].function.name == "add"
assert calls[0].function.arguments == '{"a":1}'
def test_merge_stream_tool_calls_multiple_indices() -> None:
lines = [
b'{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"a","type":"function","function":{"name":"f1","arguments":"{}"}}]}}]}',
b'{"choices":[{"delta":{"tool_calls":[{"index":1,"id":"b","type":"function","function":{"name":"f2","arguments":"[]"}}]}}]}',
]
calls = merge_stream_tool_calls(lines)
assert {c.function.name for c in calls} == {"f1", "f2"}