File size: 3,186 Bytes
95f8c4b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | import json
from pathlib import Path
from aistudio_api.domain.models import parse_chunk_usage, parse_response_chunk, parse_text_output
from aistudio_api.infrastructure.gateway.stream_parser import IncrementalJSONStreamParser, classify_chunk
ROOT = Path(__file__).resolve().parents[2]
FIXTURES = ROOT / "tests"
def test_parse_text_output_from_stream_bundle():
raw = (FIXTURES / "test_output.json").read_text()
output = parse_text_output(raw)
assert output.text == "你好！有什么我可以帮你的吗？"
assert output.thinking.startswith('The user said "你好"')
assert "Option 3 (Friendly)" in output.thinking
assert output.usage["prompt_tokens"] == 5
assert output.usage["completion_tokens"] == 161
assert output.usage["total_tokens"] == 166
assert output.usage["completion_tokens_details"]["reasoning_tokens"] == 153
assert output.response_id
assert output.candidates[0].finish_reason == 1
def test_parse_response_chunk_and_classify_chunk():
raw = json.loads((FIXTURES / "test_output.json").read_text())
final_chunk = raw[0][-1]
candidate = parse_response_chunk(final_chunk)
usage = parse_chunk_usage(final_chunk)
assert candidate.finish_reason == 1
assert candidate.safety_ratings
assert usage["prompt_tokens"] == 5
assert usage["completion_tokens"] == 161
assert usage["completion_tokens_details"]["reasoning_tokens"] == 153
ctype, text = classify_chunk(raw[0][1])
assert ctype == "thinking"
assert 'standard Chinese greeting meaning "Hello."' in text
def test_stream_parser_extracts_real_chunks():
raw = (FIXTURES / "test_output.json").read_text()
parser = IncrementalJSONStreamParser()
chunks = list(parser.feed(raw))
assert len(chunks) == 10
assert classify_chunk(chunks[2])[0] == "thinking"
assert classify_chunk(chunks[7])[0] == "body"
def test_parse_text_output_handles_double_wrapped_chunk_bundle():
raw = json.dumps(
[
[
[
[[[[[None, "先想一会儿", None, None, None, None, None, None, None, None, None, None, 1]], "model"]]],
None,
[10, None, 10, None, [[1, 10]]],
None,
None,
None,
None,
"resp_test",
],
[
[[[[[None, "答案", None, None, None, None, None, None, None, None, None, None, 0]], "model"]]],
None,
[10, 2, 17, None, [[1, 10]], None, None, None, None, 5],
None,
None,
None,
None,
"resp_test",
],
]
]
)
output = parse_text_output(raw)
assert output.thinking == "先想一会儿"
assert output.text == "答案"
assert output.usage["prompt_tokens"] == 10
assert output.usage["completion_tokens"] == 7
assert output.usage["total_tokens"] == 17
assert output.usage["completion_tokens_details"]["reasoning_tokens"] == 5
assert output.response_id == "resp_test"
|