395 lines
14 KiB
Python
395 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
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:
|
|
assert chat_completions_url("www.caoxiaozhu.com") == (
|
|
"https://www.caoxiaozhu.com/v1/chat/completions"
|
|
)
|
|
assert chat_completions_url("https://model.example/v1") == (
|
|
"https://model.example/v1/chat/completions"
|
|
)
|
|
complete = "https://model.example/openai/v1/chat/completions"
|
|
assert chat_completions_url(complete) == complete
|
|
|
|
|
|
def test_generate_model_records_uses_prompt_auth_and_stable_split() -> None:
|
|
requests: list[httpx.Request] = []
|
|
progress_updates: list[tuple[int, int]] = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
requests.append(request)
|
|
payload = json.loads(request.content)
|
|
assert payload["model"] == "qwen-plus"
|
|
assert payload["response_format"] == {"type": "json_object"}
|
|
assert "客户反馈页面加载慢" in payload["messages"][1]["content"]
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": json.dumps(
|
|
{
|
|
"items": [
|
|
{
|
|
"instruction": "请生成简洁客服回复",
|
|
"input": "客户反馈页面加载慢",
|
|
"output": "已收到反馈，我们正在排查。",
|
|
}
|
|
]
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
}
|
|
}
|
|
]
|
|
},
|
|
)
|
|
|
|
client = httpx.Client(transport=httpx.MockTransport(handler))
|
|
records = generate_model_records(
|
|
[{"id": "preview-1", "edited_content": "客户反馈页面加载慢"}],
|
|
model={
|
|
"name": "Qwen",
|
|
"online_model_name": "qwen-plus",
|
|
"api_url": "model.example",
|
|
"api_key": "test-secret",
|
|
},
|
|
config={
|
|
"generation_prompt": "请处理：{{ content }}",
|
|
"json_mode": True,
|
|
"temperature": 0.2,
|
|
"max_tokens": 512,
|
|
},
|
|
task_id="task-1",
|
|
split={"train": 100, "validation": 0, "test": 0},
|
|
qa_pairs_per_item=1,
|
|
client=client,
|
|
on_progress=lambda processed, total: progress_updates.append((processed, total)),
|
|
)
|
|
|
|
assert len(records) == 1
|
|
assert records[0]["status"] == "valid"
|
|
assert records[0]["split"] == "train"
|
|
assert requests[0].headers["Authorization"] == "Bearer test-secret"
|
|
assert progress_updates == [(1, 1)]
|
|
|
|
|
|
def test_generate_model_records_builds_reasoning_output_with_think_tags() -> None:
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
payload = json.loads(request.content)
|
|
system_prompt = payload["messages"][0]["content"]
|
|
assert '"reasoning":"...","answer":"..."' in system_prompt
|
|
assert "系统会在保存时统一组装" in system_prompt
|
|
content = "<think>模型接口自己的分析</think>" + json.dumps(
|
|
{
|
|
"items": [
|
|
{
|
|
"instruction": "计算两项费用合计",
|
|
"input": "交通费 30 元，餐费 20 元",
|
|
"reasoning": "先识别两项费用，再计算 30 + 20。",
|
|
"answer": "合计 50 元。",
|
|
}
|
|
]
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
return httpx.Response(
|
|
200,
|
|
json={"choices": [{"message": {"content": content}}]},
|
|
)
|
|
|
|
records = generate_model_records(
|
|
[{"id": "preview-reasoning", "edited_content": "交通费 30 元，餐费 20 元"}],
|
|
model={"name": "model", "api_url": "https://model.example/v1"},
|
|
config={"output_type": "reasoning"},
|
|
task_id="task-reasoning",
|
|
split={"train": 100, "validation": 0, "test": 0},
|
|
qa_pairs_per_item=1,
|
|
client=httpx.Client(transport=httpx.MockTransport(handler)),
|
|
)
|
|
|
|
assert records[0]["status"] == "valid"
|
|
assert records[0]["output"] == (
|
|
"<think>\n先识别两项费用,再计算 30 + 20。\n</think>\n合计 50 元。"
|
|
)
|
|
|
|
|
|
def test_generate_model_records_marks_reasoning_without_reasoning_field_invalid() -> None:
|
|
response = {
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": json.dumps(
|
|
{
|
|
"items": [
|
|
{
|
|
"instruction": "问题",
|
|
"answer": "只有最终答案",
|
|
}
|
|
]
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
}
|
|
}
|
|
]
|
|
}
|
|
client = httpx.Client(
|
|
transport=httpx.MockTransport(lambda _: httpx.Response(200, json=response))
|
|
)
|
|
|
|
records = generate_model_records(
|
|
[{"id": "preview-missing-reasoning", "edited_content": "来源正文"}],
|
|
model={"name": "model", "api_url": "https://model.example/v1"},
|
|
config={"output_type": "reasoning"},
|
|
task_id="task-missing-reasoning",
|
|
split={"train": 100, "validation": 0, "test": 0},
|
|
qa_pairs_per_item=1,
|
|
client=client,
|
|
)
|
|
|
|
assert records[0]["status"] == "invalid"
|
|
assert records[0]["output"] == "只有最终答案"
|
|
assert "reasoning" in records[0]["error"]
|
|
|
|
|
|
def test_standard_output_removes_model_think_block() -> None:
|
|
content = json.dumps(
|
|
{
|
|
"items": [
|
|
{
|
|
"instruction": "问题",
|
|
"output": "<think>不应保存的分析</think>最终答案",
|
|
}
|
|
]
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
client = httpx.Client(
|
|
transport=httpx.MockTransport(
|
|
lambda _: httpx.Response(
|
|
200,
|
|
json={"choices": [{"message": {"content": content}}]},
|
|
)
|
|
)
|
|
)
|
|
|
|
records = generate_model_records(
|
|
[{"id": "preview-standard", "edited_content": "来源正文"}],
|
|
model={"name": "model", "api_url": "https://model.example/v1"},
|
|
config={"output_type": "standard"},
|
|
task_id="task-standard",
|
|
split={"train": 100, "validation": 0, "test": 0},
|
|
qa_pairs_per_item=1,
|
|
client=client,
|
|
)
|
|
|
|
assert records[0]["output"] == "最终答案"
|
|
|
|
|
|
def test_generate_model_records_keeps_partial_failure_for_manual_repair() -> None:
|
|
client = httpx.Client(
|
|
transport=httpx.MockTransport(
|
|
lambda _: httpx.Response(200, json={"choices": [{"message": {"content": "not-json"}}]})
|
|
)
|
|
)
|
|
records = generate_model_records(
|
|
[{"id": "preview-1", "edited_content": "来源正文"}],
|
|
model={"name": "model", "api_url": "https://model.example/v1"},
|
|
config={"generation_retries": 1},
|
|
task_id="task-1",
|
|
split={"train": 80, "validation": 10, "test": 10},
|
|
qa_pairs_per_item=1,
|
|
client=client,
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
|
|
def test_generate_model_records_rejects_unknown_output_type() -> None:
|
|
with pytest.raises(ModelGenerationError, match="output_type"):
|
|
generate_model_records(
|
|
[],
|
|
model={"name": "model", "api_url": "https://model.example/v1"},
|
|
config={"output_type": "unknown"},
|
|
task_id="task-invalid-output",
|
|
split={"train": 100, "validation": 0, "test": 0},
|
|
qa_pairs_per_item=1,
|
|
)
|