chore: track claude skills, tools, templates, reference code and research-wiki
- Add all claude skills (brainstorming, commit, debugging, TDD, etc.) - Add claude hooks (pre-commit-guard, post-edit-quality) - Add research templates (experiment plan, research brief, etc.) - Add claude tools (arxiv/semantic_scholar/openalex fetch, wiki, exa) - Add TRM4 reference implementation as algorithm fidelity baseline - Add research-wiki content (plans, index, graph, query_pack) - Update .gitignore to exclude .graphify_version runtime state
This commit is contained in:
@@ -0,0 +1,89 @@
|
||||
"""
|
||||
全局测试配置
|
||||
============
|
||||
- 代理修复: 从 .env 读取 EMBED_API_URL,将其中的 CGNAT 段 IP 加入 no_proxy。
|
||||
- real_config fixture: 从真实 config/default.yaml + .env 加载 Config,供远程测试复用。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ipaddress
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from dotenv import dotenv_values
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 代理修复: 将内网 API 地址加入 no_proxy(模块加载时立即执行)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_CGNAT_NETWORK = ipaddress.IPv4Network("100.64.0.0/10")
|
||||
|
||||
|
||||
def _fix_no_proxy() -> None:
|
||||
"""从 .env 读取 API URL,将 CGNAT 段 IP 加入 no_proxy / NO_PROXY。
|
||||
|
||||
httpx/urllib 的 no_proxy 不支持 CIDR 记法,必须逐个添加具体 IP。
|
||||
"""
|
||||
if not (os.environ.get("http_proxy") or os.environ.get("HTTP_PROXY")):
|
||||
return
|
||||
|
||||
project_root = Path(__file__).resolve().parent.parent
|
||||
env_file = project_root / ".env"
|
||||
if not env_file.exists():
|
||||
return
|
||||
|
||||
env_vars = dotenv_values(str(env_file))
|
||||
|
||||
# 收集所有需要绕过代理的 IP
|
||||
hosts_to_add: list[str] = []
|
||||
for url_key in ("EMBED_API_URL", "LLM_API_URL", "VLM_API_URL"):
|
||||
url = env_vars.get(url_key, "")
|
||||
if not url:
|
||||
continue
|
||||
match = re.match(r"https?://([^:/]+)", url)
|
||||
if not match:
|
||||
continue
|
||||
host = match.group(1)
|
||||
try:
|
||||
if ipaddress.IPv4Address(host) in _CGNAT_NETWORK:
|
||||
hosts_to_add.append(host)
|
||||
except (ipaddress.AddressValueError, ValueError):
|
||||
pass
|
||||
|
||||
if not hosts_to_add:
|
||||
return
|
||||
|
||||
for var in ("no_proxy", "NO_PROXY"):
|
||||
current = os.environ.get(var, "")
|
||||
for host in hosts_to_add:
|
||||
if host not in current:
|
||||
separator = "," if current else ""
|
||||
current = f"{current}{separator}{host}"
|
||||
os.environ[var] = current
|
||||
|
||||
|
||||
# 模块加载时执行代理修复
|
||||
_fix_no_proxy()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def real_config():
|
||||
"""从真实 config/default.yaml + .env 加载 Config(session 级缓存)。
|
||||
|
||||
用于远程 API 测试,需要 .env 中配置有效的 API 密钥。
|
||||
"""
|
||||
from video_tree_trm.config import Config
|
||||
|
||||
project_root = Path(__file__).resolve().parent.parent
|
||||
yaml_path = project_root / "config" / "default.yaml"
|
||||
env_path = project_root / ".env"
|
||||
|
||||
return Config.load(str(yaml_path), env_path=str(env_path))
|
||||
@@ -0,0 +1,276 @@
|
||||
"""
|
||||
配置管理模块单元测试
|
||||
====================
|
||||
覆盖: YAML 加载、.env 覆盖、CLI 覆盖、缺字段报错、优先级验证、embed_dim 一致性。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
from video_tree_trm.config import Config, TreeConfig
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_FULL_YAML = {
|
||||
"tree": {
|
||||
"max_paragraphs_per_l2": 5,
|
||||
"l1_segment_duration": 600.0,
|
||||
"l2_clip_duration": 20.0,
|
||||
"l3_fps": 1.0,
|
||||
"l2_representative_frames": 3,
|
||||
"cache_dir": "cache/trees",
|
||||
},
|
||||
"embed": {
|
||||
"backend": "local",
|
||||
"model_name": "test-model",
|
||||
"embed_dim": 2560,
|
||||
"device": "cpu",
|
||||
"api_key": "",
|
||||
"api_url": "",
|
||||
},
|
||||
"llm": {
|
||||
"backend": "qwen",
|
||||
"api_key": "yaml-llm-key",
|
||||
"model": "qwen-plus",
|
||||
"api_url": "https://example.com/llm",
|
||||
"max_tokens": 256,
|
||||
"temperature": 0.1,
|
||||
},
|
||||
"vlm": {
|
||||
"backend": "qwen",
|
||||
"api_key": "yaml-vlm-key",
|
||||
"model": "qwen-vl-plus",
|
||||
"api_url": "https://example.com/vlm",
|
||||
"max_tokens": 256,
|
||||
"temperature": 0.1,
|
||||
},
|
||||
"retriever": {
|
||||
"embed_dim": 2560,
|
||||
"num_heads": 4,
|
||||
"L_layers": 2,
|
||||
"L_cycles": 4,
|
||||
"max_rounds": 5,
|
||||
"ffn_expansion": 2.0,
|
||||
"checkpoint": None,
|
||||
},
|
||||
"train": {
|
||||
"lr": 1e-4,
|
||||
"weight_decay": 1e-5,
|
||||
"batch_size": 1,
|
||||
"max_epochs_phase1": 30,
|
||||
"max_epochs_phase2": 20,
|
||||
"nav_loss_weight": 1.0,
|
||||
"act_loss_weight": 0.1,
|
||||
"act_lambda_step": 0.1,
|
||||
"act_gamma": 0.9,
|
||||
"eval_interval": 5,
|
||||
"save_dir": "checkpoints",
|
||||
"dataset": "longbench",
|
||||
"dataset_path": "data/longbench",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def yaml_path(tmp_path: Path) -> Path:
|
||||
"""创建完整配置的临时 YAML 文件。"""
|
||||
p = tmp_path / "config" / "default.yaml"
|
||||
p.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(p, "w", encoding="utf-8") as f:
|
||||
yaml.dump(_FULL_YAML, f, allow_unicode=True)
|
||||
return p
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def env_path(tmp_path: Path) -> Path:
|
||||
"""创建临时 .env 文件。"""
|
||||
p = tmp_path / ".env"
|
||||
p.write_text(
|
||||
"LLM_API_KEY=env-llm-key\n"
|
||||
"LLM_MODEL=env-llm-model\n"
|
||||
"LLM_API_URL=https://env.example.com/llm\n"
|
||||
"VLM_API_KEY=env-vlm-key\n"
|
||||
"VLM_MODEL=env-vlm-model\n"
|
||||
"VLM_API_URL=https://env.example.com/vlm\n"
|
||||
"EMBED_BACKEND=remote\n"
|
||||
"EMBED_MODEL=env-embed-model\n"
|
||||
"EMBED_API_KEY=env-embed-key\n"
|
||||
"EMBED_API_URL=https://env.example.com/embed\n"
|
||||
)
|
||||
return p
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试: YAML 加载
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestYAMLLoad:
|
||||
"""YAML 基础加载测试。"""
|
||||
|
||||
def test_load_full_yaml(self, yaml_path: Path) -> None:
|
||||
"""完整 YAML 应成功加载所有字段。"""
|
||||
cfg = Config.load(
|
||||
str(yaml_path), env_path=str(yaml_path.parent / ".env.nonexist")
|
||||
)
|
||||
assert isinstance(cfg.tree, TreeConfig)
|
||||
assert cfg.tree.max_paragraphs_per_l2 == 5
|
||||
assert cfg.tree.l1_segment_duration == 600.0
|
||||
assert cfg.embed.embed_dim == 2560
|
||||
assert cfg.retriever.checkpoint is None
|
||||
assert cfg.train.dataset == "longbench"
|
||||
|
||||
def test_file_not_found(self, tmp_path: Path) -> None:
|
||||
"""不存在的 YAML 应抛出 FileNotFoundError。"""
|
||||
with pytest.raises(FileNotFoundError, match="配置文件不存在"):
|
||||
Config.load(str(tmp_path / "nonexist.yaml"))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试: 缺字段报错
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMissingField:
|
||||
"""缺少必需字段时应抛出 TypeError。"""
|
||||
|
||||
def test_missing_tree_field(self, tmp_path: Path) -> None:
|
||||
"""tree 节缺少字段应报 TypeError。"""
|
||||
bad_yaml = _FULL_YAML.copy()
|
||||
bad_yaml = {
|
||||
k: (v.copy() if isinstance(v, dict) else v) for k, v in _FULL_YAML.items()
|
||||
}
|
||||
del bad_yaml["tree"]["cache_dir"]
|
||||
|
||||
p = tmp_path / "bad.yaml"
|
||||
with open(p, "w") as f:
|
||||
yaml.dump(bad_yaml, f)
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
Config.load(str(p), env_path=str(tmp_path / ".env.nonexist"))
|
||||
|
||||
def test_missing_section(self, tmp_path: Path) -> None:
|
||||
"""缺少整个配置节应报 TypeError。"""
|
||||
bad_yaml = {k: v for k, v in _FULL_YAML.items() if k != "train"}
|
||||
|
||||
p = tmp_path / "bad2.yaml"
|
||||
with open(p, "w") as f:
|
||||
yaml.dump(bad_yaml, f)
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
Config.load(str(p), env_path=str(tmp_path / ".env.nonexist"))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试: .env 覆盖
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEnvOverride:
|
||||
""".env 文件应覆盖 YAML 中的 api_key。"""
|
||||
|
||||
def test_env_overrides_api_keys(self, yaml_path: Path, env_path: Path) -> None:
|
||||
"""api_key/model/api_url 应优先使用 .env 中的值。"""
|
||||
cfg = Config.load(str(yaml_path), env_path=str(env_path))
|
||||
assert cfg.llm.api_key == "env-llm-key"
|
||||
assert cfg.llm.model == "env-llm-model"
|
||||
assert cfg.llm.api_url == "https://env.example.com/llm"
|
||||
assert cfg.vlm.api_key == "env-vlm-key"
|
||||
assert cfg.vlm.model == "env-vlm-model"
|
||||
assert cfg.vlm.api_url == "https://env.example.com/vlm"
|
||||
|
||||
def test_env_overrides_embed(self, yaml_path: Path, env_path: Path) -> None:
|
||||
"""embed 相关字段应优先使用 .env 中的值。"""
|
||||
cfg = Config.load(str(yaml_path), env_path=str(env_path))
|
||||
assert cfg.embed.backend == "remote"
|
||||
assert cfg.embed.model_name == "env-embed-model"
|
||||
assert cfg.embed.api_key == "env-embed-key"
|
||||
assert cfg.embed.api_url == "https://env.example.com/embed"
|
||||
|
||||
def test_yaml_fallback_when_no_env(self, yaml_path: Path) -> None:
|
||||
"""无 .env 时应使用 YAML 中的值。"""
|
||||
cfg = Config.load(
|
||||
str(yaml_path), env_path=str(yaml_path.parent / ".env.nonexist")
|
||||
)
|
||||
assert cfg.llm.api_key == "yaml-llm-key"
|
||||
assert cfg.vlm.api_key == "yaml-vlm-key"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试: CLI 覆盖
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCLIOverride:
|
||||
"""CLI args 应覆盖 YAML 和 .env 的值。"""
|
||||
|
||||
def test_cli_overrides_yaml(self, yaml_path: Path) -> None:
|
||||
"""CLI 点路径覆盖应生效。"""
|
||||
cfg = Config.load(
|
||||
str(yaml_path),
|
||||
cli_args={"retriever.num_heads": 8, "train.lr": 0.001},
|
||||
env_path=str(yaml_path.parent / ".env.nonexist"),
|
||||
)
|
||||
assert cfg.retriever.num_heads == 8
|
||||
assert cfg.train.lr == 0.001
|
||||
|
||||
def test_cli_overrides_env(self, yaml_path: Path, env_path: Path) -> None:
|
||||
"""CLI 应覆盖 .env 中的 api_key。"""
|
||||
cfg = Config.load(
|
||||
str(yaml_path),
|
||||
cli_args={"llm.api_key": "cli-key"},
|
||||
env_path=str(env_path),
|
||||
)
|
||||
assert cfg.llm.api_key == "cli-key"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试: 优先级
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPriority:
|
||||
"""三层优先级: CLI > .env > YAML。"""
|
||||
|
||||
def test_full_priority_chain(self, yaml_path: Path, env_path: Path) -> None:
|
||||
"""CLI > .env > YAML 的完整优先级链。"""
|
||||
cfg = Config.load(
|
||||
str(yaml_path),
|
||||
cli_args={"llm.api_key": "cli-key"},
|
||||
env_path=str(env_path),
|
||||
)
|
||||
# CLI 覆盖 .env
|
||||
assert cfg.llm.api_key == "cli-key"
|
||||
# .env 覆盖 YAML(vlm 未被 CLI 覆盖)
|
||||
assert cfg.vlm.api_key == "env-vlm-key"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试: embed_dim 一致性校验
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEmbedDimConsistency:
|
||||
"""embed.embed_dim 与 retriever.embed_dim 必须一致。"""
|
||||
|
||||
def test_inconsistent_embed_dim(self, tmp_path: Path) -> None:
|
||||
"""embed_dim 不一致应抛出 ValueError。"""
|
||||
bad_yaml = {
|
||||
k: (v.copy() if isinstance(v, dict) else v) for k, v in _FULL_YAML.items()
|
||||
}
|
||||
bad_yaml["retriever"] = bad_yaml["retriever"].copy()
|
||||
bad_yaml["retriever"]["embed_dim"] = 512 # 与 embed.embed_dim=768 不一致
|
||||
|
||||
p = tmp_path / "bad_dim.yaml"
|
||||
with open(p, "w") as f:
|
||||
yaml.dump(bad_yaml, f)
|
||||
|
||||
with pytest.raises(ValueError, match="不一致"):
|
||||
Config.load(str(p), env_path=str(tmp_path / ".env.nonexist"))
|
||||
@@ -0,0 +1,193 @@
|
||||
"""
|
||||
Embedding 持久化单元测试
|
||||
=======================
|
||||
测试 TreeIndex 的 embedding 序列化/反序列化功能。
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from video_tree_trm.tree_index import (
|
||||
L1Node,
|
||||
L2Node,
|
||||
L3Node,
|
||||
IndexMeta,
|
||||
TreeIndex,
|
||||
_embed_to_str,
|
||||
_embed_from_str,
|
||||
)
|
||||
|
||||
|
||||
class TestEmbeddingSerialization:
|
||||
"""测试 embedding 序列化辅助函数。"""
|
||||
|
||||
def test_embed_to_str_and_back(self):
|
||||
"""测试 base64 序列化/反序列化往返正确。"""
|
||||
arr = np.random.randn(768).astype(np.float32)
|
||||
s = _embed_to_str(arr)
|
||||
recovered = _embed_from_str(s)
|
||||
|
||||
assert recovered is not None
|
||||
assert recovered.dtype == np.float32
|
||||
assert recovered.shape == (768,)
|
||||
np.testing.assert_array_almost_equal(arr, recovered, decimal=6)
|
||||
|
||||
def test_embed_to_str_handles_none(self):
|
||||
"""测试 None 输入返回 None。"""
|
||||
assert _embed_to_str(None) is None
|
||||
assert _embed_from_str(None) is None
|
||||
assert _embed_from_str("") is None
|
||||
|
||||
|
||||
class TestL3NodeEmbedding:
|
||||
"""测试 L3Node embedding 序列化。"""
|
||||
|
||||
def test_l3_to_dict_without_embedding(self):
|
||||
"""测试不带 embedding 序列化。"""
|
||||
node = L3Node(
|
||||
id="l3_0",
|
||||
description="测试描述",
|
||||
timestamp=1.0,
|
||||
frame_path="frame.jpg",
|
||||
raw_content="原始内容",
|
||||
)
|
||||
d = {
|
||||
"id": node.id,
|
||||
"description": node.description,
|
||||
"timestamp": node.timestamp,
|
||||
"frame_path": node.frame_path,
|
||||
"raw_content": node.raw_content,
|
||||
}
|
||||
# 无 embedding 字段
|
||||
assert "embedding" not in d
|
||||
|
||||
def test_l3_embedding_roundtrip(self):
|
||||
"""测试 L3 embedding 序列化往返。"""
|
||||
embed = np.random.randn(768).astype(np.float32)
|
||||
s = _embed_to_str(embed)
|
||||
recovered = _embed_from_str(s)
|
||||
np.testing.assert_array_almost_equal(embed, recovered, decimal=6)
|
||||
|
||||
|
||||
class TestTreeIndexEmbedding:
|
||||
"""测试 TreeIndex 完整序列化/反序列化。"""
|
||||
|
||||
def test_save_load_without_embedding(self):
|
||||
"""测试不带 embedding 保存/加载(向后兼容)。"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
path = os.path.join(tmpdir, "test.json")
|
||||
|
||||
# 创建简单树
|
||||
l3 = L3Node(id="l3_0", description="L3 描述")
|
||||
l2 = L2Node(id="l2_0", description="L2 描述", children=[l3])
|
||||
l1 = L1Node(id="l1_0", summary="L1 摘要", children=[l2])
|
||||
meta = IndexMeta(source_path="test.mp4", modality="video")
|
||||
tree = TreeIndex(metadata=meta, roots=[l1])
|
||||
|
||||
# 保存(不含 embedding)
|
||||
tree.save_json(path, include_embedding=False)
|
||||
|
||||
# 加载
|
||||
loaded = TreeIndex.load_json(path)
|
||||
|
||||
assert len(loaded.roots) == 1
|
||||
assert loaded.roots[0].summary == "L1 摘要"
|
||||
assert not loaded.is_embedded # 无 embedding
|
||||
|
||||
def test_save_load_with_embedding(self):
|
||||
"""测试带 embedding 保存/加载。"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
path = os.path.join(tmpdir, "test_embed.json")
|
||||
|
||||
# 创建带 embedding 的树
|
||||
embed_l1 = np.random.randn(768).astype(np.float32)
|
||||
embed_l2 = np.random.randn(768).astype(np.float32)
|
||||
embed_l3 = np.random.randn(768).astype(np.float32)
|
||||
|
||||
l3 = L3Node(id="l3_0", description="L3 描述", embedding=embed_l3)
|
||||
l2 = L2Node(id="l2_0", description="L2 描述", embedding=embed_l2, children=[l3])
|
||||
l1 = L1Node(id="l1_0", summary="L1 摘要", embedding=embed_l1, children=[l2])
|
||||
meta = IndexMeta(
|
||||
source_path="test.mp4",
|
||||
modality="video",
|
||||
embed_model="test-model",
|
||||
embed_dim=768,
|
||||
)
|
||||
tree = TreeIndex(metadata=meta, roots=[l1])
|
||||
|
||||
# 保存(含 embedding)
|
||||
tree.save_json(path, include_embedding=True)
|
||||
|
||||
# 验证 JSON 中有 embedding 字段
|
||||
with open(path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
assert "embedding" in data["roots"][0]
|
||||
assert data["metadata"]["embed_model"] == "test-model"
|
||||
|
||||
# 加载
|
||||
loaded = TreeIndex.load_json(path)
|
||||
|
||||
assert loaded.is_embedded
|
||||
np.testing.assert_array_almost_equal(
|
||||
loaded.roots[0].embedding, embed_l1, decimal=6
|
||||
)
|
||||
np.testing.assert_array_almost_equal(
|
||||
loaded.roots[0].children[0].embedding, embed_l2, decimal=6
|
||||
)
|
||||
np.testing.assert_array_almost_equal(
|
||||
loaded.roots[0].children[0].children[0].embedding, embed_l3, decimal=6
|
||||
)
|
||||
|
||||
def test_load_old_format_compatible(self):
|
||||
"""测试加载旧格式(无 embedding 字段)JSON 兼容。"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
path = os.path.join(tmpdir, "old_format.json")
|
||||
|
||||
# 手动创建旧格式 JSON(无 embedding 字段)
|
||||
old_data = {
|
||||
"metadata": {
|
||||
"source_path": "test.mp4",
|
||||
"modality": "video",
|
||||
"created_at": "2024-01-01T00:00:00",
|
||||
},
|
||||
"roots": [
|
||||
{
|
||||
"id": "l1_0",
|
||||
"summary": "L1 摘要",
|
||||
"time_range": None,
|
||||
"children": [
|
||||
{
|
||||
"id": "l2_0",
|
||||
"description": "L2 描述",
|
||||
"time_range": None,
|
||||
"children": [
|
||||
{
|
||||
"id": "l3_0",
|
||||
"description": "L3 描述",
|
||||
"timestamp": 1.0,
|
||||
"frame_path": None,
|
||||
"raw_content": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
}
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
json.dump(old_data, f)
|
||||
|
||||
# 加载
|
||||
loaded = TreeIndex.load_json(path)
|
||||
|
||||
assert len(loaded.roots) == 1
|
||||
assert not loaded.is_embedded # 无 embedding,向后兼容
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -0,0 +1,199 @@
|
||||
"""
|
||||
嵌入服务模块单元测试
|
||||
====================
|
||||
覆盖: 本地模式(embed/embed_tensor/归一化/dim/冻结)、远程模式(真实 API 调用)。
|
||||
|
||||
本地测试使用轻量模型 all-MiniLM-L6-v2 (dim=384) 加速。
|
||||
远程测试使用 .env 中配置的真实 API(需有效密钥)。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from video_tree_trm.config import Config, EmbedConfig
|
||||
from video_tree_trm.embeddings import EmbeddingModel
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_LOCAL_CONFIG = EmbedConfig(
|
||||
backend="local",
|
||||
model_name="all-MiniLM-L6-v2",
|
||||
embed_dim=384,
|
||||
device="cpu",
|
||||
api_key="",
|
||||
api_url="",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def local_model() -> EmbeddingModel:
|
||||
"""本地嵌入模型(模块级缓存,避免重复加载)。"""
|
||||
return EmbeddingModel(_LOCAL_CONFIG)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def remote_model(real_config: Config) -> EmbeddingModel:
|
||||
"""远程嵌入模型(模块级缓存),使用真实 API。"""
|
||||
return EmbeddingModel(real_config.embed)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试: 本地模式
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestLocalEmbed:
|
||||
"""本地 sentence-transformers 后端测试。"""
|
||||
|
||||
def test_embed_single_text(self, local_model: EmbeddingModel) -> None:
|
||||
"""单条文本,验证形状 [1, D]。"""
|
||||
result = local_model.embed("hello world")
|
||||
assert isinstance(result, np.ndarray)
|
||||
assert result.shape == (1, 384)
|
||||
|
||||
def test_embed_batch(self, local_model: EmbeddingModel) -> None:
|
||||
"""批量文本,验证形状 [N, D]。"""
|
||||
texts = ["hello", "world", "test"]
|
||||
result = local_model.embed(texts)
|
||||
assert result.shape == (3, 384)
|
||||
|
||||
def test_embed_tensor_type(self, local_model: EmbeddingModel) -> None:
|
||||
"""验证 embed_tensor() 返回 Tensor。"""
|
||||
result = local_model.embed_tensor("hello")
|
||||
assert isinstance(result, torch.Tensor)
|
||||
assert result.shape == (1, 384)
|
||||
assert result.dtype == torch.float32
|
||||
|
||||
def test_embeddings_normalized(self, local_model: EmbeddingModel) -> None:
|
||||
"""验证 L2 范数 ≈ 1.0。"""
|
||||
result = local_model.embed(["hello", "world"])
|
||||
norms = np.linalg.norm(result, axis=1)
|
||||
np.testing.assert_allclose(norms, 1.0, atol=1e-5)
|
||||
|
||||
def test_dim_property(self, local_model: EmbeddingModel) -> None:
|
||||
"""验证 dim 属性一致。"""
|
||||
assert local_model.dim == 384
|
||||
|
||||
def test_local_model_frozen(self, local_model: EmbeddingModel) -> None:
|
||||
"""本地模式参数 requires_grad=False。"""
|
||||
for param in local_model._model.parameters():
|
||||
assert not param.requires_grad
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试: 远程模式(真实 API)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRemoteEmbed:
|
||||
"""远程 OpenAI 兼容 API 后端测试(真实 API 调用)。"""
|
||||
|
||||
def test_remote_embed_single(self, remote_model: EmbeddingModel) -> None:
|
||||
"""单条中文文本,验证形状 [1, D]。"""
|
||||
result = remote_model.embed("你好世界")
|
||||
assert isinstance(result, np.ndarray)
|
||||
assert result.shape == (1, remote_model.dim)
|
||||
|
||||
def test_remote_embed_batch(self, remote_model: EmbeddingModel) -> None:
|
||||
"""5 条文本,验证形状 [5, D]。"""
|
||||
texts = ["机器学习", "深度学习", "自然语言处理", "计算机视觉", "强化学习"]
|
||||
result = remote_model.embed(texts)
|
||||
assert result.shape == (5, remote_model.dim)
|
||||
|
||||
def test_remote_embed_normalized(self, remote_model: EmbeddingModel) -> None:
|
||||
"""验证 L2 范数 ≈ 1.0。"""
|
||||
result = remote_model.embed(["归一化测试文本", "另一段测试文本"])
|
||||
norms = np.linalg.norm(result, axis=1)
|
||||
np.testing.assert_allclose(norms, 1.0, atol=1e-5)
|
||||
|
||||
def test_remote_embed_tensor(self, remote_model: EmbeddingModel) -> None:
|
||||
"""验证返回 torch.Tensor + float32。"""
|
||||
result = remote_model.embed_tensor("张量测试")
|
||||
assert isinstance(result, torch.Tensor)
|
||||
assert result.dtype == torch.float32
|
||||
assert result.shape == (1, remote_model.dim)
|
||||
|
||||
def test_remote_dim(self, remote_model: EmbeddingModel, real_config: Config) -> None:
|
||||
"""验证 dim 属性与配置一致。"""
|
||||
assert remote_model.dim == real_config.embed.embed_dim
|
||||
|
||||
def test_remote_semantic_similarity(self, remote_model: EmbeddingModel) -> None:
|
||||
"""语义相近文本相似度 > 语义无关文本。"""
|
||||
# 语义相近对
|
||||
vec_cat = remote_model.embed("猫咪在沙发上睡觉")
|
||||
vec_kitten = remote_model.embed("小猫趴在沙发上打盹")
|
||||
# 语义无关
|
||||
vec_math = remote_model.embed("二次方程的求解公式")
|
||||
|
||||
sim_close = float(np.dot(vec_cat[0], vec_kitten[0]))
|
||||
sim_far = float(np.dot(vec_cat[0], vec_math[0]))
|
||||
assert sim_close > sim_far, (
|
||||
f"语义相近对相似度 ({sim_close:.4f}) 应大于语义无关对 ({sim_far:.4f})"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试: 配置校验
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestConfigValidation:
|
||||
"""配置校验测试。"""
|
||||
|
||||
def test_invalid_backend(self) -> None:
|
||||
"""无效 backend 应抛出 ValueError。"""
|
||||
config = EmbedConfig(
|
||||
backend="invalid",
|
||||
model_name="test",
|
||||
embed_dim=384,
|
||||
device="cpu",
|
||||
api_key="",
|
||||
api_url="",
|
||||
)
|
||||
with pytest.raises(ValueError, match="backend"):
|
||||
EmbeddingModel(config)
|
||||
|
||||
def test_remote_missing_api_key(self) -> None:
|
||||
"""远程模式缺少 api_key 应抛出 ValueError。"""
|
||||
config = EmbedConfig(
|
||||
backend="remote",
|
||||
model_name="test",
|
||||
embed_dim=384,
|
||||
device="cpu",
|
||||
api_key="",
|
||||
api_url="http://localhost:8080/v1",
|
||||
)
|
||||
with pytest.raises(ValueError, match="api_key"):
|
||||
EmbeddingModel(config)
|
||||
|
||||
def test_remote_missing_api_url(self) -> None:
|
||||
"""远程模式缺少 api_url 应抛出 ValueError。"""
|
||||
config = EmbedConfig(
|
||||
backend="remote",
|
||||
model_name="test",
|
||||
embed_dim=384,
|
||||
device="cpu",
|
||||
api_key="test-key",
|
||||
api_url="",
|
||||
)
|
||||
with pytest.raises(ValueError, match="api_url"):
|
||||
EmbeddingModel(config)
|
||||
|
||||
def test_local_dim_mismatch(self) -> None:
|
||||
"""本地模式维度不一致应抛出 ValueError。"""
|
||||
config = EmbedConfig(
|
||||
backend="local",
|
||||
model_name="all-MiniLM-L6-v2",
|
||||
embed_dim=999, # 实际是 384
|
||||
device="cpu",
|
||||
api_key="",
|
||||
api_url="",
|
||||
)
|
||||
with pytest.raises(ValueError, match="维度"):
|
||||
EmbeddingModel(config)
|
||||
@@ -0,0 +1,321 @@
|
||||
"""
|
||||
LLMClient 单元测试
|
||||
==================
|
||||
测试覆盖:
|
||||
- TestLLMChatMock: mock OpenAI 客户端的纯文本对话功能
|
||||
- TestVLMChatMock: mock 环境下的多模态图像对话功能
|
||||
- TestConfigValidation: 配置校验(api_key/api_url 缺失)
|
||||
- TestRealLLMChat: 真实 API 集成测试(需 .env 配置)
|
||||
|
||||
运行::
|
||||
|
||||
conda run -n Video-Tree-TRM python -m pytest tests/unit/test_llm_client.py -v \\
|
||||
--cov=video_tree_trm/llm_client --cov-report=term-missing
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import os
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from video_tree_trm.config import LLMConfig, VLMConfig
|
||||
from video_tree_trm.llm_client import LLMClient
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 辅助:构造最小配置对象(避免加载真实 YAML)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_llm_config(
|
||||
api_key: str = "sk-test",
|
||||
api_url: str = "https://api.example.com/v1",
|
||||
model: str = "test-model",
|
||||
max_tokens: int = 128,
|
||||
temperature: float = 0.1,
|
||||
) -> LLMConfig:
|
||||
"""构造测试用 LLMConfig,所有字段可覆盖。"""
|
||||
return LLMConfig(
|
||||
backend="openai",
|
||||
api_key=api_key,
|
||||
api_url=api_url,
|
||||
model=model,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
)
|
||||
|
||||
|
||||
def _make_vlm_config(
|
||||
api_key: str = "sk-test",
|
||||
api_url: str = "https://api.example.com/v1",
|
||||
model: str = "test-vlm",
|
||||
max_tokens: int = 128,
|
||||
temperature: float = 0.1,
|
||||
) -> VLMConfig:
|
||||
"""构造测试用 VLMConfig。"""
|
||||
return VLMConfig(
|
||||
backend="openai",
|
||||
api_key=api_key,
|
||||
api_url=api_url,
|
||||
model=model,
|
||||
max_tokens=max_tokens,
|
||||
temperature=temperature,
|
||||
)
|
||||
|
||||
|
||||
def _mock_completion(content: str) -> MagicMock:
|
||||
"""构造 openai.ChatCompletion 返回值的 Mock。"""
|
||||
choice = MagicMock()
|
||||
choice.message.content = content
|
||||
response = MagicMock()
|
||||
response.choices = [choice]
|
||||
return response
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestLLMChatMock — 纯文本对话(Mock)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestLLMChatMock:
|
||||
"""使用 mock openai.OpenAI 测试 chat() 和 batch_chat()。"""
|
||||
|
||||
@patch("video_tree_trm.llm_client.openai.OpenAI")
|
||||
def test_chat_returns_string(self, mock_openai_cls: MagicMock) -> None:
|
||||
"""chat() 应返回 API 返回内容的字符串。"""
|
||||
mock_client = MagicMock()
|
||||
mock_openai_cls.return_value = mock_client
|
||||
mock_client.chat.completions.create.return_value = _mock_completion("你好!")
|
||||
|
||||
llm = LLMClient(_make_llm_config())
|
||||
result = llm.chat("你好")
|
||||
|
||||
assert result == "你好!"
|
||||
assert isinstance(result, str)
|
||||
|
||||
@patch("video_tree_trm.llm_client.openai.OpenAI")
|
||||
def test_chat_uses_config_max_tokens(self, mock_openai_cls: MagicMock) -> None:
|
||||
"""未传 max_tokens 时,应使用 config.max_tokens 的值。"""
|
||||
mock_client = MagicMock()
|
||||
mock_openai_cls.return_value = mock_client
|
||||
mock_client.chat.completions.create.return_value = _mock_completion("ok")
|
||||
|
||||
cfg = _make_llm_config(max_tokens=256)
|
||||
llm = LLMClient(cfg)
|
||||
llm.chat("test")
|
||||
|
||||
call_kwargs = mock_client.chat.completions.create.call_args[1]
|
||||
assert call_kwargs["max_tokens"] == 256
|
||||
|
||||
@patch("video_tree_trm.llm_client.openai.OpenAI")
|
||||
def test_chat_overrides_max_tokens(self, mock_openai_cls: MagicMock) -> None:
|
||||
"""显式传入 max_tokens 时,应覆盖 config.max_tokens。"""
|
||||
mock_client = MagicMock()
|
||||
mock_openai_cls.return_value = mock_client
|
||||
mock_client.chat.completions.create.return_value = _mock_completion("ok")
|
||||
|
||||
cfg = _make_llm_config(max_tokens=256)
|
||||
llm = LLMClient(cfg)
|
||||
llm.chat("test", max_tokens=64)
|
||||
|
||||
call_kwargs = mock_client.chat.completions.create.call_args[1]
|
||||
assert call_kwargs["max_tokens"] == 64
|
||||
|
||||
@patch("video_tree_trm.llm_client.openai.OpenAI")
|
||||
def test_batch_chat_order_preserved(self, mock_openai_cls: MagicMock) -> None:
|
||||
"""batch_chat() 应按输入顺序返回结果,即使并发完成顺序不同。"""
|
||||
mock_client = MagicMock()
|
||||
mock_openai_cls.return_value = mock_client
|
||||
|
||||
# 每次调用返回不同内容
|
||||
responses = ["结果0", "结果1", "结果2"]
|
||||
mock_client.chat.completions.create.side_effect = [
|
||||
_mock_completion(r) for r in responses
|
||||
]
|
||||
|
||||
llm = LLMClient(_make_llm_config())
|
||||
results = llm.batch_chat(["prompt0", "prompt1", "prompt2"])
|
||||
|
||||
assert len(results) == 3
|
||||
assert results == responses
|
||||
|
||||
@patch("video_tree_trm.llm_client.openai.OpenAI")
|
||||
def test_batch_chat_empty_list(self, mock_openai_cls: MagicMock) -> None:
|
||||
"""batch_chat() 传入空列表时,应返回空列表。"""
|
||||
mock_openai_cls.return_value = MagicMock()
|
||||
llm = LLMClient(_make_llm_config())
|
||||
assert llm.batch_chat([]) == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestVLMChatMock — 多模态对话(Mock)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestVLMChatMock:
|
||||
"""使用 mock openai.OpenAI 测试 chat_with_images()。"""
|
||||
|
||||
@patch("video_tree_trm.llm_client.openai.OpenAI")
|
||||
def test_chat_with_images_encodes_local_path(self, mock_openai_cls: MagicMock) -> None:
|
||||
"""传入本地文件路径时,消息中应包含 data:image/jpeg;base64, 前缀。"""
|
||||
mock_client = MagicMock()
|
||||
mock_openai_cls.return_value = mock_client
|
||||
mock_client.chat.completions.create.return_value = _mock_completion("图中有猫")
|
||||
|
||||
# 创建临时 JPEG 文件
|
||||
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
|
||||
f.write(b"\xff\xd8\xff\xe0" + b"\x00" * 20) # 最小 JPEG header
|
||||
tmp_path = f.name
|
||||
|
||||
try:
|
||||
vlm = LLMClient(_make_vlm_config())
|
||||
result = vlm.chat_with_images("图中有什么?", images=[tmp_path])
|
||||
|
||||
assert result == "图中有猫"
|
||||
# 验证消息结构包含 base64 编码图像
|
||||
call_kwargs = mock_client.chat.completions.create.call_args[1]
|
||||
content = call_kwargs["messages"][0]["content"]
|
||||
assert isinstance(content, list)
|
||||
image_items = [c for c in content if c.get("type") == "image_url"]
|
||||
assert len(image_items) == 1
|
||||
assert image_items[0]["image_url"]["url"].startswith("data:image/jpeg;base64,")
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
@patch("video_tree_trm.llm_client.openai.OpenAI")
|
||||
def test_chat_with_images_accepts_b64(self, mock_openai_cls: MagicMock) -> None:
|
||||
"""传入已有 base64 字符串时,不应重复编码,直接透传。"""
|
||||
mock_client = MagicMock()
|
||||
mock_openai_cls.return_value = mock_client
|
||||
mock_client.chat.completions.create.return_value = _mock_completion("ok")
|
||||
|
||||
b64_str = "data:image/jpeg;base64," + base64.b64encode(b"fake").decode()
|
||||
vlm = LLMClient(_make_vlm_config())
|
||||
vlm.chat_with_images("描述图片", images=[b64_str])
|
||||
|
||||
call_kwargs = mock_client.chat.completions.create.call_args[1]
|
||||
content = call_kwargs["messages"][0]["content"]
|
||||
image_items = [c for c in content if c.get("type") == "image_url"]
|
||||
assert image_items[0]["image_url"]["url"] == b64_str
|
||||
|
||||
@patch("video_tree_trm.llm_client.openai.OpenAI")
|
||||
def test_chat_with_images_png_mime(self, mock_openai_cls: MagicMock) -> None:
|
||||
"""PNG 文件应编码为 data:image/png;base64, 格式。"""
|
||||
mock_client = MagicMock()
|
||||
mock_openai_cls.return_value = mock_client
|
||||
mock_client.chat.completions.create.return_value = _mock_completion("ok")
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
|
||||
f.write(b"\x89PNG\r\n\x1a\n" + b"\x00" * 20)
|
||||
tmp_path = f.name
|
||||
|
||||
try:
|
||||
vlm = LLMClient(_make_vlm_config())
|
||||
vlm.chat_with_images("描述图片", images=[tmp_path])
|
||||
|
||||
call_kwargs = mock_client.chat.completions.create.call_args[1]
|
||||
content = call_kwargs["messages"][0]["content"]
|
||||
image_items = [c for c in content if c.get("type") == "image_url"]
|
||||
assert image_items[0]["image_url"]["url"].startswith("data:image/png;base64,")
|
||||
finally:
|
||||
os.unlink(tmp_path)
|
||||
|
||||
@patch("video_tree_trm.llm_client.openai.OpenAI")
|
||||
def test_chat_with_images_message_structure(self, mock_openai_cls: MagicMock) -> None:
|
||||
"""多模态消息中图像应在文本之前。"""
|
||||
mock_client = MagicMock()
|
||||
mock_openai_cls.return_value = mock_client
|
||||
mock_client.chat.completions.create.return_value = _mock_completion("ok")
|
||||
|
||||
b64_str = "data:image/jpeg;base64," + base64.b64encode(b"img").decode()
|
||||
vlm = LLMClient(_make_vlm_config())
|
||||
vlm.chat_with_images("提问", images=[b64_str])
|
||||
|
||||
call_kwargs = mock_client.chat.completions.create.call_args[1]
|
||||
content = call_kwargs["messages"][0]["content"]
|
||||
# 最后一项为 text
|
||||
assert content[-1]["type"] == "text"
|
||||
assert content[-1]["text"] == "提问"
|
||||
# 前面各项为 image_url
|
||||
for item in content[:-1]:
|
||||
assert item["type"] == "image_url"
|
||||
|
||||
def test_encode_image_file_not_found(self) -> None:
|
||||
"""_encode_image 传入不存在的路径时,应抛出 FileNotFoundError。"""
|
||||
with patch("video_tree_trm.llm_client.openai.OpenAI"):
|
||||
llm = LLMClient(_make_llm_config())
|
||||
|
||||
with pytest.raises(FileNotFoundError):
|
||||
llm._encode_image("/nonexistent/path/image.jpg")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestConfigValidation — 配置校验
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestConfigValidation:
|
||||
"""测试 LLMClient 初始化时的配置校验逻辑。"""
|
||||
|
||||
def test_missing_api_key_raises(self) -> None:
|
||||
"""api_key 为空时应抛出 ValueError。"""
|
||||
cfg = _make_llm_config(api_key="")
|
||||
with pytest.raises(ValueError, match="api_key"):
|
||||
LLMClient(cfg)
|
||||
|
||||
def test_missing_api_url_raises(self) -> None:
|
||||
"""api_url 为空时应抛出 ValueError。"""
|
||||
cfg = _make_llm_config(api_url="")
|
||||
with pytest.raises(ValueError, match="api_url"):
|
||||
LLMClient(cfg)
|
||||
|
||||
@patch("video_tree_trm.llm_client.openai.OpenAI")
|
||||
def test_valid_config_initializes(self, mock_openai_cls: MagicMock) -> None:
|
||||
"""有效配置应正常初始化,不抛出异常。"""
|
||||
mock_openai_cls.return_value = MagicMock()
|
||||
cfg = _make_llm_config()
|
||||
client = LLMClient(cfg)
|
||||
assert client is not None
|
||||
|
||||
@patch("video_tree_trm.llm_client.openai.OpenAI")
|
||||
def test_vlm_config_accepted(self, mock_openai_cls: MagicMock) -> None:
|
||||
"""VLMConfig 也应被正常接受。"""
|
||||
mock_openai_cls.return_value = MagicMock()
|
||||
cfg = _make_vlm_config()
|
||||
client = LLMClient(cfg)
|
||||
assert client is not None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TestRealLLMChat — 真实 API 集成测试(需 .env)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRealLLMChat:
|
||||
"""调用真实 LLM API 进行集成测试。
|
||||
|
||||
需要 .env 中配置有效的 LLM_API_KEY / LLM_API_URL / LLM_MODEL。
|
||||
"""
|
||||
|
||||
def test_real_chat(self, real_config) -> None: # noqa: ANN001
|
||||
"""真实 API 单轮对话,应返回非空字符串。"""
|
||||
llm = LLMClient(real_config.llm)
|
||||
result = llm.chat("请用一句话回答:天空是什么颜色?", max_tokens=32)
|
||||
assert isinstance(result, str)
|
||||
assert len(result) > 0
|
||||
|
||||
def test_real_batch_chat(self, real_config) -> None: # noqa: ANN001
|
||||
"""真实 API 批量对话,应返回与输入等长的非空字符串列表。"""
|
||||
llm = LLMClient(real_config.llm)
|
||||
prompts = ["1+1等于几?请只回答数字。", "2+2等于几?请只回答数字。"]
|
||||
results = llm.batch_chat(prompts, max_tokens=16)
|
||||
assert len(results) == 2
|
||||
for r in results:
|
||||
assert isinstance(r, str)
|
||||
assert len(r) > 0
|
||||
@@ -0,0 +1,357 @@
|
||||
"""
|
||||
test_pipeline.py — Pipeline 单元测试
|
||||
======================================
|
||||
使用 unittest.mock.MagicMock + patch 隔离所有外部依赖(无真实 API / 文件 IO)。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from video_tree_trm.pipeline import Pipeline
|
||||
from video_tree_trm.tree_index import IndexMeta, L1Node, L2Node, L3Node, TreeIndex
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 辅助:构造最小 Config Mock
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
D = 8 # 嵌入维度(与 RetrieverConfig 一致)
|
||||
|
||||
|
||||
def _make_config(checkpoint: str | None = None) -> MagicMock:
|
||||
"""返回一个 Mock Config,字段值与实际 dataclass 对齐。"""
|
||||
cfg = MagicMock()
|
||||
cfg.embed.model_name = "test-embed"
|
||||
cfg.embed.embed_dim = D
|
||||
cfg.retriever.checkpoint = checkpoint
|
||||
cfg.retriever.embed_dim = D
|
||||
cfg.retriever.num_heads = 2
|
||||
cfg.retriever.L_layers = 2
|
||||
cfg.retriever.L_cycles = 2
|
||||
cfg.retriever.max_rounds = 2
|
||||
cfg.retriever.ffn_expansion = 2.0
|
||||
cfg.tree.cache_dir = "/tmp/test_pipeline_cache"
|
||||
return cfg
|
||||
|
||||
|
||||
def _make_small_tree() -> TreeIndex:
|
||||
"""构造最小 1×1×1 TreeIndex,用于 query() 测试。"""
|
||||
meta = IndexMeta(
|
||||
source_path="dummy",
|
||||
modality="text",
|
||||
embed_model="test",
|
||||
embed_dim=D,
|
||||
)
|
||||
l3 = L3Node(
|
||||
id="l3_0",
|
||||
description="节点描述",
|
||||
embedding=np.zeros(D, dtype=np.float32),
|
||||
raw_content="节点内容",
|
||||
)
|
||||
l2 = L2Node(
|
||||
id="l2_0",
|
||||
description="L2",
|
||||
embedding=np.zeros(D, dtype=np.float32),
|
||||
children=[l3],
|
||||
)
|
||||
l1 = L1Node(
|
||||
id="l1_0",
|
||||
summary="L1",
|
||||
embedding=np.zeros(D, dtype=np.float32),
|
||||
children=[l2],
|
||||
)
|
||||
return TreeIndex(metadata=meta, roots=[l1])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Patch 工厂:将所有子模块构造函数替换为 MagicMock
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_PATCHES = [
|
||||
"video_tree_trm.pipeline.EmbeddingModel",
|
||||
"video_tree_trm.pipeline.LLMClient",
|
||||
"video_tree_trm.pipeline.RecursiveRetriever",
|
||||
"video_tree_trm.pipeline.AnswerGenerator",
|
||||
"video_tree_trm.pipeline.TextTreeBuilder",
|
||||
"video_tree_trm.pipeline.VideoTreeBuilder",
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pipeline.__init__ 测试
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_pipeline_init_components() -> None:
|
||||
"""__init__ 后各属性(embed_model/llm/vlm/retriever/generator)均存在。"""
|
||||
cfg = _make_config()
|
||||
with patch.multiple(
|
||||
"video_tree_trm.pipeline",
|
||||
EmbeddingModel=MagicMock(),
|
||||
LLMClient=MagicMock(),
|
||||
RecursiveRetriever=MagicMock(),
|
||||
AnswerGenerator=MagicMock(),
|
||||
):
|
||||
p = Pipeline(cfg)
|
||||
|
||||
assert hasattr(p, "embed_model"), "缺少 embed_model 属性"
|
||||
assert hasattr(p, "llm"), "缺少 llm 属性"
|
||||
assert hasattr(p, "vlm"), "缺少 vlm 属性"
|
||||
assert hasattr(p, "retriever"), "缺少 retriever 属性"
|
||||
assert hasattr(p, "generator"), "缺少 generator 属性"
|
||||
|
||||
|
||||
def test_pipeline_init_no_checkpoint() -> None:
|
||||
"""checkpoint=None 时 load_state_dict 不被调用。"""
|
||||
cfg = _make_config(checkpoint=None)
|
||||
mock_retriever_instance = MagicMock()
|
||||
MockRetriever = MagicMock(return_value=mock_retriever_instance)
|
||||
|
||||
with patch.multiple(
|
||||
"video_tree_trm.pipeline",
|
||||
EmbeddingModel=MagicMock(),
|
||||
LLMClient=MagicMock(),
|
||||
RecursiveRetriever=MockRetriever,
|
||||
AnswerGenerator=MagicMock(),
|
||||
):
|
||||
Pipeline(cfg)
|
||||
|
||||
mock_retriever_instance.load_state_dict.assert_not_called()
|
||||
|
||||
|
||||
def test_pipeline_init_with_checkpoint(tmp_path: Path) -> None:
|
||||
"""checkpoint 非 None 时 load_state_dict 被调用一次。"""
|
||||
ckpt_file = tmp_path / "model.pt"
|
||||
ckpt_file.write_bytes(b"") # 创建空文件,使 os.path.isfile 返回 True
|
||||
|
||||
cfg = _make_config(checkpoint=str(ckpt_file))
|
||||
mock_retriever_instance = MagicMock()
|
||||
MockRetriever = MagicMock(return_value=mock_retriever_instance)
|
||||
|
||||
fake_state_dict = {"weight": torch.zeros(1)}
|
||||
|
||||
with patch.multiple(
|
||||
"video_tree_trm.pipeline",
|
||||
EmbeddingModel=MagicMock(),
|
||||
LLMClient=MagicMock(),
|
||||
RecursiveRetriever=MockRetriever,
|
||||
AnswerGenerator=MagicMock(),
|
||||
), patch("video_tree_trm.pipeline.torch.load", return_value=fake_state_dict):
|
||||
Pipeline(cfg)
|
||||
|
||||
mock_retriever_instance.load_state_dict.assert_called_once_with(fake_state_dict)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pipeline.build_index 测试
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_build_index_text_calls_builder(tmp_path: Path) -> None:
|
||||
"""文本模式调用 TextTreeBuilder.build,参数含文件内容。"""
|
||||
src = tmp_path / "doc.txt"
|
||||
src.write_text("文档内容", encoding="utf-8")
|
||||
|
||||
cfg = _make_config()
|
||||
cfg.tree.cache_dir = str(tmp_path / "cache")
|
||||
|
||||
mock_tree = MagicMock(spec=TreeIndex)
|
||||
mock_builder_instance = MagicMock()
|
||||
mock_builder_instance.build.return_value = mock_tree
|
||||
MockTextBuilder = MagicMock(return_value=mock_builder_instance)
|
||||
|
||||
with patch.multiple(
|
||||
"video_tree_trm.pipeline",
|
||||
EmbeddingModel=MagicMock(),
|
||||
LLMClient=MagicMock(),
|
||||
RecursiveRetriever=MagicMock(),
|
||||
AnswerGenerator=MagicMock(),
|
||||
TextTreeBuilder=MockTextBuilder,
|
||||
):
|
||||
p = Pipeline(cfg)
|
||||
result = p.build_index(str(src), modality="text")
|
||||
|
||||
mock_builder_instance.build.assert_called_once()
|
||||
call_args = mock_builder_instance.build.call_args
|
||||
assert "文档内容" in call_args[0][0], "TextTreeBuilder.build 应传入文件内容"
|
||||
assert result is mock_tree
|
||||
|
||||
|
||||
def test_build_index_video_calls_builder(tmp_path: Path) -> None:
|
||||
"""视频模式调用 VideoTreeBuilder.build,参数为 source_path。"""
|
||||
cfg = _make_config()
|
||||
cfg.tree.cache_dir = str(tmp_path / "cache")
|
||||
|
||||
mock_tree = MagicMock(spec=TreeIndex)
|
||||
mock_builder_instance = MagicMock()
|
||||
mock_builder_instance.build.return_value = mock_tree
|
||||
MockVideoBuilder = MagicMock(return_value=mock_builder_instance)
|
||||
|
||||
video_path = "/fake/video.mp4"
|
||||
|
||||
with patch.multiple(
|
||||
"video_tree_trm.pipeline",
|
||||
EmbeddingModel=MagicMock(),
|
||||
LLMClient=MagicMock(),
|
||||
RecursiveRetriever=MagicMock(),
|
||||
AnswerGenerator=MagicMock(),
|
||||
VideoTreeBuilder=MockVideoBuilder,
|
||||
):
|
||||
p = Pipeline(cfg)
|
||||
result = p.build_index(video_path, modality="video")
|
||||
|
||||
mock_builder_instance.build.assert_called_once_with(video_path)
|
||||
assert result is mock_tree
|
||||
|
||||
|
||||
def test_build_index_cache_hit(tmp_path: Path) -> None:
|
||||
"""缓存文件存在时直接 TreeIndex.load,不重新构建。"""
|
||||
cfg = _make_config()
|
||||
cache_dir = tmp_path / "cache"
|
||||
cache_dir.mkdir()
|
||||
cfg.tree.cache_dir = str(cache_dir)
|
||||
|
||||
# 手动创建缓存文件(空文件即可让 isfile 返回 True)
|
||||
cache_file = cache_dir / "doc_text.pkl"
|
||||
cache_file.write_bytes(b"")
|
||||
|
||||
mock_tree = MagicMock(spec=TreeIndex)
|
||||
mock_text_builder = MagicMock()
|
||||
|
||||
with patch.multiple(
|
||||
"video_tree_trm.pipeline",
|
||||
EmbeddingModel=MagicMock(),
|
||||
LLMClient=MagicMock(),
|
||||
RecursiveRetriever=MagicMock(),
|
||||
AnswerGenerator=MagicMock(),
|
||||
TextTreeBuilder=mock_text_builder,
|
||||
), patch("video_tree_trm.pipeline.TreeIndex.load", return_value=mock_tree) as mock_load:
|
||||
p = Pipeline(cfg)
|
||||
result = p.build_index(str(tmp_path / "doc.txt"), modality="text")
|
||||
|
||||
mock_load.assert_called_once_with(str(cache_file))
|
||||
mock_text_builder.return_value.build.assert_not_called()
|
||||
assert result is mock_tree
|
||||
|
||||
|
||||
def test_build_index_saves_cache(tmp_path: Path) -> None:
|
||||
"""缓存不存在时构建后调用 tree.save。"""
|
||||
cfg = _make_config()
|
||||
cfg.tree.cache_dir = str(tmp_path / "cache")
|
||||
|
||||
src = tmp_path / "doc.txt"
|
||||
src.write_text("内容", encoding="utf-8")
|
||||
|
||||
mock_tree = MagicMock(spec=TreeIndex)
|
||||
mock_builder_instance = MagicMock()
|
||||
mock_builder_instance.build.return_value = mock_tree
|
||||
|
||||
with patch.multiple(
|
||||
"video_tree_trm.pipeline",
|
||||
EmbeddingModel=MagicMock(),
|
||||
LLMClient=MagicMock(),
|
||||
RecursiveRetriever=MagicMock(),
|
||||
AnswerGenerator=MagicMock(),
|
||||
TextTreeBuilder=MagicMock(return_value=mock_builder_instance),
|
||||
):
|
||||
p = Pipeline(cfg)
|
||||
p.build_index(str(src), modality="text")
|
||||
|
||||
mock_tree.save.assert_called_once()
|
||||
saved_path: str = mock_tree.save.call_args[0][0]
|
||||
assert "doc_text.pkl" in saved_path, f"保存路径应含 'doc_text.pkl',实际={saved_path}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pipeline.query 测试
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_query_embeds_question() -> None:
|
||||
"""query() 调用 embed_model.embed_tensor(question)。"""
|
||||
cfg = _make_config()
|
||||
tree = _make_small_tree()
|
||||
|
||||
mock_embed = MagicMock()
|
||||
mock_embed.embed_tensor.return_value = torch.zeros(1, D)
|
||||
MockEmbed = MagicMock(return_value=mock_embed)
|
||||
|
||||
mock_retriever_instance = MagicMock()
|
||||
mock_retriever_instance.return_value = {"paths": [(0, 0, 0)], "num_rounds": 1}
|
||||
|
||||
with patch.multiple(
|
||||
"video_tree_trm.pipeline",
|
||||
EmbeddingModel=MockEmbed,
|
||||
LLMClient=MagicMock(),
|
||||
RecursiveRetriever=MagicMock(return_value=mock_retriever_instance),
|
||||
AnswerGenerator=MagicMock(),
|
||||
):
|
||||
p = Pipeline(cfg)
|
||||
p.query("测试问题", tree)
|
||||
|
||||
mock_embed.embed_tensor.assert_called_once_with("测试问题")
|
||||
|
||||
|
||||
def test_query_calls_retriever() -> None:
|
||||
"""query() 调用 retriever(q, tree)。"""
|
||||
cfg = _make_config()
|
||||
tree = _make_small_tree()
|
||||
|
||||
q_tensor = torch.zeros(1, D)
|
||||
mock_embed = MagicMock()
|
||||
mock_embed.embed_tensor.return_value = q_tensor
|
||||
|
||||
mock_retriever_instance = MagicMock()
|
||||
mock_retriever_instance.return_value = {"paths": [(0, 0, 0)], "num_rounds": 1}
|
||||
|
||||
with patch.multiple(
|
||||
"video_tree_trm.pipeline",
|
||||
EmbeddingModel=MagicMock(return_value=mock_embed),
|
||||
LLMClient=MagicMock(),
|
||||
RecursiveRetriever=MagicMock(return_value=mock_retriever_instance),
|
||||
AnswerGenerator=MagicMock(),
|
||||
):
|
||||
p = Pipeline(cfg)
|
||||
p.query("测试问题", tree)
|
||||
|
||||
mock_retriever_instance.assert_called_once()
|
||||
call_args = mock_retriever_instance.call_args
|
||||
# 第一个位置参数应为嵌入 Tensor,第二个为 tree
|
||||
assert call_args[0][1] is tree, "retriever 第二个参数应为 tree"
|
||||
|
||||
|
||||
def test_query_returns_answer() -> None:
|
||||
"""query() 返回 generator.generate 的返回值。"""
|
||||
cfg = _make_config()
|
||||
tree = _make_small_tree()
|
||||
|
||||
mock_embed = MagicMock()
|
||||
mock_embed.embed_tensor.return_value = torch.zeros(1, D)
|
||||
|
||||
mock_retriever_instance = MagicMock()
|
||||
mock_retriever_instance.return_value = {"paths": [(0, 0, 0)], "num_rounds": 1}
|
||||
|
||||
mock_generator_instance = MagicMock()
|
||||
mock_generator_instance.generate.return_value = "生成的答案"
|
||||
|
||||
with patch.multiple(
|
||||
"video_tree_trm.pipeline",
|
||||
EmbeddingModel=MagicMock(return_value=mock_embed),
|
||||
LLMClient=MagicMock(),
|
||||
RecursiveRetriever=MagicMock(return_value=mock_retriever_instance),
|
||||
AnswerGenerator=MagicMock(return_value=mock_generator_instance),
|
||||
):
|
||||
p = Pipeline(cfg)
|
||||
answer = p.query("问题", tree)
|
||||
|
||||
assert answer == "生成的答案", f"query() 应返回 generator 的结果,实际='{answer}'"
|
||||
mock_generator_instance.generate.assert_called_once_with(
|
||||
"问题", [(0, 0, 0)], tree
|
||||
)
|
||||
@@ -0,0 +1,546 @@
|
||||
"""
|
||||
TextTreeBuilder 单元测试
|
||||
========================
|
||||
使用 MagicMock 替代真实 LLM 和 Embed,测试文本树构建各阶段的正确性。
|
||||
|
||||
覆盖范围:
|
||||
- _detect_toc:检测 Markdown 标题
|
||||
- _segment_with_regex:正则切分 L1/L2 边界、超限分块
|
||||
- _segment_with_llm:LLM JSON 分段解析、异常处理
|
||||
- _build_l3_from_paragraphs:L3 节点字段验证
|
||||
- _build_l2:L2 节点字段验证(通过 build() 间接调用)
|
||||
- _build_l1:L1 节点字段验证(通过 build() 间接调用)
|
||||
- build():完整树结构与 IndexMeta 验证、MD 输出文件生成
|
||||
|
||||
MD 输出位置: tests/outputs/text_tree_builder/build_<timestamp>.md
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from video_tree_trm.config import TreeConfig
|
||||
from video_tree_trm.text_tree_builder import TextTreeBuilder, _chunk
|
||||
from video_tree_trm.tree_index import L1Node, L2Node, L3Node, TreeIndex
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 常量
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_EMBED_DIM = 16 # 轻量维度,加速测试
|
||||
|
||||
_SAMPLE_TOC_TEXT = """\
|
||||
# 第一章 引言
|
||||
|
||||
本章介绍研究背景。
|
||||
|
||||
信息检索系统的发展历程悠久,早期以关键词匹配为主。
|
||||
|
||||
## 1.1 研究背景
|
||||
|
||||
随着互联网数据量急剧增长,传统检索方法面临挑战。
|
||||
|
||||
语义理解成为关键技术突破口。
|
||||
|
||||
## 1.2 研究意义
|
||||
|
||||
本研究具有重要的理论和实践价值。
|
||||
|
||||
# 第二章 相关工作
|
||||
|
||||
本章回顾相关研究成果。
|
||||
|
||||
## 2.1 稠密检索
|
||||
|
||||
DPR 等模型实现了端到端稠密向量检索。
|
||||
|
||||
## 2.2 树状索引
|
||||
|
||||
PageIndex 引入了层次化树状检索结构。
|
||||
"""
|
||||
|
||||
_SAMPLE_PLAIN_TEXT = (
|
||||
"信息检索是计算机科学的重要分支,研究如何从大规模数据中找到相关信息。"
|
||||
"\n\n"
|
||||
"传统方法以 TF-IDF 和 BM25 为代表,基于词频统计进行相关性计算。"
|
||||
"\n\n"
|
||||
"近年来,基于深度学习的稠密检索方法取得了显著进展,DPR 是代表性工作之一。"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tree_config() -> TreeConfig:
|
||||
"""树构建配置(小 max_paragraphs_per_l2 便于测试分块)。"""
|
||||
return TreeConfig(
|
||||
max_paragraphs_per_l2=2,
|
||||
l1_segment_duration=600.0,
|
||||
l2_clip_duration=20.0,
|
||||
l3_fps=1.0,
|
||||
l2_representative_frames=3,
|
||||
cache_dir="cache/trees",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_embed() -> MagicMock:
|
||||
"""Mock EmbeddingModel,embed() 返回固定维度的随机向量。"""
|
||||
m = MagicMock()
|
||||
m.dim = _EMBED_DIM
|
||||
m._model_name = "mock-embed-model"
|
||||
|
||||
def _embed(texts):
|
||||
if isinstance(texts, str):
|
||||
texts = [texts]
|
||||
n = len(texts)
|
||||
return np.ones((n, _EMBED_DIM), dtype=np.float32)
|
||||
|
||||
m.embed.side_effect = _embed
|
||||
return m
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_llm() -> MagicMock:
|
||||
"""Mock LLMClient,chat() 返回固定字符串,batch_chat() 逐条映射。"""
|
||||
m = MagicMock()
|
||||
m.chat.return_value = "这是一段模拟的摘要描述。"
|
||||
m.batch_chat.side_effect = lambda prompts: [
|
||||
f"模拟摘要_{i}" for i in range(len(prompts))
|
||||
]
|
||||
return m
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def builder(mock_embed, mock_llm, tree_config) -> TextTreeBuilder:
|
||||
"""标准 TextTreeBuilder(使用 mock 依赖)。"""
|
||||
return TextTreeBuilder(
|
||||
embed_model=mock_embed,
|
||||
llm=mock_llm,
|
||||
config=tree_config,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="session", autouse=True)
|
||||
def ensure_output_dir():
|
||||
"""确保 MD 输出目录存在。"""
|
||||
Path("tests/outputs/text_tree_builder").mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 辅助函数:保存 MD 输出
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _save_md_output(test_name: str, content: str) -> str:
|
||||
"""将测试执行过程保存为 Markdown 文件。
|
||||
|
||||
参数:
|
||||
test_name: 测试名称(用于文件名)。
|
||||
content: Markdown 内容。
|
||||
|
||||
返回:
|
||||
保存的文件路径。
|
||||
"""
|
||||
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
path = Path(f"tests/outputs/text_tree_builder/{test_name}_{ts}.md")
|
||||
path.write_text(content, encoding="utf-8")
|
||||
return str(path)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 辅助函数:_chunk
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestChunk:
|
||||
"""测试 _chunk 等长分块辅助函数。"""
|
||||
|
||||
def test_exact_multiple(self):
|
||||
"""整除时每组大小相等。"""
|
||||
result = _chunk(["a", "b", "c", "d"], 2)
|
||||
assert result == [["a", "b"], ["c", "d"]]
|
||||
|
||||
def test_remainder(self):
|
||||
"""不整除时最后一组为余数。"""
|
||||
result = _chunk(["a", "b", "c", "d", "e"], 2)
|
||||
assert result == [["a", "b"], ["c", "d"], ["e"]]
|
||||
|
||||
def test_single_chunk(self):
|
||||
"""列表长度 <= size 时只有一组。"""
|
||||
result = _chunk(["a", "b"], 5)
|
||||
assert result == [["a", "b"]]
|
||||
|
||||
def test_empty_list(self):
|
||||
"""空列表返回空列表。"""
|
||||
result = _chunk([], 3)
|
||||
assert result == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _detect_toc
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDetectToc:
|
||||
"""测试 Markdown 标题检测。"""
|
||||
|
||||
def test_with_h1_header(self, builder):
|
||||
"""一级标题返回 True。"""
|
||||
assert builder._detect_toc("# 第一章\n\n内容") is True
|
||||
|
||||
def test_with_h2_header(self, builder):
|
||||
"""二级标题返回 True。"""
|
||||
assert builder._detect_toc("## 1.1 小节\n\n内容") is True
|
||||
|
||||
def test_with_both_headers(self, builder):
|
||||
"""同时含 # 和 ## 返回 True。"""
|
||||
assert builder._detect_toc(_SAMPLE_TOC_TEXT) is True
|
||||
|
||||
def test_without_headers(self, builder):
|
||||
"""纯段落文本返回 False。"""
|
||||
assert builder._detect_toc(_SAMPLE_PLAIN_TEXT) is False
|
||||
|
||||
def test_hash_in_content_not_header(self, builder):
|
||||
"""行中间的 # 不算标题。"""
|
||||
text = "颜色代码 #FF0000 是红色。\n\n这是普通段落。"
|
||||
assert builder._detect_toc(text) is False
|
||||
|
||||
def test_h3_not_counted(self, builder):
|
||||
"""三级标题 ### 也被检测为有 ToC(属于 Markdown 结构文本)。
|
||||
|
||||
注:当前实现只检测 # 和 ##,所以 ### 开头应返回 False。
|
||||
"""
|
||||
text = "### 三级标题\n\n内容"
|
||||
# _detect_toc 只匹配 #{1,2},### 不匹配
|
||||
assert builder._detect_toc(text) is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _segment_with_regex
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSegmentWithRegex:
|
||||
"""测试正则切分 L1/L2 边界。"""
|
||||
|
||||
def test_basic_structure(self, builder):
|
||||
"""含 2 个 L1 章节,正确返回 2 个外层元素。"""
|
||||
sections = builder._segment_with_regex(_SAMPLE_TOC_TEXT)
|
||||
assert len(sections) == 2
|
||||
|
||||
def test_all_sections_nonempty(self, builder):
|
||||
"""每个 section 至少包含一个段落。"""
|
||||
sections = builder._segment_with_regex(_SAMPLE_TOC_TEXT)
|
||||
for s in sections:
|
||||
assert len(s) > 0
|
||||
|
||||
def test_overflow_chunking_via_build(self, builder):
|
||||
"""当段落数超过 max_paragraphs_per_l2=2 时,build() 会等长分块为多个 L2。"""
|
||||
# 第一章有 4 个段落(含标题),超过 max=2 → 应有 2 个 L2
|
||||
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
|
||||
# 每个 L1 的 L2 数 >= 1
|
||||
for l1 in index.roots:
|
||||
assert len(l1.children) >= 1
|
||||
|
||||
def test_paragraphs_are_strings(self, builder):
|
||||
"""所有段落应为非空字符串。"""
|
||||
sections = builder._segment_with_regex(_SAMPLE_TOC_TEXT)
|
||||
for section in sections:
|
||||
for para in section:
|
||||
assert isinstance(para, str)
|
||||
assert para.strip() != ""
|
||||
|
||||
def test_single_chapter_text(self, builder):
|
||||
"""只含一个 L1 章节的文本,返回 1 个外层元素。"""
|
||||
text = "# 唯一章节\n\n第一段内容。\n\n第二段内容。"
|
||||
sections = builder._segment_with_regex(text)
|
||||
assert len(sections) == 1
|
||||
|
||||
def test_no_l2_header(self, builder):
|
||||
"""只有 L1 无 L2 标题时,段落直接收集到 L1 组。"""
|
||||
text = "# 第一章\n\n段落一。\n\n段落二。\n\n段落三。"
|
||||
sections = builder._segment_with_regex(text)
|
||||
assert len(sections) == 1
|
||||
assert len(sections[0]) >= 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _segment_with_llm
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSegmentWithLLM:
|
||||
"""测试 LLM 语义分段。"""
|
||||
|
||||
def test_json_array_parsing(self, builder):
|
||||
"""mock LLM 返回合法 JSON 数组,应正确解析为段落列表。"""
|
||||
paragraphs = ["第一段描述内容。", "第二段描述内容。", "第三段描述内容。"]
|
||||
builder.llm.chat.return_value = json.dumps(paragraphs, ensure_ascii=False)
|
||||
|
||||
sections = builder._segment_with_llm(_SAMPLE_PLAIN_TEXT)
|
||||
assert len(sections) == 1 # 单个 L1
|
||||
assert len(sections[0]) == 3
|
||||
assert sections[0][0] == "第一段描述内容。"
|
||||
|
||||
def test_json_wrapped_in_markdown(self, builder):
|
||||
"""JSON 数组被 markdown 代码块包裹时也能正确解析。"""
|
||||
paragraphs = ["段落A", "段落B"]
|
||||
json_str = json.dumps(paragraphs, ensure_ascii=False)
|
||||
builder.llm.chat.return_value = f"```json\n{json_str}\n```"
|
||||
|
||||
sections = builder._segment_with_llm(_SAMPLE_PLAIN_TEXT)
|
||||
assert sections[0] == ["段落A", "段落B"]
|
||||
|
||||
def test_invalid_json_raises(self, builder):
|
||||
"""LLM 返回非法 JSON 时应抛出 ValueError。"""
|
||||
builder.llm.chat.return_value = "这不是 JSON 格式的内容"
|
||||
with pytest.raises((ValueError, AssertionError)):
|
||||
builder._segment_with_llm(_SAMPLE_PLAIN_TEXT)
|
||||
|
||||
def test_empty_paragraphs_filtered(self, builder):
|
||||
"""空段落应被过滤掉。"""
|
||||
paragraphs = ["有效段落", "", " ", "另一有效段落"]
|
||||
builder.llm.chat.return_value = json.dumps(paragraphs, ensure_ascii=False)
|
||||
sections = builder._segment_with_llm(_SAMPLE_PLAIN_TEXT)
|
||||
assert len(sections[0]) == 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _build_l3_from_paragraphs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBuildL3:
|
||||
"""测试 L3 节点构建。"""
|
||||
|
||||
def test_description_equals_raw_content(self, builder):
|
||||
"""L3 节点的 description 应等于 raw_content,等于原始段落。"""
|
||||
paragraphs = ["段落一内容", "段落二内容"]
|
||||
nodes = builder._build_l3_from_paragraphs(paragraphs, l1_i=0, l2_j=0)
|
||||
for node, para in zip(nodes, paragraphs):
|
||||
assert node.description == para
|
||||
assert node.raw_content == para
|
||||
|
||||
def test_embedding_shape(self, builder):
|
||||
"""每个 L3 节点的 embedding 形状应为 (dim,)。"""
|
||||
paragraphs = ["A", "B", "C"]
|
||||
nodes = builder._build_l3_from_paragraphs(paragraphs, l1_i=0, l2_j=0)
|
||||
for node in nodes:
|
||||
assert node.embedding.shape == (_EMBED_DIM,)
|
||||
|
||||
def test_node_count(self, builder):
|
||||
"""返回的 L3 节点数应与输入段落数相等。"""
|
||||
paragraphs = ["A", "B", "C"]
|
||||
nodes = builder._build_l3_from_paragraphs(paragraphs, l1_i=0, l2_j=0)
|
||||
assert len(nodes) == 3
|
||||
|
||||
def test_node_id_format(self, builder):
|
||||
"""节点 ID 格式应为 l1_{i}_l2_{j}_l3_{k}。"""
|
||||
nodes = builder._build_l3_from_paragraphs(["A", "B"], l1_i=1, l2_j=2)
|
||||
assert nodes[0].id == "l1_1_l2_2_l3_0"
|
||||
assert nodes[1].id == "l1_1_l2_2_l3_1"
|
||||
|
||||
def test_video_fields_are_none(self, builder):
|
||||
"""文本模式下 frame_path 和 timestamp 应为 None。"""
|
||||
nodes = builder._build_l3_from_paragraphs(["测试段落"], l1_i=0, l2_j=0)
|
||||
assert nodes[0].frame_path is None
|
||||
assert nodes[0].timestamp is None
|
||||
|
||||
def test_embedding_dtype(self, builder):
|
||||
"""嵌入向量应为 float32。"""
|
||||
nodes = builder._build_l3_from_paragraphs(["A"], l1_i=0, l2_j=0)
|
||||
assert nodes[0].embedding.dtype == np.float32
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# build() — 完整树结构验证
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBuild:
|
||||
"""测试 build() 完整流程与 TreeIndex 结构。"""
|
||||
|
||||
def test_returns_tree_index(self, builder):
|
||||
"""build() 应返回 TreeIndex 实例。"""
|
||||
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
|
||||
assert isinstance(index, TreeIndex)
|
||||
|
||||
def test_roots_nonempty(self, builder):
|
||||
"""roots 列表不为空。"""
|
||||
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
|
||||
assert len(index.roots) > 0
|
||||
|
||||
def test_l1_has_l2_children(self, builder):
|
||||
"""每个 L1 节点至少有一个 L2 子节点。"""
|
||||
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
|
||||
for l1 in index.roots:
|
||||
assert isinstance(l1, L1Node)
|
||||
assert len(l1.children) > 0
|
||||
|
||||
def test_l2_has_l3_children(self, builder):
|
||||
"""每个 L2 节点至少有一个 L3 子节点。"""
|
||||
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
|
||||
for l1 in index.roots:
|
||||
for l2 in l1.children:
|
||||
assert isinstance(l2, L2Node)
|
||||
assert len(l2.children) > 0
|
||||
|
||||
def test_l3_are_leaf_nodes(self, builder):
|
||||
"""L3 节点是叶子层,类型为 L3Node。"""
|
||||
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
|
||||
for l1 in index.roots:
|
||||
for l2 in l1.children:
|
||||
for l3 in l2.children:
|
||||
assert isinstance(l3, L3Node)
|
||||
|
||||
def test_index_meta_modality(self, builder):
|
||||
"""IndexMeta.modality 应为 'text'。"""
|
||||
index = builder.build(_SAMPLE_TOC_TEXT, source_path="doc.txt")
|
||||
assert index.metadata.modality == "text"
|
||||
|
||||
def test_index_meta_source_path(self, builder):
|
||||
"""IndexMeta.source_path 应与传入参数一致。"""
|
||||
index = builder.build(_SAMPLE_TOC_TEXT, source_path="my_doc.txt")
|
||||
assert index.metadata.source_path == "my_doc.txt"
|
||||
|
||||
def test_index_meta_embed_dim(self, builder):
|
||||
"""IndexMeta.embed_dim 应与 embed.dim 一致。"""
|
||||
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
|
||||
assert index.metadata.embed_dim == _EMBED_DIM
|
||||
|
||||
def test_embedding_shapes(self, builder):
|
||||
"""所有节点 embedding 形状应为 (dim,)。"""
|
||||
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
|
||||
for l1 in index.roots:
|
||||
assert l1.embedding.shape == (_EMBED_DIM,)
|
||||
for l2 in l1.children:
|
||||
assert l2.embedding.shape == (_EMBED_DIM,)
|
||||
for l3 in l2.children:
|
||||
assert l3.embedding.shape == (_EMBED_DIM,)
|
||||
|
||||
def test_l1_summary_nonempty(self, builder):
|
||||
"""L1 summary 不为空。"""
|
||||
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
|
||||
for l1 in index.roots:
|
||||
assert l1.summary.strip() != ""
|
||||
|
||||
def test_l2_description_nonempty(self, builder):
|
||||
"""L2 description 不为空。"""
|
||||
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
|
||||
for l1 in index.roots:
|
||||
for l2 in l1.children:
|
||||
assert l2.description.strip() != ""
|
||||
|
||||
def test_plain_text_build(self, builder):
|
||||
"""无 ToC 的纯段落文本也能正确构建。"""
|
||||
# 修正 mock:返回合法 JSON
|
||||
paragraphs = ["信息检索是计算机科学的重要分支。", "传统方法以 TF-IDF 为代表。", "近年来稠密检索方法兴起。"]
|
||||
builder.llm.chat.return_value = json.dumps(paragraphs, ensure_ascii=False)
|
||||
builder.llm.batch_chat.side_effect = lambda prompts: [
|
||||
f"摘要_{i}" for i in range(len(prompts))
|
||||
]
|
||||
index = builder.build(_SAMPLE_PLAIN_TEXT, source_path="plain.txt")
|
||||
assert len(index.roots) > 0
|
||||
|
||||
def test_empty_text_raises(self, builder):
|
||||
"""空文本应抛出 ValueError(通过 ensure 检查)。"""
|
||||
with pytest.raises((ValueError, AssertionError)):
|
||||
builder.build(" ", source_path="empty.txt")
|
||||
|
||||
def test_batch_chat_called_once(self, builder):
|
||||
"""build() 应调用 batch_chat() 一次(所有 L2 并发处理)。"""
|
||||
builder.llm.batch_chat.reset_mock()
|
||||
builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
|
||||
builder.llm.batch_chat.assert_called_once()
|
||||
|
||||
def test_l1_node_ids_unique(self, builder):
|
||||
"""所有 L1 节点 ID 唯一。"""
|
||||
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
|
||||
ids = [l1.id for l1 in index.roots]
|
||||
assert len(ids) == len(set(ids))
|
||||
|
||||
def test_l2_node_ids_unique(self, builder):
|
||||
"""所有 L2 节点 ID 全局唯一。"""
|
||||
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
|
||||
ids = [l2.id for l1 in index.roots for l2 in l1.children]
|
||||
assert len(ids) == len(set(ids))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MD 输出文件
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMDOutput:
|
||||
"""测试 MD 输出文件生成(Agent 测试规范)。"""
|
||||
|
||||
def test_md_output_saved(self, builder):
|
||||
"""build() 后应能保存 Markdown 执行记录文件。"""
|
||||
index = builder.build(_SAMPLE_TOC_TEXT, source_path="test.txt")
|
||||
|
||||
# 统计树结构信息
|
||||
total_l2 = sum(len(r.children) for r in index.roots)
|
||||
total_l3 = sum(len(l2.children) for r in index.roots for l2 in r.children)
|
||||
|
||||
# 构造 MD 内容
|
||||
l2_details = []
|
||||
for l1 in index.roots:
|
||||
for l2 in l1.children:
|
||||
l2_details.append(f" - {l2.id}: {l2.description[:40]}...")
|
||||
|
||||
md_content = f"""\
|
||||
# Agent 测试: TextTreeBuilder.build
|
||||
## 任务: 长文本 → TreeIndex
|
||||
|
||||
## 输入信息
|
||||
- **文本长度**: {len(_SAMPLE_TOC_TEXT)} 字符
|
||||
- **是否含 ToC**: {builder._detect_toc(_SAMPLE_TOC_TEXT)}
|
||||
- **source_path**: test.txt
|
||||
- **max_paragraphs_per_l2**: {builder.config.max_paragraphs_per_l2}
|
||||
- **embed_dim**: {_EMBED_DIM}
|
||||
|
||||
## Step 1: 结构切分
|
||||
- **策略**: {'ToC 正则切分' if builder._detect_toc(_SAMPLE_TOC_TEXT) else 'LLM 语义分段'}
|
||||
- **L1 数量**: {len(index.roots)}
|
||||
- **各 L1 的 L2 数**: {[len(r.children) for r in index.roots]}
|
||||
|
||||
## Step 2: L2 先行(批量 LLM)
|
||||
- **L2 节点总数**: {total_l2}
|
||||
- **调用方式**: batch_chat()(并发生成所有 L2 摘要)
|
||||
- **L2 描述示例**:
|
||||
{chr(10).join(l2_details[:5])}
|
||||
|
||||
## Step 3: L3 向下(原始段落直接复用)
|
||||
- **L3 节点总数**: {total_l3}
|
||||
- **L3 特性**: description == raw_content(无 LLM 调用)
|
||||
|
||||
## Step 4: L1 向上(聚合摘要)
|
||||
- **L1 摘要示例**:
|
||||
{chr(10).join(f' - {r.id}: {r.summary[:60]}...' for r in index.roots)}
|
||||
|
||||
## 最终结果
|
||||
- **roots 数量**: {len(index.roots)}
|
||||
- **总 L2 节点**: {total_l2}
|
||||
- **总 L3 节点**: {total_l3}
|
||||
- **modality**: {index.metadata.modality}
|
||||
- **embed_dim**: {index.metadata.embed_dim}
|
||||
- **embedding shape 检查**: {'PASS' if all(r.embedding.shape == (_EMBED_DIM,) for r in index.roots) else 'FAIL'}
|
||||
- **L3 description==raw_content 检查**: {'PASS' if all(l3.description == l3.raw_content for r in index.roots for l2 in r.children for l3 in l2.children) else 'FAIL'}
|
||||
"""
|
||||
|
||||
out_path = _save_md_output("build_toc", md_content)
|
||||
assert os.path.exists(out_path)
|
||||
print(f"\n[MD 输出] 已保存到: {out_path}")
|
||||
@@ -0,0 +1,212 @@
|
||||
"""
|
||||
tree_index 单元测试
|
||||
===================
|
||||
覆盖: 嵌入矩阵提取、节点访问、边界检查、序列化往返、空树处理。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from video_tree_trm.tree_index import (
|
||||
IndexMeta,
|
||||
L1Node,
|
||||
L2Node,
|
||||
L3Node,
|
||||
TreeIndex,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
EMBED_DIM = 64
|
||||
|
||||
|
||||
def _make_embed(seed: int = 0) -> np.ndarray:
|
||||
"""生成固定种子的随机嵌入向量 [D]。"""
|
||||
rng = np.random.RandomState(seed)
|
||||
return rng.randn(EMBED_DIM).astype(np.float32)
|
||||
|
||||
|
||||
def _make_meta() -> IndexMeta:
|
||||
return IndexMeta(
|
||||
source_path="test.mp4",
|
||||
modality="video",
|
||||
embed_model="test-model",
|
||||
embed_dim=EMBED_DIM,
|
||||
)
|
||||
|
||||
|
||||
def _make_tree() -> TreeIndex:
|
||||
"""构建一棵 2 x 2 x 3 的测试树。"""
|
||||
meta = _make_meta()
|
||||
roots = []
|
||||
seed = 0
|
||||
for i in range(2):
|
||||
l2_nodes = []
|
||||
for j in range(2):
|
||||
l3_nodes = [
|
||||
L3Node(
|
||||
id=f"l3_{i}_{j}_{k}",
|
||||
description=f"帧描述 {i}-{j}-{k}",
|
||||
embedding=_make_embed(seed := seed + 1),
|
||||
)
|
||||
for k in range(3)
|
||||
]
|
||||
l2_nodes.append(
|
||||
L2Node(
|
||||
id=f"l2_{i}_{j}",
|
||||
description=f"片段描述 {i}-{j}",
|
||||
embedding=_make_embed(seed := seed + 1),
|
||||
children=l3_nodes,
|
||||
)
|
||||
)
|
||||
roots.append(
|
||||
L1Node(
|
||||
id=f"l1_{i}",
|
||||
summary=f"摘要 {i}",
|
||||
embedding=_make_embed(seed := seed + 1),
|
||||
children=l2_nodes,
|
||||
)
|
||||
)
|
||||
return TreeIndex(metadata=meta, roots=roots)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试: 嵌入矩阵提取
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEmbeddings:
|
||||
"""嵌入矩阵提取方法测试。"""
|
||||
|
||||
def test_l1_embeddings_shape(self) -> None:
|
||||
"""l1_embeddings() 返回 [N1, D]。"""
|
||||
tree = _make_tree()
|
||||
emb = tree.l1_embeddings()
|
||||
assert emb.shape == (2, EMBED_DIM)
|
||||
assert emb.dtype == np.float32
|
||||
|
||||
def test_l2_embeddings_of_shape(self) -> None:
|
||||
"""l2_embeddings_of(idx) 返回 [N2, D]。"""
|
||||
tree = _make_tree()
|
||||
emb = tree.l2_embeddings_of(0)
|
||||
assert emb.shape == (2, EMBED_DIM)
|
||||
assert emb.dtype == np.float32
|
||||
|
||||
def test_l3_embeddings_of_shape(self) -> None:
|
||||
"""l3_embeddings_of(l1, l2) 返回 [N3, D]。"""
|
||||
tree = _make_tree()
|
||||
emb = tree.l3_embeddings_of(0, 1)
|
||||
assert emb.shape == (3, EMBED_DIM)
|
||||
assert emb.dtype == np.float32
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试: 节点访问
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestGetNode:
|
||||
"""节点访问方法测试。"""
|
||||
|
||||
def test_get_node(self) -> None:
|
||||
"""正确返回目标 L3Node。"""
|
||||
tree = _make_tree()
|
||||
node = tree.get_node(1, 0, 2)
|
||||
assert isinstance(node, L3Node)
|
||||
assert node.id == "l3_1_0_2"
|
||||
assert node.description == "帧描述 1-0-2"
|
||||
|
||||
def test_get_node_boundary_error(self) -> None:
|
||||
"""越界索引抛出 IndexError。"""
|
||||
tree = _make_tree()
|
||||
with pytest.raises(IndexError):
|
||||
tree.get_node(5, 0, 0)
|
||||
with pytest.raises(IndexError):
|
||||
tree.get_node(0, 5, 0)
|
||||
with pytest.raises(IndexError):
|
||||
tree.get_node(0, 0, 5)
|
||||
|
||||
def test_get_node_negative_index_error(self) -> None:
|
||||
"""负数索引抛出 IndexError。"""
|
||||
tree = _make_tree()
|
||||
with pytest.raises(IndexError):
|
||||
tree.get_node(-1, 0, 0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试: 序列化
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSerialization:
|
||||
"""pickle 序列化测试。"""
|
||||
|
||||
def test_save_load_roundtrip(self) -> None:
|
||||
"""pickle 序列化后反序列化,数据完整一致。"""
|
||||
tree = _make_tree()
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
path = os.path.join(tmpdir, "test.pkl")
|
||||
tree.save(path)
|
||||
loaded = TreeIndex.load(path)
|
||||
|
||||
# 元数据一致
|
||||
assert loaded.metadata.source_path == tree.metadata.source_path
|
||||
assert loaded.metadata.embed_dim == tree.metadata.embed_dim
|
||||
|
||||
# 结构一致
|
||||
assert len(loaded.roots) == len(tree.roots)
|
||||
for orig_l1, load_l1 in zip(tree.roots, loaded.roots):
|
||||
assert orig_l1.id == load_l1.id
|
||||
assert len(orig_l1.children) == len(load_l1.children)
|
||||
|
||||
# 嵌入一致
|
||||
np.testing.assert_array_equal(loaded.l1_embeddings(), tree.l1_embeddings())
|
||||
np.testing.assert_array_equal(
|
||||
loaded.l3_embeddings_of(0, 1), tree.l3_embeddings_of(0, 1)
|
||||
)
|
||||
|
||||
def test_load_nonexistent_file(self) -> None:
|
||||
"""加载不存在的文件抛出 FileNotFoundError。"""
|
||||
with pytest.raises(FileNotFoundError):
|
||||
TreeIndex.load("/tmp/nonexistent_tree_index_abc123.pkl")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试: 空树边界
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEmptyTree:
|
||||
"""空树边界情况测试。"""
|
||||
|
||||
def test_empty_tree_l1_embeddings(self) -> None:
|
||||
"""空树的 l1_embeddings() 返回 [0, D]。"""
|
||||
tree = TreeIndex(metadata=_make_meta(), roots=[])
|
||||
emb = tree.l1_embeddings()
|
||||
assert emb.shape == (0, EMBED_DIM)
|
||||
assert emb.dtype == np.float32
|
||||
|
||||
def test_empty_tree_get_node_raises(self) -> None:
|
||||
"""空树访问节点抛出 IndexError。"""
|
||||
tree = TreeIndex(metadata=_make_meta(), roots=[])
|
||||
with pytest.raises(IndexError):
|
||||
tree.get_node(0, 0, 0)
|
||||
|
||||
def test_l2_embeddings_of_boundary(self) -> None:
|
||||
"""l2_embeddings_of 越界抛出 ValueError。"""
|
||||
tree = _make_tree()
|
||||
with pytest.raises(ValueError):
|
||||
tree.l2_embeddings_of(10)
|
||||
|
||||
def test_l3_embeddings_of_boundary(self) -> None:
|
||||
"""l3_embeddings_of 越界抛出 ValueError。"""
|
||||
tree = _make_tree()
|
||||
with pytest.raises(ValueError):
|
||||
tree.l3_embeddings_of(0, 10)
|
||||
@@ -0,0 +1,504 @@
|
||||
"""
|
||||
VideoTreeBuilder 单元测试
|
||||
=========================
|
||||
覆盖视频树构建的各个子方法和完整流程。
|
||||
|
||||
测试策略:
|
||||
- mock cv2.VideoCapture 避免依赖真实视频(_segment_video 等)
|
||||
- 使用 cv2 合成小视频进行帧提取和集成测试(tiny_video fixture)
|
||||
- mock VLM/embed 依赖,隔离外部 API
|
||||
- 测试 L3 批量降级路径(JSON 解析失败时退回逐帧调用)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from video_tree_trm.config import TreeConfig
|
||||
from video_tree_trm.tree_index import L1Node, L2Node, L3Node
|
||||
from video_tree_trm.video_tree_builder import VideoTreeBuilder
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tree_config(tmp_path: Path) -> TreeConfig:
|
||||
"""构建测试用 TreeConfig,cache_dir 指向临时目录。"""
|
||||
return TreeConfig(
|
||||
max_paragraphs_per_l2=5,
|
||||
l1_segment_duration=30.0,
|
||||
l2_clip_duration=10.0,
|
||||
l3_fps=1.0,
|
||||
l2_representative_frames=3,
|
||||
cache_dir=str(tmp_path / "cache"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_embed() -> MagicMock:
|
||||
"""返回 mock 嵌入模型,embed() 返回全 1 向量。"""
|
||||
embed = MagicMock()
|
||||
embed.dim = 4
|
||||
embed._model_name = "mock-embed"
|
||||
|
||||
def _embed(texts):
|
||||
"""根据输入类型返回 [1, D] 或 [N, D] 的 float32 数组。"""
|
||||
if isinstance(texts, str):
|
||||
return np.ones((1, 4), dtype=np.float32)
|
||||
n = len(texts)
|
||||
return np.ones((n, 4), dtype=np.float32)
|
||||
|
||||
embed.embed.side_effect = _embed
|
||||
return embed
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_vlm() -> MagicMock:
|
||||
"""返回 mock VLM 客户端。"""
|
||||
vlm = MagicMock()
|
||||
vlm.chat.return_value = "这是一段精彩的视频内容摘要。"
|
||||
vlm.chat_with_images.return_value = "这帧画面中有人物在移动。"
|
||||
return vlm
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def builder(
|
||||
mock_embed: MagicMock, mock_vlm: MagicMock, tree_config: TreeConfig
|
||||
) -> VideoTreeBuilder:
|
||||
"""构建测试用 VideoTreeBuilder 实例。"""
|
||||
return VideoTreeBuilder(
|
||||
embed_model=mock_embed,
|
||||
vlm=mock_vlm,
|
||||
config=tree_config,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tiny_video(tmp_path: Path) -> str:
|
||||
"""用 cv2 生成 30 帧合成彩色视频(10fps,时长 3 秒),返回路径。
|
||||
|
||||
视频规格:
|
||||
- 分辨率: 64×48
|
||||
- 帧率: 10fps
|
||||
- 时长: 3 秒(30 帧)
|
||||
- 内容: 随机彩色帧
|
||||
"""
|
||||
video_path = str(tmp_path / "tiny.mp4")
|
||||
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
||||
writer = cv2.VideoWriter(video_path, fourcc, 10.0, (64, 48))
|
||||
for _ in range(30):
|
||||
frame = np.random.randint(0, 256, (48, 64, 3), dtype=np.uint8)
|
||||
writer.write(frame)
|
||||
writer.release()
|
||||
return video_path
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试:_segment_video — 时间切分
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_segment_video_fixed_step(builder: VideoTreeBuilder, tmp_path: Path) -> None:
|
||||
"""mock cv2.VideoCapture(总时长=60s,l1_segment_duration=30s),
|
||||
验证切分出 2 个均等 L1 区间:(0,30),(30,60)。"""
|
||||
mock_cap = MagicMock()
|
||||
mock_cap.isOpened.return_value = True
|
||||
mock_cap.get.side_effect = lambda prop: (
|
||||
10.0 if prop == cv2.CAP_PROP_FPS else 600.0 # 600帧/10fps = 60s
|
||||
)
|
||||
|
||||
with patch("video_tree_trm.video_tree_builder.cv2.VideoCapture", return_value=mock_cap):
|
||||
ranges = builder._segment_video("fake.mp4")
|
||||
|
||||
assert len(ranges) == 2
|
||||
assert ranges[0] == (0.0, 30.0)
|
||||
assert ranges[1] == (30.0, 60.0)
|
||||
|
||||
|
||||
def test_segment_video_uneven(builder: VideoTreeBuilder) -> None:
|
||||
"""总时长不能被 l1_segment_duration 整除时,最后一段应短于步长。"""
|
||||
mock_cap = MagicMock()
|
||||
mock_cap.isOpened.return_value = True
|
||||
# 75s = 750帧 / 10fps
|
||||
mock_cap.get.side_effect = lambda prop: (
|
||||
10.0 if prop == cv2.CAP_PROP_FPS else 750.0
|
||||
)
|
||||
|
||||
with patch("video_tree_trm.video_tree_builder.cv2.VideoCapture", return_value=mock_cap):
|
||||
ranges = builder._segment_video("fake.mp4")
|
||||
|
||||
assert len(ranges) == 3
|
||||
assert ranges[0] == (0.0, 30.0)
|
||||
assert ranges[1] == (30.0, 60.0)
|
||||
assert abs(ranges[2][0] - 60.0) < 1e-6
|
||||
assert abs(ranges[2][1] - 75.0) < 1e-6
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试:_get_l2_clips — L2 切分
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_get_l2_clips_even(builder: VideoTreeBuilder) -> None:
|
||||
"""l1=(0,30),l2_duration=10 → 3 clips 均等:(0,10),(10,20),(20,30)。"""
|
||||
clips = builder._get_l2_clips((0.0, 30.0))
|
||||
assert len(clips) == 3
|
||||
assert clips[0] == (0.0, 10.0)
|
||||
assert clips[1] == (10.0, 20.0)
|
||||
assert clips[2] == (20.0, 30.0)
|
||||
|
||||
|
||||
def test_get_l2_clips_uneven(builder: VideoTreeBuilder) -> None:
|
||||
"""l1=(0,25),l2_duration=10 → 3 clips,最后一段为 5s。"""
|
||||
clips = builder._get_l2_clips((0.0, 25.0))
|
||||
assert len(clips) == 3
|
||||
assert clips[2] == (20.0, 25.0)
|
||||
|
||||
|
||||
def test_get_l2_clips_shorter_than_step(builder: VideoTreeBuilder) -> None:
|
||||
"""L1 区间短于 l2_clip_duration 时,返回 1 个 clip。"""
|
||||
clips = builder._get_l2_clips((0.0, 5.0))
|
||||
assert len(clips) == 1
|
||||
assert clips[0] == (0.0, 5.0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试:_extract_frames — 帧提取
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_extract_frames_saves_files(
|
||||
builder: VideoTreeBuilder, tiny_video: str, tmp_path: Path
|
||||
) -> None:
|
||||
"""使用真实合成视频(3s),提取 1fps 帧,验证返回路径和文件存在。"""
|
||||
frames = builder._extract_frames(tiny_video, (0.0, 3.0), fps=1.0)
|
||||
|
||||
assert len(frames) >= 1
|
||||
for frame_path, ts in frames:
|
||||
assert os.path.isfile(frame_path), f"帧文件不存在: {frame_path}"
|
||||
assert ts >= 0.0
|
||||
|
||||
|
||||
def test_extract_frames_cache_reuse(
|
||||
builder: VideoTreeBuilder, tiny_video: str
|
||||
) -> None:
|
||||
"""第二次提取同一区间时,帧文件应直接复用(不重复写磁盘)。"""
|
||||
frames1 = builder._extract_frames(tiny_video, (0.0, 2.0), fps=1.0)
|
||||
assert len(frames1) >= 1
|
||||
|
||||
# 记录文件修改时间
|
||||
mtimes_before = [os.path.getmtime(fp) for fp, _ in frames1]
|
||||
|
||||
frames2 = builder._extract_frames(tiny_video, (0.0, 2.0), fps=1.0)
|
||||
mtimes_after = [os.path.getmtime(fp) for fp, _ in frames2]
|
||||
|
||||
assert frames1 == frames2
|
||||
assert mtimes_before == mtimes_after, "缓存帧文件被重复写入"
|
||||
|
||||
|
||||
def test_extract_frames_empty_range(
|
||||
builder: VideoTreeBuilder, tiny_video: str
|
||||
) -> None:
|
||||
"""时间范围内无有效时间戳时,返回空列表。"""
|
||||
# 起始=结束,无时间戳
|
||||
frames = builder._extract_frames(tiny_video, (1.0, 1.0), fps=1.0)
|
||||
assert frames == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试:_build_l2_video — L2 节点构建
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_build_l2_video_node_structure(
|
||||
builder: VideoTreeBuilder, tiny_video: str, mock_vlm: MagicMock
|
||||
) -> None:
|
||||
"""验证 L2Node 字段:description 非空、embedding shape 正确、time_range 正确。"""
|
||||
mock_vlm.chat_with_images.return_value = "片段展示了室内场景的变化。"
|
||||
l2_node = builder._build_l2_video(tiny_video, (0.0, 2.0), "l1_0_l2_0")
|
||||
|
||||
assert isinstance(l2_node, L2Node)
|
||||
assert l2_node.id == "l1_0_l2_0"
|
||||
assert len(l2_node.description) > 0
|
||||
assert l2_node.embedding.shape == (4,)
|
||||
assert l2_node.embedding.dtype == np.float32
|
||||
assert l2_node.time_range == (0.0, 2.0)
|
||||
assert l2_node.children == [] # 调用方填充
|
||||
|
||||
|
||||
def test_build_l2_video_representative_frames_count(
|
||||
builder: VideoTreeBuilder, tiny_video: str, mock_vlm: MagicMock
|
||||
) -> None:
|
||||
"""验证 VLM 被调用时传入的图像数不超过 l2_representative_frames。"""
|
||||
mock_vlm.chat_with_images.return_value = "描述内容。"
|
||||
builder._build_l2_video(tiny_video, (0.0, 3.0), "l1_0_l2_0")
|
||||
|
||||
call_args = mock_vlm.chat_with_images.call_args
|
||||
images_passed = call_args.kwargs.get("images", call_args.args[1] if len(call_args.args) > 1 else [])
|
||||
assert len(images_passed) <= builder.config.l2_representative_frames
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试:_build_l3_video — L3 节点构建
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_frames(n: int, tmp_path: Path) -> List[tuple]:
|
||||
"""创建 n 个临时 JPEG 帧文件,返回 [(path, ts), ...]。"""
|
||||
frames = []
|
||||
frame_dir = tmp_path / "frames"
|
||||
frame_dir.mkdir(exist_ok=True)
|
||||
for i in range(n):
|
||||
frame_path = str(frame_dir / f"frame_{i}.jpg")
|
||||
img = np.zeros((48, 64, 3), dtype=np.uint8)
|
||||
cv2.imwrite(frame_path, img)
|
||||
frames.append((frame_path, float(i)))
|
||||
return frames
|
||||
|
||||
|
||||
def test_build_l3_video_batch_success(
|
||||
builder: VideoTreeBuilder, mock_vlm: MagicMock, tmp_path: Path
|
||||
) -> None:
|
||||
"""mock VLM 返回合法 JSON 数组,验证 L3Node 列表结构正确。"""
|
||||
frames = _make_frames(2, tmp_path)
|
||||
mock_vlm.chat_with_images.return_value = json.dumps(["帧1描述内容", "帧2描述内容"])
|
||||
|
||||
nodes = builder._build_l3_video(frames, "片段整体描述", l1_i=0, l2_j=0)
|
||||
|
||||
assert len(nodes) == 2
|
||||
for k, node in enumerate(nodes):
|
||||
assert isinstance(node, L3Node)
|
||||
assert node.id == f"l1_0_l2_0_l3_{k}"
|
||||
assert len(node.description) > 0
|
||||
assert node.embedding.shape == (4,)
|
||||
assert node.embedding.dtype == np.float32
|
||||
assert node.frame_path == frames[k][0]
|
||||
assert node.timestamp == float(k)
|
||||
assert node.raw_content is None
|
||||
|
||||
|
||||
def test_build_l3_video_batch_fallback(
|
||||
builder: VideoTreeBuilder, mock_vlm: MagicMock, tmp_path: Path
|
||||
) -> None:
|
||||
"""mock VLM 第一次返回非 JSON 字符串,验证降级逐帧调用(call_count == n+1)。
|
||||
|
||||
第一次 = 批量调用(失败),后 n 次 = 逐帧调用。
|
||||
"""
|
||||
n = 3
|
||||
frames = _make_frames(n, tmp_path)
|
||||
# 第一次返回无效 JSON,后续逐帧返回正常描述
|
||||
mock_vlm.chat_with_images.side_effect = (
|
||||
["这不是一个JSON数组,无法解析"] + [f"第{i}帧描述" for i in range(n)]
|
||||
)
|
||||
|
||||
nodes = builder._build_l3_video(frames, "片段整体描述", l1_i=0, l2_j=0)
|
||||
|
||||
assert len(nodes) == n
|
||||
# 1次批量 + n次逐帧
|
||||
assert mock_vlm.chat_with_images.call_count == n + 1
|
||||
for node in nodes:
|
||||
assert len(node.description) > 0
|
||||
|
||||
|
||||
def test_build_l3_video_json_length_mismatch_fallback(
|
||||
builder: VideoTreeBuilder, mock_vlm: MagicMock, tmp_path: Path
|
||||
) -> None:
|
||||
"""VLM 返回 JSON 但长度不匹配时,也应降级逐帧调用。"""
|
||||
n = 3
|
||||
frames = _make_frames(n, tmp_path)
|
||||
# 只返回 2 项,但期望 3 项
|
||||
mock_vlm.chat_with_images.side_effect = (
|
||||
[json.dumps(["描述1", "描述2"])] + [f"帧{i}" for i in range(n)]
|
||||
)
|
||||
|
||||
nodes = builder._build_l3_video(frames, "片段描述", l1_i=0, l2_j=0)
|
||||
|
||||
assert len(nodes) == n
|
||||
assert mock_vlm.chat_with_images.call_count == n + 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试:_build_l1_video — L1 节点构建
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_build_l1_video_node_structure(
|
||||
builder: VideoTreeBuilder, mock_vlm: MagicMock, mock_embed: MagicMock
|
||||
) -> None:
|
||||
"""验证 L1Node 字段:summary 非空、time_range 正确、children 已赋值。"""
|
||||
mock_vlm.chat.return_value = "这段视频涵盖了户外活动和室内场景的切换。"
|
||||
|
||||
l2_children = [
|
||||
L2Node(
|
||||
id=f"l1_0_l2_{j}",
|
||||
description=f"L2描述{j}",
|
||||
embedding=np.ones(4, dtype=np.float32),
|
||||
time_range=(j * 10.0, (j + 1) * 10.0),
|
||||
)
|
||||
for j in range(3)
|
||||
]
|
||||
|
||||
l1_node = builder._build_l1_video(l2_children, "l1_0", (0.0, 30.0))
|
||||
|
||||
assert isinstance(l1_node, L1Node)
|
||||
assert l1_node.id == "l1_0"
|
||||
assert len(l1_node.summary) > 0
|
||||
assert l1_node.time_range == (0.0, 30.0)
|
||||
assert l1_node.children is l2_children
|
||||
assert l1_node.embedding.shape == (4,)
|
||||
assert l1_node.embedding.dtype == np.float32
|
||||
|
||||
|
||||
def test_build_l1_video_prompt_contains_l2_descriptions(
|
||||
builder: VideoTreeBuilder, mock_vlm: MagicMock
|
||||
) -> None:
|
||||
"""验证 L1 摘要的 prompt 包含所有 L2 描述文本。"""
|
||||
mock_vlm.chat.return_value = "综合摘要内容。"
|
||||
l2_descriptions = ["片段A描述", "片段B描述", "片段C描述"]
|
||||
l2_children = [
|
||||
L2Node(
|
||||
id=f"l1_0_l2_{j}",
|
||||
description=desc,
|
||||
embedding=np.ones(4, dtype=np.float32),
|
||||
time_range=(j * 10.0, (j + 1) * 10.0),
|
||||
)
|
||||
for j, desc in enumerate(l2_descriptions)
|
||||
]
|
||||
|
||||
builder._build_l1_video(l2_children, "l1_0", (0.0, 30.0))
|
||||
|
||||
call_prompt = mock_vlm.chat.call_args.args[0]
|
||||
for desc in l2_descriptions:
|
||||
assert desc in call_prompt, f"L2 描述 '{desc}' 未出现在 L1 prompt 中"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 测试:build 完整流程(集成测试,mock VLM)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_build_full_integration(
|
||||
builder: VideoTreeBuilder,
|
||||
tiny_video: str,
|
||||
mock_vlm: MagicMock,
|
||||
mock_embed: MagicMock,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""用合成视频(3s)+ mock VLM 验证完整 TreeIndex 三层结构。
|
||||
|
||||
配置:l1_segment_duration=2s,l2_clip_duration=1s
|
||||
预期:至少 1 个 L1,每 L1 至少 1 个 L2,每 L2 至少 1 个 L3。
|
||||
"""
|
||||
# 调整 config 使 3s 视频能切出多个节点
|
||||
builder.config.l1_segment_duration = 2.0
|
||||
builder.config.l2_clip_duration = 1.0
|
||||
|
||||
# VLM 批量调用返回 JSON 数组(按帧数动态生成)
|
||||
def vlm_side_effect(prompt, images=None):
|
||||
if images and len(images) > 1:
|
||||
# L3 批量调用:返回 JSON
|
||||
return json.dumps([f"帧{i}描述" for i in range(len(images))])
|
||||
return "这是 VLM 的描述文本。"
|
||||
|
||||
mock_vlm.chat_with_images.side_effect = vlm_side_effect
|
||||
mock_vlm.chat.return_value = "L1 整体摘要内容。"
|
||||
|
||||
index = builder.build(tiny_video)
|
||||
|
||||
# 验证元数据
|
||||
assert index.metadata.modality == "video"
|
||||
assert index.metadata.source_path == tiny_video
|
||||
assert index.metadata.embed_dim == 4
|
||||
|
||||
# 验证三层结构非空
|
||||
assert len(index.roots) >= 1, "应有至少 1 个 L1 节点"
|
||||
for l1 in index.roots:
|
||||
assert len(l1.children) >= 1, f"L1 {l1.id} 应有至少 1 个 L2 子节点"
|
||||
assert l1.time_range is not None
|
||||
for l2 in l1.children:
|
||||
assert len(l2.children) >= 1, f"L2 {l2.id} 应有至少 1 个 L3 子节点"
|
||||
assert l2.time_range is not None
|
||||
for l3 in l2.children:
|
||||
assert l3.frame_path is not None
|
||||
assert l3.timestamp is not None
|
||||
assert l3.embedding.shape == (4,)
|
||||
|
||||
|
||||
def test_build_saves_output_md(
|
||||
builder: VideoTreeBuilder,
|
||||
tiny_video: str,
|
||||
mock_vlm: MagicMock,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""构建完成后,将执行摘要保存为 Markdown(CLAUDE.md 规范)。"""
|
||||
builder.config.l1_segment_duration = 2.0
|
||||
builder.config.l2_clip_duration = 1.0
|
||||
|
||||
def vlm_side_effect(prompt, images=None):
|
||||
if images and len(images) > 1:
|
||||
return json.dumps([f"帧{i}描述" for i in range(len(images))])
|
||||
return "VLM 描述内容。"
|
||||
|
||||
mock_vlm.chat_with_images.side_effect = vlm_side_effect
|
||||
mock_vlm.chat.return_value = "L1 摘要。"
|
||||
|
||||
index = builder.build(tiny_video)
|
||||
|
||||
# 保存 Markdown 输出
|
||||
output_dir = Path(__file__).resolve().parent.parent / "outputs" / "video_tree_builder"
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
output_path = output_dir / f"build_video_{ts}.md"
|
||||
|
||||
total_l2 = sum(len(r.children) for r in index.roots)
|
||||
total_l3 = sum(len(l2.children) for r in index.roots for l2 in r.children)
|
||||
|
||||
lines = [
|
||||
f"# Agent 测试: test_build_saves_output_md",
|
||||
f"## 任务: VideoTreeBuilder.build() 完整流程验证",
|
||||
f"",
|
||||
f"## 输入",
|
||||
f"- 视频路径: `{tiny_video}`",
|
||||
f"- l1_segment_duration: {builder.config.l1_segment_duration}s",
|
||||
f"- l2_clip_duration: {builder.config.l2_clip_duration}s",
|
||||
f"- l3_fps: {builder.config.l3_fps}",
|
||||
f"",
|
||||
f"## 输出结构",
|
||||
f"- L1 节点数: {len(index.roots)}",
|
||||
f"- L2 节点数: {total_l2}",
|
||||
f"- L3 节点数: {total_l3}",
|
||||
f"- embed_dim: {index.metadata.embed_dim}",
|
||||
f"",
|
||||
f"## L1 详情",
|
||||
]
|
||||
for l1 in index.roots:
|
||||
lines.append(f"### {l1.id} (time_range={l1.time_range})")
|
||||
lines.append(f"- summary: {l1.summary[:80]}...")
|
||||
for l2 in l1.children:
|
||||
lines.append(
|
||||
f" - {l2.id} [{l2.time_range}]: {l2.description[:60]}... "
|
||||
f"({len(l2.children)} L3)"
|
||||
)
|
||||
lines += [
|
||||
"",
|
||||
"## 最终结果",
|
||||
"✅ TreeIndex 构建成功,三层结构完整。",
|
||||
f"",
|
||||
f"输出文件: `{output_path}`",
|
||||
]
|
||||
|
||||
output_path.write_text("\n".join(lines), encoding="utf-8")
|
||||
print(f"\n[测试输出] {output_path}")
|
||||
assert output_path.is_file()
|
||||
Reference in New Issue
Block a user