kplam
e2e8e6829d
1. 新增将文档元数据存入L1向量库的功能 2. 新增文档文件下载API接口，支持路径安全校验 3. 后端文档详情接口新增文件信息返回字段 4. 管理后台页面新增原始文件信息展示与下载链接
308 lines
12 KiB
Python
308 lines
12 KiB
Python
"""文档入库模块
|
||
|
||
入库流程：文档输入 → 三级总结(Ollama) → 分类判定(L1总结) → 切分 chunk
|
||
→ 构建 L2/L3 大纲节点 → 批量向量化（dense + sparse）→ 写入 Qdrant 四层集合
|
||
|
||
L2/L3 大纲节点的构建策略见 _build_l2_nodes / _build_l3_nodes。
|
||
"""
|
||
|
||
import uuid
|
||
from collections.abc import Callable
|
||
from typing import Any
|
||
|
||
import structlog
|
||
|
||
from app.config import settings
|
||
from app.core.chunker import Chunker
|
||
from app.core.classifier import Classifier
|
||
from app.core.embeddings import EmbeddingService, create_embedding_service
|
||
from app.core.headings import Heading, parse_headings
|
||
from app.core.sparse import SparseEncoder
|
||
from app.core.summarizer import Summarizer
|
||
from app.models.document import ChunkModel, DocumentInput, DocumentSummary, IngestionResult, SummaryLevel
|
||
from app.models.knowledge import CategoryResult
|
||
from app.services.qdrant import COLLECTION_L2, COLLECTION_L3, QdrantService, SparseVectorTuple
|
||
|
||
logger = structlog.get_logger()
|
||
|
||
# IngestionError 阶段标识
|
||
STAGE_SUMMARIZE = "summarize"
|
||
STAGE_CLASSIFY = "classify"
|
||
STAGE_EMBED = "embed"
|
||
STAGE_QDRANT = "qdrant"
|
||
|
||
|
||
class IngestionError(Exception):
|
||
"""入库失败异常
|
||
|
||
携带失败阶段（stage）与已产出的总结（summary，如有），
|
||
Qdrant 写入失败时总结不丢，上层可按阶段重试。
|
||
"""
|
||
|
||
def __init__(self, stage: str, message: str, summary: DocumentSummary | None = None) -> None:
|
||
super().__init__(message)
|
||
self.stage = stage
|
||
self.summary = summary
|
||
|
||
|
||
def _heading_paths(headings: list[Heading]) -> list[tuple[str, str]]:
|
||
"""按文档顺序计算每个标题的 (标题文本, 祖先标题链含自身)，链用 " / " 连接"""
|
||
paths: list[tuple[str, str]] = []
|
||
stack: list[Heading] = []
|
||
for heading in headings:
|
||
# 遇到同级或更高级标题时弹栈，维护当前祖先链
|
||
while stack and stack[-1].level >= heading.level:
|
||
stack.pop()
|
||
stack.append(heading)
|
||
paths.append((heading.title, " / ".join(h.title for h in stack)))
|
||
return paths
|
||
|
||
|
||
def _split_l3_blocks(outline: str) -> list[tuple[str, str]]:
|
||
"""将内容大纲按 "## " 行分块，返回 (块标题, 块文本) 列表
|
||
|
||
无 "## " 行时整块作为一个节点（块标题为空）；
|
||
首个 "## " 之前的引导内容直接忽略。
|
||
"""
|
||
stripped = outline.strip()
|
||
if not stripped:
|
||
return []
|
||
|
||
blocks: list[tuple[str, list[str]]] = []
|
||
for line in stripped.splitlines():
|
||
if line.startswith("## "):
|
||
blocks.append((line[3:].strip(), [line]))
|
||
elif blocks:
|
||
blocks[-1][1].append(line)
|
||
if not blocks:
|
||
return [("", stripped)]
|
||
return [(title, "\n".join(lines).strip()) for title, lines in blocks]
|
||
|
||
|
||
class Ingester:
|
||
"""文档入库器：编排总结、分类、切分、向量化与 Qdrant 写入全链路"""
|
||
|
||
def __init__(
|
||
self,
|
||
summarizer: Summarizer | None = None,
|
||
classifier: Classifier | None = None,
|
||
chunker: Chunker | None = None,
|
||
embedding: EmbeddingService | None = None,
|
||
sparse: SparseEncoder | None = None,
|
||
qdrant: QdrantService | None = None,
|
||
) -> None:
|
||
self.summarizer = summarizer or Summarizer()
|
||
self.classifier = classifier or Classifier()
|
||
self.chunker = chunker or Chunker()
|
||
self.embedding = embedding or create_embedding_service()
|
||
self.sparse = sparse or SparseEncoder()
|
||
self.qdrant = qdrant or QdrantService()
|
||
|
||
async def ingest(self, doc: DocumentInput, progress_cb: Callable[[str], None] | None = None) -> IngestionResult:
|
||
"""执行文档入库
|
||
|
||
Args:
|
||
doc: 文档输入（文本内容 + 元数据）
|
||
progress_cb: 可选的阶段进度回调（同步函数），在各阶段边界以
|
||
"summarizing" / "classifying" / "embedding" / "writing" 调用
|
||
|
||
Returns:
|
||
IngestionResult: 入库结果
|
||
|
||
Raises:
|
||
IngestionError: 任一阶段失败时抛出，携带 stage 与已产出总结
|
||
"""
|
||
|
||
def _report(stage: str) -> None:
|
||
if progress_cb is not None:
|
||
progress_cb(stage)
|
||
|
||
logger.info("开始文档入库", title=doc.title, text_length=len(doc.text))
|
||
doc_id = uuid.uuid4().hex
|
||
|
||
# 1. 三级总结
|
||
_report("summarizing")
|
||
try:
|
||
summary = await self.summarizer.summarize(doc.text, title=doc.title)
|
||
except Exception as exc:
|
||
logger.error("入库失败：三级总结", stage=STAGE_SUMMARIZE, error=str(exc))
|
||
raise IngestionError(STAGE_SUMMARIZE, f"三级总结失败: {exc}") from exc
|
||
logger.info("三级总结完成", doc_id=doc_id, level=summary.level.value)
|
||
|
||
# 2. 分类判定（基于 L1 总结）
|
||
_report("classifying")
|
||
try:
|
||
category = await self.classifier.classify(summary.l1_summary, title=doc.title)
|
||
except Exception as exc:
|
||
logger.error("入库失败：分类判定", stage=STAGE_CLASSIFY, error=str(exc))
|
||
raise IngestionError(STAGE_CLASSIFY, f"分类判定失败: {exc}", summary=summary) from exc
|
||
logger.info("分类判定完成", doc_id=doc_id, category=category.main_category, confidence=category.confidence)
|
||
|
||
# 3. 切分 chunk 并构建 L2/L3 大纲节点（(text, section_path) 列表）
|
||
chunks = self.chunker.chunk(doc.text, doc_id)
|
||
l2_nodes = self._build_l2_nodes(doc.text, summary)
|
||
l3_nodes = self._build_l3_nodes(doc.text, summary)
|
||
|
||
# 4. 批量 embedding：L1 + L2 + L3 + chunks 一次调用，按序切片取向量
|
||
texts = [
|
||
summary.l1_summary,
|
||
*(node_text for node_text, _ in l2_nodes),
|
||
*(node_text for node_text, _ in l3_nodes),
|
||
*(c.text for c in chunks),
|
||
]
|
||
try:
|
||
_report("embedding")
|
||
vectors = await self.embedding.embed(texts)
|
||
except Exception as exc:
|
||
logger.error("入库失败：向量化", stage=STAGE_EMBED, error=str(exc))
|
||
raise IngestionError(STAGE_EMBED, f"向量化失败: {exc}", summary=summary) from exc
|
||
l1_vector = vectors[0]
|
||
l2_vectors = vectors[1 : 1 + len(l2_nodes)]
|
||
l3_vectors = vectors[1 + len(l2_nodes) : 1 + len(l2_nodes) + len(l3_nodes)]
|
||
chunk_vectors = vectors[1 + len(l2_nodes) + len(l3_nodes) :]
|
||
|
||
# 5. sparse 向量（仅 L1 与 chunks 需要）
|
||
l1_sparse: SparseVectorTuple | None = None
|
||
chunk_sparses: list[SparseVectorTuple | None] = [None] * len(chunks)
|
||
if settings.sparse_enabled:
|
||
l1_sparse = self.sparse.encode(summary.l1_summary)
|
||
chunk_sparses = [self.sparse.encode(c.text) for c in chunks]
|
||
|
||
# 6. 写入 Qdrant 四层集合
|
||
_report("writing")
|
||
try:
|
||
await self._write_qdrant(
|
||
doc_id,
|
||
doc,
|
||
summary,
|
||
category,
|
||
chunks,
|
||
l2_nodes,
|
||
l3_nodes,
|
||
l1_vector,
|
||
l2_vectors,
|
||
l3_vectors,
|
||
chunk_vectors,
|
||
l1_sparse,
|
||
chunk_sparses,
|
||
)
|
||
except Exception as exc:
|
||
logger.error("入库失败：Qdrant 写入", stage=STAGE_QDRANT, doc_id=doc_id, error=str(exc))
|
||
raise IngestionError(STAGE_QDRANT, f"Qdrant 写入失败: {exc}", summary=summary) from exc
|
||
|
||
logger.info(
|
||
"文档入库完成",
|
||
doc_id=doc_id,
|
||
chunks_count=len(chunks),
|
||
l2_nodes=len(l2_nodes),
|
||
l3_nodes=len(l3_nodes),
|
||
category=category.main_category,
|
||
)
|
||
return IngestionResult(
|
||
document_id=doc_id,
|
||
summary=summary,
|
||
category=category.main_category,
|
||
tags=category.tags,
|
||
category_confidence=category.confidence,
|
||
collection="四层集合",
|
||
chunks_count=len(chunks),
|
||
)
|
||
|
||
def _build_l2_nodes(self, text: str, summary: DocumentSummary) -> list[tuple[str, str]]:
|
||
"""构建 L2 大纲节点，返回 (text, section_path) 列表
|
||
|
||
- 有标题结构（L3 级且标题数 >= 2）：每个标题一个节点，
|
||
text 与 section_path 均为该节点的祖先标题链
|
||
- 否则若 l2_outline 非空（LLM 生成的大纲）：按非空行拆节点，section_path 为空
|
||
- 2.5 级文档（l2_outline 为 None）：无 L2 节点
|
||
"""
|
||
headings = parse_headings(text)
|
||
if summary.level == SummaryLevel.L3 and len(headings) >= 2:
|
||
return [(path, path) for _, path in _heading_paths(headings)]
|
||
if summary.l2_outline:
|
||
return [(line.strip(), "") for line in summary.l2_outline.splitlines() if line.strip()]
|
||
return []
|
||
|
||
def _build_l3_nodes(self, text: str, summary: DocumentSummary) -> list[tuple[str, str]]:
|
||
"""构建 L3 内容大纲节点，返回 (text, section_path) 列表
|
||
|
||
按 "## " 分块（无 "## " 则整块一个节点）；
|
||
section_path 尽力匹配文档标题链（块标题与文档标题文本精确匹配），匹配不到用 ""。
|
||
"""
|
||
path_by_title: dict[str, str] = {}
|
||
for title, path in _heading_paths(parse_headings(text)):
|
||
path_by_title.setdefault(title, path)
|
||
return [
|
||
(block_text, path_by_title.get(block_title, ""))
|
||
for block_title, block_text in _split_l3_blocks(summary.l3_content_outline)
|
||
]
|
||
|
||
async def _write_qdrant(
|
||
self,
|
||
doc_id: str,
|
||
doc: DocumentInput,
|
||
summary: DocumentSummary,
|
||
category: CategoryResult,
|
||
chunks: list[ChunkModel],
|
||
l2_nodes: list[tuple[str, str]],
|
||
l3_nodes: list[tuple[str, str]],
|
||
l1_vector: list[float],
|
||
l2_vectors: list[list[float]],
|
||
l3_vectors: list[list[float]],
|
||
chunk_vectors: list[list[float]],
|
||
l1_sparse: SparseVectorTuple | None,
|
||
chunk_sparses: list[SparseVectorTuple | None],
|
||
) -> None:
|
||
"""将 L1/L2/L3/chunks 四层数据写入 Qdrant（任一失败向上抛出）"""
|
||
await self.qdrant.upsert_l1(
|
||
doc_id=doc_id,
|
||
title=doc.title,
|
||
summary=summary.l1_summary,
|
||
category=category.main_category,
|
||
tags=category.tags,
|
||
dense_vector=l1_vector,
|
||
sparse_vector=l1_sparse,
|
||
metadata=doc.metadata,
|
||
)
|
||
|
||
# L2/L3 大纲节点（为空时跳过对应集合的 upsert）
|
||
for collection, nodes, vectors in (
|
||
(COLLECTION_L2, l2_nodes, l2_vectors),
|
||
(COLLECTION_L3, l3_nodes, l3_vectors),
|
||
):
|
||
if not nodes:
|
||
continue
|
||
await self.qdrant.upsert_nodes(
|
||
collection,
|
||
[
|
||
{
|
||
"doc_id": doc_id,
|
||
"section_path": section_path,
|
||
"text": node_text,
|
||
"category": category.main_category,
|
||
"tags": category.tags,
|
||
"dense_vector": vector,
|
||
}
|
||
for (node_text, section_path), vector in zip(nodes, vectors, strict=True)
|
||
],
|
||
)
|
||
|
||
if chunks:
|
||
# chunk dict 额外携带 doc_summary（= L1 总结），检索侧直接取用，不用回查 L1
|
||
chunk_dicts: list[dict[str, Any]] = [
|
||
{
|
||
"doc_id": doc_id,
|
||
"chunk_index": chunk.chunk_index,
|
||
"text": chunk.text,
|
||
"section_path": chunk.section_path,
|
||
"title": doc.title,
|
||
"category": category.main_category,
|
||
"tags": category.tags,
|
||
"dense_vector": vector,
|
||
"sparse_vector": sparse,
|
||
"doc_summary": summary.l1_summary,
|
||
}
|
||
for chunk, vector, sparse in zip(chunks, chunk_vectors, chunk_sparses, strict=True)
|
||
]
|
||
await self.qdrant.upsert_chunks(chunk_dicts)
|