42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
"""基准测试共享 fixture."""
|
|
import pytest
|
|
|
|
from src.core.config import ChunkConfig, EmbedConfig
|
|
from src.core.db import VectorDB
|
|
from src.core.embedder import create_embedder
|
|
from src.core.ingest import DocumentIngestor
|
|
from src.core.search import Searcher
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def benchmark_db(tmp_path_factory):
|
|
"""模块级共享 ChromaDB 实例."""
|
|
persist_dir = tmp_path_factory.mktemp("bench_data")
|
|
return VectorDB(persist_dir=str(persist_dir))
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def benchmark_embedder():
|
|
"""模块级共享 LocalEmbedder."""
|
|
config = EmbedConfig(mode="local", local_model="BAAI/bge-small-zh-v1.5")
|
|
return create_embedder(config)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def benchmark_searcher(benchmark_db, benchmark_embedder):
|
|
"""预填充数据的 Searcher."""
|
|
ingestor = DocumentIngestor(
|
|
benchmark_db, benchmark_embedder, "bench_collection",
|
|
chunk_config=ChunkConfig(max_size=1000, overlap=100),
|
|
)
|
|
for i in range(100):
|
|
content = (
|
|
f"# 文档{i}\n\n"
|
|
+ "\n\n".join(
|
|
f"第{j}段用于基准测试。关键词: Python, Rust, GPU, 向量数据库。"
|
|
for j in range(5)
|
|
)
|
|
)
|
|
ingestor.ingest_content(content, f"bench_{i}.md")
|
|
return Searcher(benchmark_db, benchmark_embedder, "bench_collection")
|