feat: 新增数据导出功能和 CLI export 命令(JSON/CSV)
This commit is contained in:
@@ -222,5 +222,30 @@ def stats(
|
|||||||
typer.echo(f" - {s}")
|
typer.echo(f" - {s}")
|
||||||
|
|
||||||
|
|
||||||
|
@app.command(help="导出 collection 数据为 JSON 或 CSV.")
|
||||||
|
def export(
|
||||||
|
output: Annotated[str, typer.Option("--output", "-o", help="输出文件路径")],
|
||||||
|
fmt: Annotated[
|
||||||
|
str, typer.Option("--format", "-f", help="导出格式: json | csv")
|
||||||
|
] = "json",
|
||||||
|
config: ConfigOpt = DEFAULT_CONFIG_PATH,
|
||||||
|
collection: CollectionOpt = None,
|
||||||
|
):
|
||||||
|
"""导出 collection 数据."""
|
||||||
|
_init_config(config)
|
||||||
|
state = get_state()
|
||||||
|
searcher = state.get_searcher(_resolve_collection(collection))
|
||||||
|
|
||||||
|
if fmt == "json":
|
||||||
|
searcher.export_json(file_path=output)
|
||||||
|
elif fmt == "csv":
|
||||||
|
searcher.export_csv(file_path=output)
|
||||||
|
else:
|
||||||
|
typer.echo(f"错误: 不支持的格式 '{fmt}',可选: json, csv", err=True)
|
||||||
|
raise typer.Exit(code=1)
|
||||||
|
|
||||||
|
typer.echo(f"[OK] 已导出到: {output}")
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
app()
|
app()
|
||||||
|
|||||||
@@ -0,0 +1,106 @@
|
|||||||
|
"""导出功能测试."""
|
||||||
|
import csv
|
||||||
|
import io
|
||||||
|
import json
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from src.core.search import Searcher
|
||||||
|
|
||||||
|
|
||||||
|
class FakeExportCollection:
|
||||||
|
def count(self):
|
||||||
|
return 2
|
||||||
|
|
||||||
|
def get(self, include=None):
|
||||||
|
return {
|
||||||
|
"ids": ["doc_0", "doc_1"],
|
||||||
|
"documents": ["内容A。\n\n段落B。", "内容C。"],
|
||||||
|
"metadatas": [
|
||||||
|
{"source_file": "a.md", "section_title": "标题A", "heading_level": 1, "chunk_index": 0},
|
||||||
|
{"source_file": "b.md", "section_title": "", "heading_level": 0, "chunk_index": 0},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class FakeExportDB:
|
||||||
|
def get_or_create_collection(self, name):
|
||||||
|
return FakeExportCollection()
|
||||||
|
|
||||||
|
def list_collections(self):
|
||||||
|
return []
|
||||||
|
|
||||||
|
|
||||||
|
class FakeExportEmbedder:
|
||||||
|
@property
|
||||||
|
def dimension(self):
|
||||||
|
return 4
|
||||||
|
|
||||||
|
def embed(self, texts):
|
||||||
|
return [[0.1, 0.2, 0.3, 0.4] for _ in texts]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def searcher():
|
||||||
|
db = FakeExportDB()
|
||||||
|
embedder = FakeExportEmbedder()
|
||||||
|
return Searcher(db, embedder, "test_export")
|
||||||
|
|
||||||
|
|
||||||
|
def test_export_json_returns_valid_json(searcher):
|
||||||
|
"""export_json 返回合法的 JSON 字符串."""
|
||||||
|
output = searcher.export_json()
|
||||||
|
data = json.loads(output)
|
||||||
|
assert isinstance(data, list)
|
||||||
|
assert len(data) == 2
|
||||||
|
|
||||||
|
|
||||||
|
def test_export_json_contains_all_fields(searcher):
|
||||||
|
"""导出包含所有必要字段."""
|
||||||
|
output = searcher.export_json()
|
||||||
|
data = json.loads(output)
|
||||||
|
first = data[0]
|
||||||
|
assert "id" in first
|
||||||
|
assert "content" in first
|
||||||
|
assert "source_file" in first
|
||||||
|
assert "section_title" in first
|
||||||
|
|
||||||
|
|
||||||
|
def test_export_csv_returns_valid_csv(searcher):
|
||||||
|
"""export_csv 返回合法的 CSV 字符串."""
|
||||||
|
output = searcher.export_csv()
|
||||||
|
reader = csv.DictReader(io.StringIO(output))
|
||||||
|
rows = list(reader)
|
||||||
|
assert len(rows) == 2
|
||||||
|
|
||||||
|
|
||||||
|
def test_export_csv_has_header(searcher):
|
||||||
|
"""CSV 包含表头."""
|
||||||
|
output = searcher.export_csv()
|
||||||
|
reader = csv.DictReader(io.StringIO(output))
|
||||||
|
assert reader.fieldnames is not None
|
||||||
|
assert "content" in reader.fieldnames
|
||||||
|
assert "source_file" in reader.fieldnames
|
||||||
|
|
||||||
|
|
||||||
|
def test_export_empty_collection():
|
||||||
|
"""空 collection 导出空列表."""
|
||||||
|
|
||||||
|
class EmptyCollection:
|
||||||
|
def count(self):
|
||||||
|
return 0
|
||||||
|
def get(self, include=None):
|
||||||
|
return {"ids": [], "documents": [], "metadatas": []}
|
||||||
|
|
||||||
|
class EmptyDB:
|
||||||
|
def get_or_create_collection(self, name):
|
||||||
|
return EmptyCollection()
|
||||||
|
def list_collections(self):
|
||||||
|
return []
|
||||||
|
|
||||||
|
s = Searcher(EmptyDB(), FakeExportEmbedder(), "empty")
|
||||||
|
json_out = s.export_json()
|
||||||
|
assert json.loads(json_out) == []
|
||||||
|
csv_out = s.export_csv()
|
||||||
|
lines = csv_out.strip().split("\n")
|
||||||
|
assert len(lines) == 1 # 仅表头
|
||||||
Reference in New Issue
Block a user