feat(data-process): 支持单项生成五十条数据

backend/tests/test_data_process_generation.py

@@ -3,8 +3,13 @@ from __future__ import annotations
 import json
 
 import httpx
+import pytest
 
-from app.modules.data_process.generation import chat_completions_url, generate_model_records
+from app.modules.data_process.generation import (
+    ModelGenerationError,
+    chat_completions_url,
+    generate_model_records,
+)
 
 
 def test_chat_completions_url_accepts_host_base_and_complete_url() -> None:
@@ -100,3 +105,165 @@ def test_generate_model_records_keeps_partial_failure_for_manual_repair() -> Non
     assert len(records) == 1
     assert records[0]["status"] == "invalid"
     assert records[0]["error"]
+
+
+def test_generate_model_records_batches_fifty_results_with_unique_ids() -> None:
+    requests: list[httpx.Request] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        requests.append(request)
+        batch_start = (len(requests) - 1) * 10 + 1
+        batch_end = batch_start + 9
+        payload = json.loads(request.content)
+        system_prompt = payload["messages"][0]["content"]
+        assert "items 必须包含 10 条" in system_prompt
+        assert f"第 {batch_start}-{batch_end} 条" in system_prompt
+        return httpx.Response(
+            200,
+            json={
+                "choices": [
+                    {
+                        "message": {
+                            "content": json.dumps(
+                                {
+                                    "items": [
+                                        {
+                                            "instruction": "同一问题",
+                                            "input": "来源正文",
+                                            "output": "同一答案",
+                                        }
+                                        for _ in range(batch_start, batch_end + 1)
+                                    ]
+                                },
+                                ensure_ascii=False,
+                            )
+                        }
+                    }
+                ]
+            },
+        )
+
+    records = generate_model_records(
+        [{"id": "preview-50", "edited_content": "来源正文"}],
+        model={"name": "model", "api_url": "https://model.example/v1"},
+        config={},
+        task_id="task-50",
+        split={"train": 100, "validation": 0, "test": 0},
+        qa_pairs_per_item=50,
+        client=httpx.Client(transport=httpx.MockTransport(handler)),
+    )
+
+    assert len(requests) == 5
+    assert len(records) == 50
+    assert len({record["id"] for record in records}) == 50
+    assert {record["instruction"] for record in records} == {"同一问题"}
+    assert all(record["status"] == "valid" for record in records)
+
+
+def test_generate_model_records_preserves_successful_batches_when_one_fails() -> None:
+    request_count = 0
+
+    def handler(_: httpx.Request) -> httpx.Response:
+        nonlocal request_count
+        request_count += 1
+        if request_count == 2:
+            return httpx.Response(500)
+        return httpx.Response(
+            200,
+            json={
+                "choices": [
+                    {
+                        "message": {
+                            "content": json.dumps(
+                                {
+                                    "items": [
+                                        {
+                                            "instruction": f"问题 {index}",
+                                            "output": f"答案 {index}",
+                                        }
+                                        for index in range(1, 11)
+                                    ]
+                                },
+                                ensure_ascii=False,
+                            )
+                        }
+                    }
+                ]
+            },
+        )
+
+    records = generate_model_records(
+        [{"id": "preview-partial", "edited_content": "来源正文"}],
+        model={"name": "model", "api_url": "https://model.example/v1"},
+        config={"generation_retries": 0},
+        task_id="task-partial",
+        split={"train": 100, "validation": 0, "test": 0},
+        qa_pairs_per_item=20,
+        client=httpx.Client(transport=httpx.MockTransport(handler)),
+    )
+
+    assert len(records) == 11
+    assert sum(record["status"] == "valid" for record in records) == 10
+    failed = next(record for record in records if record["status"] == "invalid")
+    assert "第 11-20 条" in failed["instruction"]
+    assert len({record["id"] for record in records}) == len(records)
+
+
+def test_generate_model_records_retries_short_batch_then_marks_it_invalid() -> None:
+    request_count = 0
+
+    def handler(_: httpx.Request) -> httpx.Response:
+        nonlocal request_count
+        request_count += 1
+        return httpx.Response(
+            200,
+            json={
+                "choices": [
+                    {
+                        "message": {
+                            "content": json.dumps(
+                                {
+                                    "items": [
+                                        {
+                                            "instruction": "只有一条",
+                                            "output": "不足本批要求数量",
+                                        }
+                                    ]
+                                },
+                                ensure_ascii=False,
+                            )
+                        }
+                    }
+                ]
+            },
+        )
+
+    records = generate_model_records(
+        [{"id": "preview-short", "edited_content": "来源正文"}],
+        model={"name": "model", "api_url": "https://model.example/v1"},
+        config={"generation_retries": 1},
+        task_id="task-short",
+        split={"train": 100, "validation": 0, "test": 0},
+        qa_pairs_per_item=10,
+        client=httpx.Client(transport=httpx.MockTransport(handler)),
+    )
+
+    assert request_count == 2
+    assert len(records) == 1
+    assert records[0]["status"] == "invalid"
+    assert "expected 10, got 1" in records[0]["error"]
+
+
+@pytest.mark.parametrize("qa_pairs_per_item", [0, 51])
+def test_generate_model_records_rejects_out_of_range_count(
+    qa_pairs_per_item: int,
+) -> None:
+    with pytest.raises(ModelGenerationError, match=r"\[1, 50\]"):
+        generate_model_records(
+            [],
+            model={"name": "model", "api_url": "https://model.example/v1"},
+            config={},
+            task_id="task-invalid",
+            split={"train": 100, "validation": 0, "test": 0},
+            qa_pairs_per_item=qa_pairs_per_item,
+        )