simonkoson
338 lines
14 KiB
Python
338 lines
14 KiB
Python
"""
|
||
知识库服务 — 写入向量 + 语义检索 + md 文件解析
|
||
使用 pgvector 原生 SQL 向量检索（<=> 余弦距离算子），不在 Python 侧计算
|
||
"""
|
||
|
||
from typing import Optional
|
||
from datetime import date
|
||
|
||
import frontmatter
|
||
from sqlalchemy import text
|
||
from sqlmodel import Session, select
|
||
from pgvector.sqlalchemy import Vector
|
||
|
||
from app.models.knowledge import KnowledgeItem, KnowledgeEmbedding
|
||
from app.services.embedding_service import EmbeddingService
|
||
from app.db.session import engine
|
||
|
||
|
||
class KnowledgeService:
|
||
"""知识库 CRUD + 语义检索 + md 解析"""
|
||
|
||
# yaml 类型字段 → source_type 枚举映射
|
||
SOURCE_TYPE_MAP = {
|
||
"杂志文章": "military_report",
|
||
"军报": "military_report",
|
||
"节目文稿": "manuscript",
|
||
"报题单": "baoti",
|
||
}
|
||
|
||
# 来源大类固定显示顺序（制片人 Obsidian 习惯）
|
||
SOURCE_TYPE_ORDER = [
|
||
"manuscript", # 节目文稿
|
||
"military_report", # 杂志文章
|
||
"baoti", # 报题单
|
||
"manual", # 其他
|
||
]
|
||
|
||
# 二级分组维度映射（与前端 useKnowledgeGrouping.js 的 SECONDARY_GROUP_FIELD 一致）
|
||
# key = source_type, value = 用来做二级分组的字段名，None = 不建二级节点
|
||
SECONDARY_GROUP_FIELD = {
|
||
"manuscript": "author", # 节目文稿 → 按作者(编导)
|
||
"military_report": "source_detail", # 杂志文章 → 按出处
|
||
"baoti": None, # 报题单 → 不分组
|
||
"manual": None, # 其他 → 不分组
|
||
}
|
||
|
||
SOURCE_TYPE_LABEL = {
|
||
"military_report": "杂志文章",
|
||
"manuscript": "节目文稿",
|
||
"baoti": "报题单",
|
||
"manual": "其他",
|
||
}
|
||
|
||
def __init__(self):
|
||
self.embedder = EmbeddingService()
|
||
|
||
def parse_md_file(self, file_content: bytes, file_name: str) -> dict:
|
||
"""
|
||
解析一个 .md 文件的 yaml frontmatter + 正文，返回入库用的字典。
|
||
严格按真实样本的字段名映射，不猜测。
|
||
|
||
Returns:
|
||
dict 含 keys: title, content_md, source_type, author, publish_date,
|
||
source_detail, metadata(JSONB), related_entities(JSONB)
|
||
"""
|
||
content = file_content.decode("utf-8", errors="replace")
|
||
parsed = frontmatter.loads(content)
|
||
fm = parsed.metadata or {}
|
||
|
||
# —— 类型 → source_type（硬映射，不猜测）——
|
||
raw_type = str(fm.get("类型", "")).strip()
|
||
source_type = self.SOURCE_TYPE_MAP.get(raw_type, "manual")
|
||
|
||
# —— 标题：名称 或 标题——
|
||
title = str(fm.get("名称", "") or fm.get("标题", "")).strip()
|
||
if not title:
|
||
# fallback: 用正文第一行或文件名
|
||
lines = [l.strip() for l in content.split("\n") if l.strip() and not l.strip().startswith("---")]
|
||
title = lines[0] if lines else file_name
|
||
|
||
# —— 作者：作者 或 编导——
|
||
author = str(fm.get("作者", "") or fm.get("编导", "") or "").strip() or None
|
||
|
||
# —— 出处详情：期刊 + 期号（拼在一起存进 JSONB 的 source_detail）——
|
||
journal = str(fm.get("期刊", "") or "").strip()
|
||
issue = str(fm.get("期号", "") or "").strip()
|
||
if journal or issue:
|
||
source_detail = f"{journal} {issue}".strip()
|
||
else:
|
||
source_detail = None
|
||
|
||
# —— 播出日期：容错 "待补充" 等非日期文本——
|
||
raw_date = str(fm.get("播出日期", "") or "").strip()
|
||
publish_date = None
|
||
if raw_date and raw_date not in ("待补充", "待确认", ""):
|
||
try:
|
||
publish_date = date.fromisoformat(raw_date)
|
||
except ValueError:
|
||
# 非 ISO 格式，尝试 common 格式
|
||
for fmt in ("%Y-%m-%d", "%Y年%m月%d日", "%Y/%m/%d"):
|
||
try:
|
||
publish_date = date.fromisoformat(raw_date.replace("年", "-").replace("月", "-").replace("日", ""))
|
||
break
|
||
except ValueError:
|
||
continue
|
||
|
||
# —— 权重（不展示，存 JSONB 备 Phase 4）——
|
||
weight = str(fm.get("权重", "") or "").strip() or None
|
||
|
||
# —— 相关实体（涉及装备/涉及技术/涉及厂商/主题）——
|
||
related_entities = []
|
||
for key in ("涉及装备", "涉及技术", "涉及厂商", "主题"):
|
||
val = fm.get(key)
|
||
if val:
|
||
if isinstance(val, list):
|
||
related_entities.extend(val)
|
||
elif isinstance(val, str):
|
||
# 可能是 "山东舰, 福建舰" 这样的逗号分隔字符串
|
||
for item in val.replace("，", ",").split(","):
|
||
item = item.strip()
|
||
if item:
|
||
related_entities.append(item)
|
||
|
||
# —— metadata JSONB：权重、出处详情、双链预留——
|
||
metadata = {}
|
||
if weight:
|
||
metadata["weight"] = weight
|
||
if source_detail:
|
||
metadata["source_detail"] = source_detail
|
||
# related_concepts 字段预留给双链解析（Phase 4），本 Task 原样存入
|
||
metadata["double_bracket_links"] = self._extract_double_brackets(parsed.content)
|
||
|
||
# —— 正文（去掉 frontmatter 的纯内容）——
|
||
content_md = parsed.content
|
||
|
||
return {
|
||
"title": title,
|
||
"content_md": content_md,
|
||
"source_type": source_type,
|
||
"author": author,
|
||
"publish_date": publish_date,
|
||
"metadata": metadata if metadata else None,
|
||
"related_entities": related_entities if related_entities else None,
|
||
"source_file_name": file_name,
|
||
}
|
||
|
||
def _extract_double_brackets(self, text: str) -> list[str]:
|
||
"""提取 [[...]] 双链标记，原样返回列表，不解析成图谱（本 Task 留门）。"""
|
||
import re
|
||
return re.findall(r"\[\[([^\]]+)\]\]", text)
|
||
|
||
def store_md_file(self, file_content: bytes, file_name: str) -> KnowledgeItem:
|
||
"""
|
||
读取一篇 md 内容，调用 embo-01 拿到向量，写入 knowledge_items + knowledge_embeddings
|
||
"""
|
||
parsed = self.parse_md_file(file_content, file_name)
|
||
|
||
# 调用 embedding（type="db" 表示存入知识库）
|
||
embedding_list = self.embedder.embed_single(parsed["content_md"], embed_type="db")
|
||
|
||
with Session(engine) as session:
|
||
item = KnowledgeItem(
|
||
title=parsed["title"],
|
||
content_md=parsed["content_md"],
|
||
source_type=parsed["source_type"],
|
||
source_file_name=parsed["source_file_name"],
|
||
author=parsed["author"],
|
||
publish_date=parsed["publish_date"],
|
||
tags=parsed["metadata"],
|
||
related_entities=parsed["related_entities"],
|
||
)
|
||
session.add(item)
|
||
session.flush()
|
||
|
||
emb = KnowledgeEmbedding(
|
||
knowledge_id=item.id,
|
||
chunk_index=0,
|
||
chunk_text=parsed["content_md"],
|
||
embedding=embedding_list,
|
||
)
|
||
session.add(emb)
|
||
session.commit()
|
||
session.refresh(item)
|
||
return item
|
||
|
||
def delete_item(self, knowledge_id: int) -> bool:
|
||
"""删除知识库条目及其向量（CASCADE 已由 DB 层配置）。"""
|
||
with Session(engine) as session:
|
||
item = session.get(KnowledgeItem, knowledge_id)
|
||
if item is None:
|
||
return False
|
||
session.delete(item)
|
||
session.commit()
|
||
return True
|
||
|
||
def list_items(self, source_type: Optional[str] = None) -> list[dict]:
|
||
"""返回知识库条目列表（含 source_detail 从 metadata 解压）。"""
|
||
with Session(engine) as session:
|
||
statement = select(KnowledgeItem).order_by(KnowledgeItem.created_at.desc())
|
||
if source_type:
|
||
statement = statement.where(KnowledgeItem.source_type == source_type)
|
||
items = session.exec(statement).all()
|
||
|
||
results = []
|
||
for item in items:
|
||
# 从 tags(JSONB) 取 source_detail
|
||
tags = item.tags or {}
|
||
source_detail = tags.get("source_detail") if isinstance(tags, dict) else None
|
||
results.append({
|
||
"id": item.id,
|
||
"title": item.title,
|
||
"author": item.author,
|
||
"publish_date": item.publish_date,
|
||
"source_type": item.source_type,
|
||
"source_file_name": item.source_file_name,
|
||
"source_detail": source_detail,
|
||
"created_at": item.created_at,
|
||
})
|
||
return results
|
||
|
||
def get_distinct_sources(self) -> list[str]:
|
||
"""返回库里所有不重复的 source_detail（动态从 JSONB 提取），供筛选下拉用。"""
|
||
with Session(engine) as session:
|
||
items = session.exec(select(KnowledgeItem)).all()
|
||
sources = set()
|
||
for item in items:
|
||
tags = item.tags or {}
|
||
if isinstance(tags, dict) and tags.get("source_detail"):
|
||
sources.add(tags["source_detail"])
|
||
return sorted(list(sources))
|
||
|
||
def search_similar(self, query_text: str, top_k: int = 5) -> list[dict]:
|
||
"""
|
||
语义检索：查询句转为向量，用 SQL 余弦距离（<=>）在数据库层检索
|
||
返回 top_k 条相似笔记，含相似度分数
|
||
"""
|
||
query_vector = self.embedder.embed_single(query_text, embed_type="query")
|
||
vec_str = "[" + ",".join(str(v) for v in query_vector) + "]"
|
||
|
||
with Session(engine) as session:
|
||
sql = f"""
|
||
SELECT
|
||
ki.id,
|
||
ki.title,
|
||
ki.source_type,
|
||
1 - (ke.embedding <=> '{vec_str}'::vector) AS similarity
|
||
FROM knowledge_embeddings ke
|
||
JOIN knowledge_items ki ON ke.knowledge_id = ki.id
|
||
WHERE ke.chunk_index = 0
|
||
ORDER BY ke.embedding <=> '{vec_str}'::vector
|
||
LIMIT {top_k}
|
||
"""
|
||
stmt = text(sql)
|
||
rows = session.execute(stmt).all()
|
||
return [
|
||
{"id": r.id, "title": r.title, "source_type": r.source_type, "similarity": round(r.similarity, 4)}
|
||
for r in rows
|
||
]
|
||
|
||
def get_item_count(self) -> int:
|
||
with Session(engine) as session:
|
||
return len(session.exec(select(KnowledgeItem)).all())
|
||
|
||
def get_embedding_count(self) -> int:
|
||
with Session(engine) as session:
|
||
return len(session.exec(select(KnowledgeEmbedding)).all())
|
||
|
||
def get_grouped_items(self) -> list[dict]:
|
||
"""
|
||
按 source_type → 二级字段（author / source_detail）两层聚合，返回树形结构数据。
|
||
按 SOURCE_TYPE_ORDER 固定顺序排列，仅显示有数据的大类（count > 0）。
|
||
|
||
二级节点 key 格式：`{source_type}|{二级字段名}|{字段值}`
|
||
例：manuscript|author|左鑫
|
||
military_report|source_detail|航空知识 2026年第1期
|
||
|
||
二级字段值为 null / 空字串 → 归入对应大类，不造空节点。
|
||
空大类（0条）不渲染。
|
||
"""
|
||
with Session(engine) as session:
|
||
items = session.exec(select(KnowledgeItem)).all()
|
||
|
||
total_count = len(items)
|
||
|
||
# 按 source_type 分组，初始化所有已知类别
|
||
type_groups: dict = {st: [] for st in self.SOURCE_TYPE_ORDER}
|
||
for item in items:
|
||
st = item.source_type or "manual"
|
||
if st not in type_groups:
|
||
type_groups[st] = []
|
||
type_groups[st].append(item)
|
||
|
||
children = []
|
||
for st in self.SOURCE_TYPE_ORDER:
|
||
st_items = type_groups.get(st, [])
|
||
# 空类别（0条）不渲染
|
||
if not st_items:
|
||
continue
|
||
|
||
secondary_field = self.SECONDARY_GROUP_FIELD.get(st)
|
||
grandchildren = []
|
||
|
||
if secondary_field is not None:
|
||
# 按二级字段分组
|
||
detail_groups: dict = {}
|
||
for item in st_items:
|
||
if secondary_field == "source_detail":
|
||
tags = item.tags or {}
|
||
sd = tags.get("source_detail") if isinstance(tags, dict) else None
|
||
field_val = sd
|
||
else:
|
||
field_val = (getattr(item, secondary_field, None) or "").strip() or None
|
||
|
||
if field_val not in detail_groups:
|
||
detail_groups[field_val] = []
|
||
detail_groups[field_val].append(item)
|
||
|
||
for sd, sd_items in detail_groups.items():
|
||
if sd is not None:
|
||
grandchildren.append({
|
||
"key": f"{st}|{secondary_field}|{sd}",
|
||
"label": f"{sd}（{len(sd_items)}条）",
|
||
"count": len(sd_items),
|
||
})
|
||
|
||
children.append({
|
||
"key": st,
|
||
"label": f"{self.SOURCE_TYPE_LABEL.get(st, st)}（{len(st_items)}条）",
|
||
"count": len(st_items),
|
||
"children": grandchildren,
|
||
})
|
||
|
||
return [{
|
||
"key": "all",
|
||
"label": f"全部（{total_count}条）",
|
||
"count": total_count,
|
||
"children": children,
|
||
}] |