kplam
51dc8dc4f6
- FastAPI + Qdrant + Redis + Ollama 技术栈 - L1→L2→L3→chunk 四层分层检索（dense + sparse RRF 融合） - 文档三级总结与 2.5 级回退 - query 解析路由与分类 - /admin 管理页面
140 lines
5.7 KiB
Python
140 lines
5.7 KiB
Python
"""文档分类器
|
||
|
||
入库链路第二步：基于 L1 总结，用 Ollama 小模型将文档判定为 taxonomy 中的
|
||
主类目 + 附加标签：
|
||
- LLM 输出解析失败 / 类目名不在 taxonomy → 归 uncategorized（confidence=0.0）
|
||
- 置信度低于阈值 → 主类目归 uncategorized，候选类目名保留进 tags（软召回用）
|
||
"""
|
||
|
||
import json
|
||
import re
|
||
|
||
import structlog
|
||
|
||
from app.config import settings
|
||
from app.models.knowledge import UNCATEGORIZED, CategoryResult, TaxonomyCategory, load_taxonomy
|
||
from app.services.ollama import OllamaClient
|
||
|
||
logger = structlog.get_logger()
|
||
|
||
# 从 LLM 输出中提取第一个 {...} JSON 块（贪婪匹配到最后的 }，兼容嵌套对象）
|
||
_JSON_BLOCK_RE = re.compile(r"\{.*\}", re.DOTALL)
|
||
|
||
|
||
def _extract_json(raw: str) -> dict | None:
|
||
"""从 LLM 输出中提取 JSON 对象
|
||
|
||
先尝试直接解析；失败则用正则提取第一个 {...} 块再解析。
|
||
返回 None 表示无法提取出合法的 JSON 对象。
|
||
"""
|
||
text = raw.strip()
|
||
try:
|
||
data = json.loads(text)
|
||
return data if isinstance(data, dict) else None
|
||
except json.JSONDecodeError:
|
||
pass
|
||
|
||
match = _JSON_BLOCK_RE.search(text)
|
||
if not match:
|
||
return None
|
||
try:
|
||
data = json.loads(match.group(0))
|
||
return data if isinstance(data, dict) else None
|
||
except json.JSONDecodeError:
|
||
return None
|
||
|
||
|
||
class Classifier:
|
||
"""文档分类器：将 L1 总结判定为 taxonomy 主类目 + 附加标签"""
|
||
|
||
def __init__(self, ollama: OllamaClient | None = None, taxonomy: list[TaxonomyCategory] | None = None) -> None:
|
||
self.ollama = ollama or OllamaClient()
|
||
self.taxonomy = taxonomy if taxonomy is not None else load_taxonomy()
|
||
# 合法类目名集合（含 uncategorized）
|
||
self._valid_names = {c.name for c in self.taxonomy}
|
||
|
||
async def classify(self, l1_summary: str, title: str = "") -> CategoryResult:
|
||
"""对文档进行分类判定
|
||
|
||
Args:
|
||
l1_summary: 文档 L1 总结
|
||
title: 文档标题（可选，辅助判定）
|
||
|
||
Returns:
|
||
CategoryResult: 主类目 / 附加标签 / 置信度
|
||
"""
|
||
prompt = self._build_prompt(l1_summary, title)
|
||
raw = await self.ollama.generate(prompt, json_mode=True)
|
||
|
||
data = _extract_json(raw)
|
||
if data is None:
|
||
logger.warning("分类失败：LLM 输出非合法 JSON", title=title, output=raw[:200])
|
||
return CategoryResult(main_category=UNCATEGORIZED, tags=[], confidence=0.0)
|
||
|
||
main_category = data.get("main_category")
|
||
confidence = data.get("confidence")
|
||
if (
|
||
not isinstance(main_category, str)
|
||
or main_category not in self._valid_names
|
||
or not isinstance(confidence, (int, float))
|
||
or not 0 <= confidence <= 1
|
||
):
|
||
logger.warning(
|
||
"分类失败：类目名不在 taxonomy 或 confidence 非法",
|
||
title=title,
|
||
main_category=main_category,
|
||
confidence=confidence,
|
||
)
|
||
return CategoryResult(main_category=UNCATEGORIZED, tags=[], confidence=0.0)
|
||
|
||
tags = self._clean_tags(data.get("tags"), main_category)
|
||
confidence = float(confidence)
|
||
|
||
# 低置信度软召回：主类目归 uncategorized，候选类目名保留进 tags
|
||
if confidence < settings.classify_confidence_threshold:
|
||
logger.info(
|
||
"分类置信度低于阈值，归入 uncategorized",
|
||
title=title,
|
||
candidate=main_category,
|
||
confidence=confidence,
|
||
threshold=settings.classify_confidence_threshold,
|
||
)
|
||
if main_category != UNCATEGORIZED and main_category not in tags:
|
||
tags.insert(0, main_category)
|
||
return CategoryResult(main_category=UNCATEGORIZED, tags=tags, confidence=confidence)
|
||
|
||
return CategoryResult(main_category=main_category, tags=tags, confidence=confidence)
|
||
|
||
def _build_prompt(self, l1_summary: str, title: str) -> str:
|
||
"""构造分类 prompt：列出全部 taxonomy 类目（含 uncategorized 及其用途说明）"""
|
||
category_lines = []
|
||
for c in self.taxonomy:
|
||
if c.name == UNCATEGORIZED:
|
||
category_lines.append(f"- {c.name}: 当文档跨多个类目或无法明确归入其他类目时选择此类目")
|
||
else:
|
||
category_lines.append(f"- {c.name}: {c.description}")
|
||
category_block = "\n".join(category_lines)
|
||
|
||
prompt = (
|
||
"你是知识库分类助手。请根据文档标题和总结，判断文档最适合归入以下哪个类目。\n\n"
|
||
f"可选类目：\n{category_block}\n\n"
|
||
"要求：\n"
|
||
"- main_category 只能从上面的类目名中选择，不要输出其他名称；\n"
|
||
"- tags 为 0~3 个附加标签（词或短语），且不能包含 main_category 本身；\n"
|
||
"- confidence 为 0~1 之间的小数，表示对主类目判断的置信度；\n"
|
||
"- 只输出 JSON，不要输出任何其他内容。\n\n"
|
||
'输出格式：{"main_category": "类目名", "tags": ["标签"], "confidence": 0.0}\n\n'
|
||
)
|
||
if title:
|
||
prompt += f"文档标题：{title}\n"
|
||
prompt += f"文档总结：{l1_summary}"
|
||
return prompt
|
||
|
||
@staticmethod
|
||
def _clean_tags(raw_tags: object, main_category: str) -> list[str]:
|
||
"""清洗 LLM 输出的 tags：仅保留非空字符串、剔除主类目名、最多 3 个"""
|
||
if not isinstance(raw_tags, list):
|
||
return []
|
||
tags = [t for t in raw_tags if isinstance(t, str) and t and t != main_category]
|
||
return tags[:3]
|