数据知识回调入库

test/test_table_profiling_parsing.py

@@ -0,0 +1,74 @@
+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