75 lines
1.8 KiB
Python
75 lines
1.8 KiB
Python
from __future__ import annotations
|
|
|
|
from app.services.table_profiling import _parse_completion_payload
|
|
from app.utils.llm_usage import extract_usage
|
|
|
|
|
|
def test_parse_completion_payload_handles_array_with_trailing_text() -> None:
|
|
response_payload = {
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": """
|
|
结果如下：
|
|
[
|
|
{"id": "snpt_a"},
|
|
{"id": "snpt_b"}
|
|
]
|
|
附加说明：模型可能会输出额外文本。
|
|
""".strip()
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
parsed = _parse_completion_payload(response_payload)
|
|
|
|
assert isinstance(parsed, list)
|
|
assert [item["id"] for item in parsed] == ["snpt_a", "snpt_b"]
|
|
|
|
|
|
def test_extract_usage_info_normalizes_numeric_fields() -> None:
|
|
response_payload = {
|
|
"raw": {
|
|
"usage": {
|
|
"prompt_tokens": 12.7,
|
|
"completion_tokens": 3,
|
|
"total_tokens": 15.7,
|
|
"prompt_tokens_details": {"cached_tokens": 8.9, "other": None},
|
|
"non_numeric": "ignored",
|
|
}
|
|
}
|
|
}
|
|
|
|
usage = extract_usage(response_payload)
|
|
|
|
assert usage == {
|
|
"prompt_tokens": 12,
|
|
"completion_tokens": 3,
|
|
"total_tokens": 15,
|
|
"prompt_tokens_details": {"cached_tokens": 8},
|
|
}
|
|
|
|
|
|
def test_extract_usage_handles_alias_keys() -> None:
|
|
response_payload = {
|
|
"raw": {
|
|
"usageMetadata": {
|
|
"input_tokens": 20,
|
|
"output_tokens": 4,
|
|
}
|
|
}
|
|
}
|
|
|
|
usage = extract_usage(response_payload)
|
|
|
|
assert usage == {
|
|
"prompt_tokens": 20,
|
|
"completion_tokens": 4,
|
|
"total_tokens": 24,
|
|
}
|
|
|
|
|
|
def test_extract_usage_returns_none_when_missing() -> None:
|
|
assert extract_usage({"raw": {}}) is None
|