| 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" |
|
|
