289 lines
12 KiB
Python
289 lines
12 KiB
Python
"""个人知识库: `workspace/users/<user_id>/.kb/<库名>/` —— 纯文件,无向量无 DB。
|
||
|
||
照「记忆」机制范式(core/memory.py):真实文件为准 + agentic search,索引只是
|
||
可重建的派生视图。每个库一个目录:
|
||
|
||
.kb/<库名>/
|
||
INDEX.md —— 库目录(单行一条,格式见 INDEX_LINE_FORMAT),注 prompt 的召回依据
|
||
docs/<x>.md —— 转换后的 markdown 正文(agent 用 read/grep 按需拉)
|
||
sources/<x> —— 原始文件(pdf/docx/...),留档 + 待入库队列
|
||
|
||
**「已入库」判据 = INDEX.md 有指向该 doc 的条目**;sources 有而 INDEX 无对应 doc
|
||
= 待入库。由此入库天然幂等(重跑只补缺口)、崩溃可恢复(半截转换重来即可)、
|
||
零 DB migration。入库编排在 core/kb_ingest.py,本模块只管状态读写与视图。
|
||
|
||
.kb 是 dotfile:被 GET /v1/files 天然隐藏、validate_task_name 拒 `.` 起头 ——
|
||
与 .memory/.skills 同款防呆。agent 无需新工具:fs 工具 user_root 内可读写,
|
||
docker 沙箱把整个 user_root bind 到 /workspace,.kb 随之可见。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
import shutil
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional
|
||
from uuid import UUID
|
||
|
||
# INDEX 单行格式(全角 ｜ 分隔,摘要/关键词内允许半角标点)。agent 对话内手动入库
|
||
# 与后台 ingest 产出同一格式 —— 契约文本(kb_block)里原样给出。
|
||
INDEX_LINE_FORMAT = "- [标题](docs/<文件名>.md)｜来源 sources/<原件名>｜摘要：<两三句>｜关键词：<逗号分隔>"
|
||
|
||
_INDEX_LINE_RE = re.compile(
|
||
r"^-\s*\[(?P<title>[^\]]*)\]\((?P<doc>docs/[^)]+)\)"
|
||
r"\s*｜\s*来源\s*(?P<source>sources/[^｜]+?)"
|
||
r"\s*｜\s*摘要[:：]\s*(?P<summary>[^｜]*)"
|
||
r"(?:\s*｜\s*关键词[:：]\s*(?P<keywords>.*))?\s*$"
|
||
)
|
||
|
||
# 库名:中文/字母/数字/-/_,拒 dotfile、路径分隔、Windows 保留字符。
|
||
_KB_NAME_RE = re.compile(r"^[\w一-鿿][\w一-鿿\-. ]{0,39}$")
|
||
# 库内文件名(docs/ 与 sources/ 下的扁平文件):拒斜杠 / `..` / dotfile。
|
||
_FILE_NAME_RE = re.compile(r"^[^/\\]{1,200}$")
|
||
|
||
|
||
def kb_root(workspace_dir: Path, user_id: UUID) -> Path:
|
||
return workspace_dir / "users" / str(user_id) / ".kb"
|
||
|
||
|
||
def is_safe_kb_name(name: str) -> bool:
|
||
if not name or name != name.strip() or name.startswith("."):
|
||
return False
|
||
if ".." in name or any(c in name for c in '/\\:*?"<>|'):
|
||
return False
|
||
return bool(_KB_NAME_RE.match(name))
|
||
|
||
|
||
def is_safe_file_name(name: str) -> bool:
|
||
if not name or name.startswith(".") or ".." in name:
|
||
return False
|
||
return bool(_FILE_NAME_RE.match(name))
|
||
|
||
|
||
def kb_dir(workspace_dir: Path, user_id: UUID, name: str) -> Optional[Path]:
|
||
"""库目录(校验名字合法 + 落在 .kb 子树内);非法返回 None(调用方转 4xx)。"""
|
||
if not is_safe_kb_name(name):
|
||
return None
|
||
root = kb_root(workspace_dir, user_id).resolve()
|
||
d = (root / name).resolve()
|
||
if d.parent != root:
|
||
return None
|
||
return d
|
||
|
||
|
||
def parse_index(text: str) -> List[Dict[str, str]]:
|
||
"""解析 INDEX.md → [{title, doc, source, summary, keywords}];不合格式的行忽略。"""
|
||
out: List[Dict[str, str]] = []
|
||
for raw in text.splitlines():
|
||
m = _INDEX_LINE_RE.match(raw.strip())
|
||
if not m:
|
||
continue
|
||
out.append({
|
||
"title": m.group("title").strip(),
|
||
"doc": m.group("doc").strip(),
|
||
"source": (m.group("source") or "").strip(),
|
||
"summary": (m.group("summary") or "").strip(),
|
||
"keywords": (m.group("keywords") or "").strip(),
|
||
})
|
||
return out
|
||
|
||
|
||
def format_index_line(*, title: str, doc: str, source: str, summary: str, keywords: str) -> str:
|
||
"""产出与 INDEX_LINE_FORMAT 一致的单行(后台 ingest 用;agent 侧照契约手写)。"""
|
||
def clean(s: str) -> str:
|
||
return " ".join((s or "").split()).replace("｜", "|")
|
||
return (
|
||
f"- [{clean(title)}]({doc})｜来源 {source}"
|
||
f"｜摘要：{clean(summary)}｜关键词：{clean(keywords)}"
|
||
)
|
||
|
||
|
||
def _read_index(d: Path) -> List[Dict[str, str]]:
|
||
p = d / "INDEX.md"
|
||
if not p.is_file():
|
||
return []
|
||
try:
|
||
return parse_index(p.read_text(encoding="utf-8"))
|
||
except (OSError, UnicodeDecodeError):
|
||
return []
|
||
|
||
|
||
def _list_sources(d: Path) -> List[str]:
|
||
src = d / "sources"
|
||
if not src.is_dir():
|
||
return []
|
||
return sorted(p.name for p in src.iterdir() if p.is_file() and not p.name.startswith("."))
|
||
|
||
|
||
def pending_sources(d: Path) -> List[str]:
|
||
"""sources 有而 INDEX 无 = 待入库(判据即幂等性来源,见模块注释)。"""
|
||
indexed = {e["source"].removeprefix("sources/") for e in _read_index(d)}
|
||
return [n for n in _list_sources(d) if n not in indexed]
|
||
|
||
|
||
def list_kbs(workspace_dir: Path, user_id: UUID) -> List[Dict[str, Any]]:
|
||
"""所有库概览:[{name, doc_count, pending_count}],按名排序。"""
|
||
root = kb_root(workspace_dir, user_id)
|
||
if not root.is_dir():
|
||
return []
|
||
out: List[Dict[str, Any]] = []
|
||
for d in sorted(root.iterdir()):
|
||
if not d.is_dir() or d.name.startswith("."):
|
||
continue
|
||
entries = _read_index(d)
|
||
out.append({
|
||
"name": d.name,
|
||
"doc_count": len(entries),
|
||
"pending_count": len(pending_sources(d)),
|
||
})
|
||
return out
|
||
|
||
|
||
def create_kb(workspace_dir: Path, user_id: UUID, name: str) -> Optional[Path]:
|
||
"""建库(幂等):目录 + 空 INDEX.md + docs/ + sources/。名字非法返回 None。"""
|
||
d = kb_dir(workspace_dir, user_id, name)
|
||
if d is None:
|
||
return None
|
||
(d / "docs").mkdir(parents=True, exist_ok=True)
|
||
(d / "sources").mkdir(parents=True, exist_ok=True)
|
||
idx = d / "INDEX.md"
|
||
if not idx.exists():
|
||
idx.write_text(f"# {name}\n\n", encoding="utf-8")
|
||
return d
|
||
|
||
|
||
def delete_kb(workspace_dir: Path, user_id: UUID, name: str) -> bool:
|
||
"""整库删除(原件 + docs + INDEX 一起没,前端已二次确认)。"""
|
||
d = kb_dir(workspace_dir, user_id, name)
|
||
if d is None or not d.is_dir():
|
||
return False
|
||
shutil.rmtree(d)
|
||
return True
|
||
|
||
|
||
def kb_detail(workspace_dir: Path, user_id: UUID, name: str) -> Optional[Dict[str, Any]]:
|
||
"""单库全貌:INDEX 条目 + 待入库 sources 列表。库不存在返回 None。"""
|
||
d = kb_dir(workspace_dir, user_id, name)
|
||
if d is None or not d.is_dir():
|
||
return None
|
||
return {
|
||
"name": name,
|
||
"entries": _read_index(d),
|
||
"pending": pending_sources(d),
|
||
}
|
||
|
||
|
||
def read_doc(workspace_dir: Path, user_id: UUID, name: str, filename: str) -> Optional[str]:
|
||
"""读单篇 docs/<filename> 原文;非法 / 不存在 → None(调用方转 404)。"""
|
||
d = kb_dir(workspace_dir, user_id, name)
|
||
if d is None or not is_safe_file_name(filename) or not filename.endswith(".md"):
|
||
return None
|
||
target = (d / "docs" / filename).resolve()
|
||
if target.parent != (d / "docs").resolve() or not target.is_file():
|
||
return None
|
||
try:
|
||
return target.read_text(encoding="utf-8")
|
||
except (OSError, UnicodeDecodeError):
|
||
return None
|
||
|
||
|
||
def delete_doc(workspace_dir: Path, user_id: UUID, name: str, filename: str) -> bool:
|
||
"""删单篇:docs 文件 + INDEX 对应行 + 对应 source 原件一起删(否则原件会被当
|
||
待入库重新转一遍 —— 判据使然)。"""
|
||
d = kb_dir(workspace_dir, user_id, name)
|
||
if d is None or not is_safe_file_name(filename) or not filename.endswith(".md"):
|
||
return False
|
||
doc_rel = f"docs/{filename}"
|
||
entries = _read_index(d)
|
||
hit = next((e for e in entries if e["doc"] == doc_rel), None)
|
||
target = (d / "docs" / filename).resolve()
|
||
if target.parent != (d / "docs").resolve():
|
||
return False
|
||
if not target.is_file() and hit is None:
|
||
return False
|
||
if target.is_file():
|
||
target.unlink()
|
||
if hit is not None:
|
||
src_name = hit["source"].removeprefix("sources/")
|
||
if is_safe_file_name(src_name):
|
||
src = d / "sources" / src_name
|
||
if src.is_file():
|
||
src.unlink()
|
||
idx = d / "INDEX.md"
|
||
try:
|
||
lines = idx.read_text(encoding="utf-8").splitlines()
|
||
kept = [
|
||
ln for ln in lines
|
||
if not (_INDEX_LINE_RE.match(ln.strip()) and f"({doc_rel})" in ln)
|
||
]
|
||
idx.write_text("\n".join(kept).rstrip() + "\n", encoding="utf-8")
|
||
except (OSError, UnicodeDecodeError):
|
||
pass
|
||
return True
|
||
|
||
|
||
def save_source(workspace_dir: Path, user_id: UUID, name: str, filename: str, data: bytes) -> Optional[str]:
|
||
"""上传原件落 sources/(同名覆盖 —— 重传即重新入库的自然语义)。
|
||
返回落盘文件名;库不存在 / 文件名非法 → None。"""
|
||
d = kb_dir(workspace_dir, user_id, name)
|
||
if d is None or not d.is_dir() or not is_safe_file_name(filename):
|
||
return None
|
||
src_dir = d / "sources"
|
||
src_dir.mkdir(parents=True, exist_ok=True)
|
||
# 覆盖旧 doc 判据:同名 source 若已在 INDEX,删掉旧条目让它重新排队入库
|
||
doc_rel_hits = [e for e in _read_index(d) if e["source"] == f"sources/{filename}"]
|
||
for e in doc_rel_hits:
|
||
delete_doc(workspace_dir, user_id, name, e["doc"].removeprefix("docs/"))
|
||
(src_dir / filename).write_bytes(data)
|
||
return filename
|
||
|
||
|
||
# ── prompt 注入(照 memory_block 范式) ────────────────────────────────
|
||
|
||
_KB_CONTRACT = """\
|
||
用法规矩:
|
||
- **主动查阅**:回答前先扫一眼下方各库条目,标题 / 摘要 / 关键词与用户问题相关就先
|
||
`read` 对应 `docs/*.md` 正文再作答(量大时 `grep` 先缩范围)。**无需用户点名
|
||
"知识库"** —— 用户资料里可能已有的内容,宁可多查一次,不要凭空回答。
|
||
- **答题标来源**:引用了哪个库哪篇就在回答里注明(标题或文件名)。
|
||
- **对话内入库**:用户在对话里给了值得长期留的资料时,可直接写入:原件放
|
||
`sources/`(没有原件就跳过)、正文转成 markdown 写 `docs/<slug>.md`、再往该库
|
||
INDEX.md 追加一行,**格式必须是**:
|
||
`{fmt}`
|
||
(与后台自动入库产出一致;摘要写准 —— 它是下次召回的依据)。"""
|
||
|
||
|
||
def kb_block(
|
||
workspace_dir: Path,
|
||
user_id: UUID,
|
||
kb_dir_display: Optional[str] = None,
|
||
) -> str:
|
||
"""构造注入 system prompt 的知识库段;用户没有任何库时返回空串(零成本)。
|
||
|
||
kb_dir_display: `.kb/` 在 agent 视角下的路径前缀(docker 传 `/workspace/.kb`,
|
||
host 传 None ⇒ 宿主绝对路径)—— 与 memory_block 的 mem_dir_display 同款约定。
|
||
注 INDEX 全文而非只注库名:INDEX 就是召回索引(照 memory extended 的
|
||
description 逻辑),几十篇的个人库体量注得起;正文仍按需 read。
|
||
"""
|
||
kbs = list_kbs(workspace_dir, user_id)
|
||
if not kbs:
|
||
return ""
|
||
root = kb_root(workspace_dir, user_id)
|
||
base = (kb_dir_display if kb_dir_display is not None else str(root)).rstrip("/")
|
||
|
||
parts = ["\n\n## 个人知识库 (user 级,跨 task 共享)\n"]
|
||
parts.append(_KB_CONTRACT.replace("{fmt}", INDEX_LINE_FORMAT))
|
||
for kb in kbs:
|
||
d = root / kb["name"]
|
||
parts.append(f"\n\n### 库「{kb['name']}」(`{base}/{kb['name']}/`)\n")
|
||
entries = _read_index(d)
|
||
if not entries:
|
||
parts.append("(空库,尚无已入库文档)\n")
|
||
continue
|
||
for e in entries:
|
||
kw = f"｜关键词：{e['keywords']}" if e["keywords"] else ""
|
||
parts.append(
|
||
f"- [{e['title']}](`{base}/{kb['name']}/{e['doc']}`)"
|
||
f"｜摘要：{e['summary']}{kw}\n"
|
||
)
|
||
return "".join(parts)
|