from __future__ import annotations import asyncio import json from hashlib import sha256 from pathlib import Path import httpx import numpy as np import pytest from script_build_host.infrastructure.decode_index_builder import ( DashScopeEmbeddingClient, EmbeddingApiError, extract_decode_chunks, rebuild_decode_index, ) def _raw_fixture(root: Path) -> Path: raw = root / "script_decode_raw_data" / "账号甲" raw.mkdir(parents=True) (raw / "post-1.json").write_text( json.dumps( { "帖子ID": "post-1", "段落列表": [ { "id": "p1", "名称": "开头", "主题原子点": [ { "维度名": "对象", "维度名描述": "面向谁", "维度值": "创业者", "维度值描述": "早期团队", } ], "子项": [ { "id": "p1.1", "名称": "子段", "形式原子点": [{"维度": "表达", "原子点": "口播"}], } ], } ], }, ensure_ascii=False, ), encoding="utf-8", ) return raw.parent def _sha256(path: Path) -> str: return sha256(path.read_bytes()).hexdigest() def test_extract_decode_chunks_matches_legacy_fields_and_nested_paragraphs( tmp_path: Path, ) -> None: raw_root = _raw_fixture(tmp_path) chunks = extract_decode_chunks(raw_root) assert len(chunks) == 6 assert [chunk["text_field"] for chunk in chunks] == [ "维度名", "维度名描述", "维度值", "维度值描述", "维度", "原子点", ] assert chunks[-1]["paragraph_id"] == "p1.1" assert chunks[-1]["json_relpath"] == "script_decode_raw_data/账号甲/post-1.json" @pytest.mark.asyncio async def test_rebuild_decode_index_is_complete_resumable_and_atomic(tmp_path: Path) -> None: raw_root = _raw_fixture(tmp_path) source_path = raw_root / "账号甲" / "post-1.json" source_before = _sha256(source_path) index_root = tmp_path / "index" index_root.mkdir() (index_root / "old-marker").write_text("old", encoding="utf-8") request_sizes: list[int] = [] def handler(request: httpx.Request) -> httpx.Response: payload = json.loads(request.content) texts = payload["input"] request_sizes.append(len(texts)) data = [ {"index": index, "embedding": [1.0, float(index + 1), 0.5, 0.25]} for index in reversed(range(len(texts))) ] return httpx.Response( 200, json={"model": "text-embedding-v4", "data": data, "usage": {"prompt_tokens": 7}}, ) async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http: client = DashScopeEmbeddingClient( http, api_key="secret", endpoint="https://dashscope.example/embeddings", model="text-embedding-v4", dimension=4, concurrency=2, ) result = await rebuild_decode_index( raw_root=raw_root, index_root=index_root, client=client, batch_size=2, shard_size=4, ) assert request_sizes == [2, 2, 2] assert not (index_root / "old-marker").exists() assert _sha256(source_path) == source_before assert result.manifest["provider"] == "aliyun-bailian" assert result.manifest["chunk_count"] == 6 assert result.manifest["request_count"] == 3 assert result.manifest["total_tokens"] == 21 assert np.load(index_root / "embeddings.npy").shape == (6, 4) assert len((index_root / "meta.jsonl").read_text(encoding="utf-8").splitlines()) == 6 assert not await asyncio.to_thread(lambda: list(tmp_path.glob(".index-build-*"))) @pytest.mark.asyncio async def test_failed_rebuild_keeps_previously_published_index(tmp_path: Path) -> None: raw_root = _raw_fixture(tmp_path) index_root = tmp_path / "index" index_root.mkdir() (index_root / "old-marker").write_text("old", encoding="utf-8") def handler(_: httpx.Request) -> httpx.Response: return httpx.Response(400, json={"message": "bad request"}) async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http: client = DashScopeEmbeddingClient( http, api_key="secret", endpoint="https://dashscope.example/embeddings", model="text-embedding-v4", dimension=4, max_retries=0, ) with pytest.raises(EmbeddingApiError, match="HTTP 400"): await rebuild_decode_index( raw_root=raw_root, index_root=index_root, client=client, batch_size=2, shard_size=4, ) assert (index_root / "old-marker").read_text(encoding="utf-8") == "old" assert await asyncio.to_thread(lambda: list(tmp_path.glob(".index-build-*")))