119 lines
3.7 KiB
Python
119 lines
3.7 KiB
Python
"""Markdown / JSON 输出写入器"""
|
||
|
||
import json
|
||
import os
|
||
from datetime import datetime
|
||
from typing import List
|
||
|
||
from models import Chunk
|
||
|
||
|
||
class MarkdownWriter:
|
||
"""将 Chunk 列表写入 Markdown 文件"""
|
||
|
||
def write(
|
||
self,
|
||
chunks: List[Chunk],
|
||
output_path: str,
|
||
source_file: str,
|
||
delimiter: str = "---",
|
||
) -> None:
|
||
"""
|
||
写入 Markdown 文件。
|
||
- 自动创建输出目录
|
||
- 文件开头添加元信息注释（源文件名、处理时间、Chunk 总数）
|
||
- 每个 Chunk 前添加摘要标题（## 级别）
|
||
- Chunk 之间用 delimiter 分隔
|
||
- 若输出文件已存在，覆盖并在控制台提示
|
||
"""
|
||
# 自动创建输出目录
|
||
output_dir = os.path.dirname(output_path)
|
||
if output_dir:
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
if os.path.exists(output_path):
|
||
print(f"警告: 输出文件已存在，将覆盖: {output_path}")
|
||
|
||
process_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
|
||
lines = []
|
||
|
||
# 元信息 HTML 注释
|
||
chunk_sizes = [len(c.content) for c in chunks]
|
||
avg_size = sum(chunk_sizes) // max(len(chunk_sizes), 1)
|
||
max_size = max(chunk_sizes) if chunk_sizes else 0
|
||
|
||
lines.append("<!-- ")
|
||
lines.append(f" 源文件: {source_file}")
|
||
lines.append(f" 处理时间: {process_time}")
|
||
lines.append(f" 分块总数: {len(chunks)}")
|
||
lines.append(f" 分块大小: 平均 {avg_size} 字, 最大 {max_size} 字")
|
||
lines.append(f" 上传提示: 使用「自定义分段」，分隔符设为 {delimiter}")
|
||
lines.append("-->")
|
||
|
||
for i, chunk in enumerate(chunks):
|
||
lines.append("")
|
||
# 标签 + 标题
|
||
if chunk.tag:
|
||
lines.append(f"## `{chunk.tag}` {chunk.title}")
|
||
else:
|
||
lines.append(f"## {chunk.title}")
|
||
lines.append("")
|
||
lines.append(chunk.content)
|
||
|
||
if i < len(chunks) - 1:
|
||
lines.append("")
|
||
lines.append(delimiter)
|
||
|
||
lines.append("")
|
||
|
||
with open(output_path, "w", encoding="utf-8") as f:
|
||
f.write("\n".join(lines))
|
||
|
||
|
||
class JsonWriter:
|
||
"""将 Chunk 列表写入 JSON 文件，方便导入向量数据库"""
|
||
|
||
def write(
|
||
self,
|
||
chunks: List[Chunk],
|
||
output_path: str,
|
||
source_file: str,
|
||
delimiter: str = "---",
|
||
) -> None:
|
||
"""
|
||
写入 JSON 文件。
|
||
- 自动创建输出目录
|
||
- 输出结构化的 JSON，包含元信息和分块列表
|
||
- 若输出文件已存在，覆盖并在控制台提示
|
||
"""
|
||
# 自动创建输出目录
|
||
output_dir = os.path.dirname(output_path)
|
||
if output_dir:
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# 将扩展名改为 .json
|
||
json_path = os.path.splitext(output_path)[0] + ".json"
|
||
|
||
if os.path.exists(json_path):
|
||
print(f"警告: 输出文件已存在，将覆盖: {json_path}")
|
||
|
||
data = {
|
||
"source_file": source_file,
|
||
"process_time": datetime.now().isoformat(),
|
||
"total_chunks": len(chunks),
|
||
"chunks": [
|
||
{
|
||
"index": i,
|
||
"tag": chunk.tag,
|
||
"title": chunk.title,
|
||
"content": chunk.content,
|
||
"char_count": len(chunk.content),
|
||
}
|
||
for i, chunk in enumerate(chunks)
|
||
],
|
||
}
|
||
|
||
with open(json_path, "w", encoding="utf-8") as f:
|
||
json.dump(data, f, ensure_ascii=False, indent=2)
|