mirror of
https://gitcode.com/JianFeeeee/HomeAgent.git
synced 2026-09-27 12:53:35 +00:00
feat(knowledge): 目录批量导入 + 派生数据批量收口
## 能力缺口 导入只能一条条 Add（knowledge_create）。agent 拿到一份 200 页的 文档目录要调 200 次工具，且每次都得自己决定分类与名字。 新增工具 `knowledge_import_dir(dir, category?, include_media?, dry_run?, max_items?)`。 语义是**复制**不是引用：源文件删改不影响已导入的副本。 - 文本经 Write 整份写入 <知识根>/<分类>/<名>/content.md - 媒体按 sha256 进媒体库（内容寻址天然去重），条目只存 digest 引用 ## 目录约定（自动适配，不要求改造资料） 1. 含 content.md 的目录 ⇒ 整体作为一个条目（与 scanDir 既有语义一致， 所以知识库自身目录能被原样再导入而不会被拆散） 2. 否则 .md/.txt 等文件各成一条，**目录路径即分类** ## 三个语义决策 - category 是**前缀叠加**（tech + 源结构），不替换：替换会丢掉源目录 自身最有价值的层级信息 - 同名冲突**跳过并计数**，绝不覆盖：Write 对同名本就是覆盖语义 （knowledge_create 靠它做更新），若直接调它，一次重导就会把手工 补充的内容悄悄抹掉，而日志只写"导入完成" - dry_run **默认 true**：批量写，agent 第一次试某目录应先看清会写什么 ## 安全边界（批量操作，缺一道就可能读到不该读的） - 必须绝对路径：agent 的 cwd 不受控，相对路径会静默导到别处 - 符号链接不跟随：否则一个软链就把知识根之外的文件导进来 - 拒绝把知识库自身当源（自导会无限自我复制） - category 复用 normalizeName（与 Write 同一道闸，两处分叉就成了绕过） - MaxItems 默认 500：防 agent 误传 "/" 把盘灌满 ## ★ 批量导入暴露的既有 O(N²) Write 每条末尾都调 flushDenseLocked，而 saveDenseCacheLocked 是 **全量序列化整个 items map 再重写整个文件**。按 512 维 float64 估， 单条约 10KB，导入 500 条累计要写约 1.4GB。 仓库里索引侧早已有 indexDirty 的「标脏+延迟收口」（实测 writeIndexLocked 6.7ms/次、占单条 Add 绝大部分），**稠密缓存却还是逐条全量重写** —— 同一类开销只修了一半。 照 indexDirty 的模式补 batchDepth：批量期只标脏，endBatch 收口一次。 用 defer 保证提前 return 也会收口 —— 否则这批向量会留成"标脏未写"， 下次启动被当作缺失而全量重算。 ## 判据：17 条 + 变异 安全边界做了 4 组变异验证（去符号链接拦截/去绝对路径要求/去自导检查/ content.md 目录不下钻）。 ★ 判据第一版有两处自己骗自己，被变异抓出来： 1. 符号链接判据造的是**目录软链**，而 WalkDir 对目录软链本来就不下钻 ⇒ 有无防护结果都一样，是假绿。改成**文件软链**后才真正判红。 2. 同名冲突判据里已有条目写成 "a/b"、源映射出的是 "b"（不同名）， 判据自己就错了 —— 修判据而不是改实现。 媒体路径用假 MediaPutter：验的是「调了 Put 且 digest 挂到条目上」， 媒体库自身的落盘去重是 media 包的判据，不该在这里重测。 全量 41 包绿。
This commit is contained in:
JianFeeeee
@ -346,6 +346,10 @@ func New(cfg AgentConfig) *Agent {
|
||||
// 若先重建后接线，首次启动算出的向量会被丢掉而不落盘。
|
||||
cfg.Knowledge.SetDenseSpace(cfg.MultimodalSpace)
|
||||
cfg.Knowledge.SetMediaGetter(cfg.MediaStore)
|
||||
// 媒体**写入**器：ImportDir 复制图片/音视频时用。
|
||||
// 不接的话 include_media=true 会静默失效（媒体全被跳过、
|
||||
// 只在导入结果里留一行"媒体库不可用"），模型无从察觉。
|
||||
cfg.Knowledge.SetMediaPutter(cfg.MediaStore)
|
||||
built, skipped := cfg.Knowledge.ReindexDense()
|
||||
if built > 0 || skipped > 0 {
|
||||
log.Printf("[knowledge] 多模态稠密索引: 新建 %d 跳过 %d（其余命中缓存）", built, skipped)
|
||||
|
||||
@ -615,6 +615,66 @@ func (a *Agent) executeKnowledgeTool(tc agentAPI.ToolCall) string {
|
||||
tree := a.knowledge.BuildTree()
|
||||
return formatTree(tree, 0)
|
||||
|
||||
case "knowledge_import_dir":
|
||||
dir, _ := tc.Arguments["dir"].(string)
|
||||
category, _ := tc.Arguments["category"].(string)
|
||||
// dry_run 默认 true：导入是批量写，agent 第一次试某个目录时
|
||||
// 应该先看清会写什么。默认直接写等于让它盲写一批数据。
|
||||
dryRun := true
|
||||
if b, ok := getBool(tc.Arguments, "dry_run"); ok {
|
||||
dryRun = b
|
||||
}
|
||||
includeMedia := false
|
||||
if b, ok := getBool(tc.Arguments, "include_media"); ok {
|
||||
includeMedia = b
|
||||
}
|
||||
st, err := a.knowledge.ImportDir(knowledge.ImportOptions{
|
||||
Dir: dir,
|
||||
Category: category,
|
||||
DryRun: dryRun,
|
||||
IncludeMedia: includeMedia,
|
||||
MaxItems: int(getFloat(tc.Arguments, "max_items")),
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Sprintf("知识导入失败: %v", err)
|
||||
}
|
||||
var b strings.Builder
|
||||
verb := "已导入"
|
||||
if dryRun {
|
||||
verb = "将导入（dry_run，未实际写入）"
|
||||
}
|
||||
fmt.Fprintf(&b, "%s %d 条", verb, st.Imported)
|
||||
if category != "" {
|
||||
fmt.Fprintf(&b, "（分类前缀 %s）", category)
|
||||
}
|
||||
if st.Media > 0 {
|
||||
fmt.Fprintf(&b, "，含 %d 个媒体", st.Media)
|
||||
}
|
||||
if st.Skipped > 0 {
|
||||
fmt.Fprintf(&b, "；跳过 %d", st.Skipped)
|
||||
}
|
||||
if st.Failed > 0 {
|
||||
fmt.Fprintf(&b, "；失败 %d", st.Failed)
|
||||
}
|
||||
if st.Truncated {
|
||||
fmt.Fprintf(&b, "；★ 超出 max_items 被截断，未导完（可调大 max_items 或分批）")
|
||||
}
|
||||
if len(st.Names) > 0 {
|
||||
show := st.Names
|
||||
if len(show) > 10 {
|
||||
show = show[:10]
|
||||
}
|
||||
fmt.Fprintf(&b, "。知识名：%s", strings.Join(show, "、"))
|
||||
if len(st.Names) > 10 {
|
||||
fmt.Fprintf(&b, " …共 %d 个", len(st.Names))
|
||||
}
|
||||
}
|
||||
// 失败原因要报给 agent：否则它只知道"失败 37 条"却无从下手。
|
||||
for _, e := range st.Errors {
|
||||
fmt.Fprintf(&b, "\n- %s", e)
|
||||
}
|
||||
return b.String()
|
||||
|
||||
case "knowledge_delete":
|
||||
name, _ := tc.Arguments["name"].(string)
|
||||
if name == "" {
|
||||
|
||||
@ -391,6 +391,25 @@ func (a *Agent) buildToolDefs() []interface{} {
|
||||
"description": "可选：关联的媒体 digest（可传前缀）。媒体作为一等节点参与跨模态检索——知识能按图本身被搜到，而不依赖生成的描述文本",
|
||||
},
|
||||
}, "name", "content"))
|
||||
tools = append(tools, toolDef("knowledge_import_dir",
|
||||
"从目录批量导入知识（复制，不是引用）。递归扫描该目录，把文档复制进知识库并自动完成索引与向量化，之后即可被 knowledge_search 检索到。\n"+
|
||||
"目录约定（自动适配，无需改造资料）：含 content.md 的目录整体作为一条知识；否则目录里的 .md/.txt 等文件各成一条，**目录路径即分类**。\n"+
|
||||
"注意：dir 必须是绝对路径（你的工作目录不等于进程 cwd）；同名条目会跳过而不覆盖。",
|
||||
map[string]interface{}{
|
||||
"dir": map[string]interface{}{"type": "string", "description": "源目录的绝对路径，如 /home/qq/notes。不接受相对路径"},
|
||||
"category": map[string]interface{}{"type": "string", "description": "可选：分类前缀，叠加在源目录结构之上。导入 ~/docs/go/x.md 配 category=tech 得到 tech/go/x。留空则沿用源目录自身的层级"},
|
||||
"include_media": map[string]interface{}{
|
||||
"type": "boolean",
|
||||
"description": "可选：是否把目录里的图片/音视频也导入（复制进媒体库，使该知识能按图本身被搜到）。默认 false",
|
||||
"default": false,
|
||||
},
|
||||
"dry_run": map[string]interface{}{
|
||||
"type": "boolean",
|
||||
"description": "可选：只报告将导入什么，不实际写入。试探性导入时建议先跑一次",
|
||||
"default": true,
|
||||
},
|
||||
"max_items": map[string]interface{}{"type": "integer", "description": "可选：本次最多导入多少条（默认 500，防误传大目录）"},
|
||||
}, "dir"))
|
||||
tools = append(tools, toolDef("knowledge_delete", "删除知识库中的指定知识条目。", map[string]interface{}{
|
||||
"name": map[string]interface{}{"type": "string", "description": "要删除的知识名称"},
|
||||
}, "name"))
|
||||
|
||||
@ -17,6 +17,33 @@ func getString(m map[string]interface{}, key string) string {
|
||||
return ""
|
||||
}
|
||||
|
||||
// getBool 从工具参数里取布尔值。
|
||||
//
|
||||
// 为何不直接 v.(bool)：LLM 的工具参数经 JSON 解码后可能是
|
||||
// bool、也可能是字符串（"true"）或数字（1）—— 实际调用里三种都出现过。
|
||||
// 只认 bool 会让模型明明传了 true 却走默认值（表现为「参数被无视」，
|
||||
// 排查时很难看出是解析问题）。
|
||||
func getBool(m map[string]interface{}, key string) (bool, bool) {
|
||||
v, ok := m[key]
|
||||
if !ok {
|
||||
return false, false
|
||||
}
|
||||
switch b := v.(type) {
|
||||
case bool:
|
||||
return b, true
|
||||
case string:
|
||||
switch strings.ToLower(strings.TrimSpace(b)) {
|
||||
case "1", "true", "yes", "on":
|
||||
return true, true
|
||||
case "0", "false", "no", "off":
|
||||
return false, true
|
||||
}
|
||||
case float64:
|
||||
return b != 0, true
|
||||
}
|
||||
return false, false
|
||||
}
|
||||
|
||||
func getFloat(m map[string]interface{}, key string) float64 {
|
||||
if v, ok := m[key]; ok {
|
||||
switch n := v.(type) {
|
||||
|
||||
484
internal/knowledge/import.go
Normal file
484
internal/knowledge/import.go
Normal file
@ -0,0 +1,484 @@
|
||||
package knowledge
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io/fs"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory/media"
|
||||
)
|
||||
|
||||
// ===== 从目录导入知识 =====
|
||||
|
||||
// ImportOptions 是 ImportDir 的参数。
|
||||
type ImportOptions struct {
|
||||
// Dir 是源目录，**必须是绝对路径**。
|
||||
//
|
||||
// 为什么强制绝对：agent 调工具时 cwd 不受控（内核 daemon 的 cwd 是
|
||||
// 工作目录，不是 agent 心智模型里的那个"项目目录"）。相对路径会
|
||||
// 静默导到错误的地方，而用户以为导的是他指定的目录。
|
||||
Dir string
|
||||
|
||||
// Category 是**前缀叠加**在源目录结构之上的分类路径。
|
||||
//
|
||||
// 选"叠加"而非"替换"：替换会丢掉源目录本身的层级信息
|
||||
//（导入 ~/docs/go/x.md 配 category=tech 得到 tech/x，
|
||||
// 而"go" 这一层正是这份资料最有价值的组织信息）。
|
||||
Category string
|
||||
|
||||
// DryRun 只统计不落盘。
|
||||
DryRun bool
|
||||
|
||||
// MaxItems 限制本次导入条目数，0 用默认值。
|
||||
// 防呆：agent 误传 "/" 或整个家目录时不会一次把盘灌满。
|
||||
MaxItems int
|
||||
|
||||
// IncludeMedia 是否把图片/音视频一并导入（进媒体库、按 digest 引用）。
|
||||
IncludeMedia bool
|
||||
}
|
||||
|
||||
// ImportStats 是导入结果。
|
||||
type ImportStats struct {
|
||||
// Imported 成功写入的条目数。
|
||||
Imported int
|
||||
// Skipped 因同名冲突跳过的条目数（**不覆盖已有知识**）。
|
||||
Skipped int
|
||||
// Failed 失败的条目数。
|
||||
Failed int
|
||||
// Media 是导入的媒体文件数。
|
||||
Media int
|
||||
// Truncated 表示因 MaxItems 被截断（agent 据此知道"没导完"）。
|
||||
Truncated bool
|
||||
// Names 是本次导入的知识名（dry_run 与成功路径都给，便于 agent 复核）。
|
||||
Names []string
|
||||
// Errors 记前若干条失败原因（不全量：失败成百上千时返回上万行没意义）。
|
||||
Errors []string
|
||||
}
|
||||
|
||||
// defaultImportMaxItems 是单次导入的默认上限。
|
||||
const defaultImportMaxItems = 500
|
||||
|
||||
// importTextExt 是被当作正文导入的文本扩展名。
|
||||
//
|
||||
// 刻意白名单而非"非二进制都算"：.exe/.zip/.so 这类即便不是二进制，
|
||||
// 塞进知识库也只是污染检索面（还会被误算进 IDF）。
|
||||
var importTextExt = map[string]bool{
|
||||
".md": true, ".markdown": true, ".txt": true, ".text": true,
|
||||
".rst": true, ".org": true, ".adoc": true,
|
||||
}
|
||||
|
||||
// importMediaExt 是被当作媒体导入的扩展名（对应 media.Kind）。
|
||||
var importMediaExt = map[string]string{
|
||||
".png": "image", ".jpg": "image", ".jpeg": "image", ".gif": "image",
|
||||
".webp": "image", ".bmp": "image",
|
||||
".mp3": "audio", ".wav": "audio", ".m4a": "audio", ".flac": "audio", ".ogg": "audio",
|
||||
".mp4": "video", ".mov": "video", ".webm": "video", ".mkv": "video", ".avi": "video",
|
||||
}
|
||||
|
||||
// maxImportErrors 限制 Errors 的条数。
|
||||
const maxImportErrors = 20
|
||||
|
||||
// ImportDir 把一个目录里的文档复制进知识库，并完成索引与向量化。
|
||||
//
|
||||
// 语义是**复制**不是引用：
|
||||
// - 文本经 Write 整份写入 <知识根>/<分类>/<名>/content.md
|
||||
// - 媒体按 sha256 进媒体库（内容寻址，天然去重），条目只存 digest 引用
|
||||
//
|
||||
// 源目录之后删掉/改动都不影响已导入的副本。
|
||||
//
|
||||
// 目录约定（自动适配两种，不要求用户改造资料）：
|
||||
// 1. 目录里有 content.md ⇒ 整个目录是一个条目（与 scanDir 的既有语义一致，
|
||||
// 所以知识库自身的目录能被原样再导入而不会被拆散）；
|
||||
// 2. 否则目录里的 .md/.txt 等文件各是一个条目，**目录路径即分类**。
|
||||
//
|
||||
// 安全边界（批量操作，缺一道就可能把不该读的东西读进来）：
|
||||
// - 必须绝对路径
|
||||
// - 符号链接不跟随（否则一个软链就能把知识根之外的文件导进来）
|
||||
// - 拒绝把知识库自身当源（自导会无限自我复制）
|
||||
func (s *Store) ImportDir(opt ImportOptions) (ImportStats, error) {
|
||||
var st ImportStats
|
||||
if strings.TrimSpace(opt.Dir) == "" {
|
||||
return st, fmt.Errorf("knowledge: 导入目录不能为空")
|
||||
}
|
||||
if !filepath.IsAbs(opt.Dir) {
|
||||
return st, fmt.Errorf("knowledge: 导入目录必须是绝对路径（agent 的 cwd 不可控，相对路径会导到别处）")
|
||||
}
|
||||
// 符号链接不跟随：Resolve 拿到的是**逻辑**路径，不碰文件系统。
|
||||
// 但目录内部可能藏软链，所以还要在 walk 里逐个 EvalSymlinks 校验。
|
||||
srcRoot := filepath.Clean(opt.Dir)
|
||||
realSrc, err := filepath.EvalSymlinks(srcRoot)
|
||||
if err != nil {
|
||||
return st, fmt.Errorf("knowledge: 导入目录不可达: %w", err)
|
||||
}
|
||||
// 拒绝自导：源目录在知识根之内（含知识根本身）。
|
||||
realRoot := filepath.Clean(s.root)
|
||||
if realSrc == realRoot || strings.HasPrefix(realSrc, realRoot+string(filepath.Separator)) {
|
||||
return st, fmt.Errorf("knowledge: 不能把知识库自身当导入源（会自我复制）")
|
||||
}
|
||||
si, err := os.Stat(realSrc)
|
||||
if err != nil || !si.IsDir() {
|
||||
return st, fmt.Errorf("knowledge: 导入源不是目录: %s", opt.Dir)
|
||||
}
|
||||
|
||||
maxItems := opt.MaxItems
|
||||
if maxItems <= 0 {
|
||||
maxItems = defaultImportMaxItems
|
||||
}
|
||||
// 批量模式：派生数据（稠密缓存 + 索引）只在末尾各写一次。
|
||||
defer s.endBatch()
|
||||
if !opt.DryRun {
|
||||
s.beginBatch()
|
||||
}
|
||||
category := strings.Trim(strings.TrimSpace(opt.Category), "/")
|
||||
// 分类本身也走同一套路径校验：category 来自 LLM，不能是 "../.." 或隐藏段。
|
||||
// 复用 normalizeName 而非另写一份：它是 Write 的同一道闸，
|
||||
// 两处校验一旦分叉，导入就会成为绕过 Write 校验的后门。
|
||||
if category != "" {
|
||||
norm, nerr := normalizeName(category)
|
||||
if nerr != nil {
|
||||
return st, fmt.Errorf("knowledge: category 不合法: %w", nerr)
|
||||
}
|
||||
category = norm
|
||||
}
|
||||
|
||||
// 先扫出全部候选（不写），再逐条导入：这样 MaxItems 的截断是
|
||||
// "按确定顺序取前 N 条"，而不是"边走边停"导致结果不可复现。
|
||||
var cands []importCandidate
|
||||
walkErr := filepath.WalkDir(realSrc, func(p string, d fs.DirEntry, err error) error {
|
||||
if err != nil {
|
||||
st.Errors = appendLimited(st.Errors, fmt.Sprintf("跳过 %s: %v", p, err))
|
||||
return nil // 不中断：个别不可读不该让整次导入失败
|
||||
}
|
||||
rel, rerr := filepath.Rel(realSrc, p)
|
||||
if rerr != nil {
|
||||
return nil
|
||||
}
|
||||
rel = filepath.ToSlash(rel)
|
||||
if rel == "." {
|
||||
return nil
|
||||
}
|
||||
// 逐层拦软链：WalkDir 不会跟随目录软链，但**文件**软链会当成普通文件。
|
||||
if isSymlink(p) {
|
||||
st.Skipped++
|
||||
st.Errors = appendLimited(st.Errors, fmt.Sprintf("跳过符号链接: %s", rel))
|
||||
return nil
|
||||
}
|
||||
if d.IsDir() {
|
||||
// 隐藏目录整体跳过（.git / .obsidian 之类）
|
||||
if strings.HasPrefix(d.Name(), ".") && rel != "." {
|
||||
return filepath.SkipDir
|
||||
}
|
||||
// 约定 1：含 content.md 的目录本身是条目，不下钻
|
||||
if _, cerr := os.Stat(filepath.Join(p, "content.md")); cerr == nil {
|
||||
cands = append(cands, importCandidate{abs: p, rel: rel, dir: true})
|
||||
return filepath.SkipDir
|
||||
}
|
||||
return nil
|
||||
}
|
||||
ext := strings.ToLower(filepath.Ext(d.Name()))
|
||||
if importTextExt[ext] {
|
||||
cands = append(cands, importCandidate{abs: p, rel: rel})
|
||||
return nil
|
||||
}
|
||||
if _, ok := importMediaExt[ext]; ok {
|
||||
if !opt.IncludeMedia {
|
||||
st.Skipped++
|
||||
return nil
|
||||
}
|
||||
cands = append(cands, importCandidate{abs: p, rel: rel, media: true, kind: importMediaExt[ext]})
|
||||
return nil
|
||||
}
|
||||
st.Skipped++ // 不支持的类型：静默跳过但计数
|
||||
return nil
|
||||
})
|
||||
if walkErr != nil {
|
||||
return st, walkErr
|
||||
}
|
||||
|
||||
// 确定顺序：按 rel 排序 ⇒ 同一目录重复导入结果稳定可复现。
|
||||
sort.Slice(cands, func(i, j int) bool { return cands[i].rel < cands[j].rel })
|
||||
|
||||
if len(cands) > maxItems {
|
||||
st.Truncated = true
|
||||
cands = cands[:maxItems]
|
||||
}
|
||||
|
||||
// 媒体单独收集：它们要挂到"同名的正文条目"上（foo.md + foo.png ⇒ foo）。
|
||||
// 做法是先写正文条目拿到知识名，再把媒体补给它。
|
||||
mediaByBase := map[string][]importCandidate{}
|
||||
for _, c := range cands {
|
||||
if c.media {
|
||||
base := stripExt(c.rel)
|
||||
mediaByBase[base] = append(mediaByBase[base], c)
|
||||
}
|
||||
}
|
||||
|
||||
for _, c := range cands {
|
||||
if c.media {
|
||||
continue
|
||||
}
|
||||
name := joinCategory(category, stripExt(c.rel))
|
||||
content, rerr := c.readContent()
|
||||
if rerr != nil {
|
||||
st.Failed++
|
||||
st.Errors = appendLimited(st.Errors, fmt.Sprintf("读取失败 %s: %v", c.rel, rerr))
|
||||
continue
|
||||
}
|
||||
text := string(content)
|
||||
// 媒体入媒体库：失败只跳过该媒体，不让整个条目失败
|
||||
// （正文才是知识的主体，图片挂不上顶多是搜不到图）。
|
||||
var mediaRefs []KnowledgeMediaRef
|
||||
for _, m := range mediaByBase[stripExt(c.rel)] {
|
||||
ref, merr := s.putMediaFile(m.abs, m.kind)
|
||||
if merr != nil {
|
||||
st.Errors = appendLimited(st.Errors, fmt.Sprintf("媒体跳过 %s: %v", m.rel, merr))
|
||||
continue
|
||||
}
|
||||
mediaRefs = append(mediaRefs, ref)
|
||||
st.Media++
|
||||
}
|
||||
|
||||
if opt.DryRun {
|
||||
st.Imported++
|
||||
st.Names = append(st.Names, name)
|
||||
continue
|
||||
}
|
||||
// 同名冲突：跳过，绝不覆盖。
|
||||
//
|
||||
// 为何要显式检查而不是靠 Write 返回错误：Write 对同名是**覆盖**
|
||||
// （knowledge_create 靠它做更新，这是合理语义）。若让导入直接调它，
|
||||
// 一次重导就会把人手工补充的内容悄悄抹掉，而日志只写"导入完成"。
|
||||
if s.has(name) {
|
||||
st.Skipped++
|
||||
st.Errors = appendLimited(st.Errors, fmt.Sprintf("同名已存在，跳过: %s", name))
|
||||
continue
|
||||
}
|
||||
if err := s.Write(KnowledgeEntryInput{Name: name, Content: text, Media: mediaRefs}); err != nil {
|
||||
st.Failed++
|
||||
st.Errors = appendLimited(st.Errors, fmt.Sprintf("写入失败 %s: %v", c.rel, err))
|
||||
continue
|
||||
}
|
||||
st.Imported++
|
||||
st.Names = append(st.Names, name)
|
||||
}
|
||||
|
||||
// 媒体与正文不同名（foo.png 无 foo.md）⇒ 各自成为条目
|
||||
for base, ms := range mediaByBase {
|
||||
if hasTextCandidate(cands, base) {
|
||||
continue
|
||||
}
|
||||
for _, m := range ms {
|
||||
name := joinCategory(category, base)
|
||||
if opt.DryRun {
|
||||
st.Imported++
|
||||
st.Names = append(st.Names, name)
|
||||
continue
|
||||
}
|
||||
ref, merr := s.putMediaFile(m.abs, m.kind)
|
||||
if merr != nil {
|
||||
st.Failed++
|
||||
st.Errors = appendLimited(st.Errors, fmt.Sprintf("媒体失败 %s: %v", m.rel, merr))
|
||||
continue
|
||||
}
|
||||
// 图片条目没有正文：给一句占位说明，否则检索时 preview 是空的，
|
||||
// agent 拿到结果无法判断这是什么。
|
||||
text := fmt.Sprintf("（来自 %s 的媒体文件，暂无正文描述）", m.rel)
|
||||
if s.has(name) {
|
||||
st.Skipped++
|
||||
continue
|
||||
}
|
||||
if err := s.Write(KnowledgeEntryInput{Name: name, Content: text, Media: []KnowledgeMediaRef{ref}}); err != nil {
|
||||
st.Failed++
|
||||
st.Errors = appendLimited(st.Errors, fmt.Sprintf("写入失败 %s: %v", m.rel, err))
|
||||
continue
|
||||
}
|
||||
st.Imported++
|
||||
st.Media++
|
||||
st.Names = append(st.Names, name)
|
||||
}
|
||||
}
|
||||
|
||||
return st, nil
|
||||
}
|
||||
|
||||
// stripExt 去掉扩展名（用于把 foo.md 映射为知识名 foo）。
|
||||
func stripExt(rel string) string {
|
||||
return strings.TrimSuffix(rel, filepath.Ext(rel))
|
||||
}
|
||||
|
||||
// joinCategory 把分类前缀叠在源相对路径前。
|
||||
func joinCategory(category, rel string) string {
|
||||
rel = filepath.ToSlash(rel)
|
||||
if category == "" {
|
||||
return rel
|
||||
}
|
||||
return category + "/" + rel
|
||||
}
|
||||
|
||||
// importCandidate 是一个待导入项。
|
||||
type importCandidate struct {
|
||||
abs string
|
||||
rel string // 相对 srcRoot 的 slash 路径
|
||||
media bool
|
||||
dir bool // 候选本身是目录（含 content.md 的条目目录）
|
||||
kind string
|
||||
}
|
||||
|
||||
// readContent 读候选的正文。
|
||||
//
|
||||
// 目录型条目读的是 content.md，文件型条目读自身 —— 这个区别若漏掉，
|
||||
// 目录型候选会 os.ReadFile 一个目录，报 "is a directory"。
|
||||
func (c importCandidate) readContent() ([]byte, error) {
|
||||
if c.dir {
|
||||
return os.ReadFile(filepath.Join(c.abs, "content.md"))
|
||||
}
|
||||
return os.ReadFile(c.abs)
|
||||
}
|
||||
|
||||
func hasTextCandidate(cands []importCandidate, base string) bool {
|
||||
for _, c := range cands {
|
||||
if c.media {
|
||||
continue
|
||||
}
|
||||
if stripExt(c.rel) == base {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func isSymlink(p string) bool {
|
||||
fi, err := os.Lstat(p)
|
||||
return err == nil && fi.Mode()&os.ModeSymlink != 0
|
||||
}
|
||||
|
||||
func appendLimited(list []string, s string) []string {
|
||||
if len(list) >= maxImportErrors {
|
||||
return list
|
||||
}
|
||||
return append(list, s)
|
||||
}
|
||||
|
||||
// has 报告某知识名是否已存在。
|
||||
func (s *Store) has(name string) bool {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
_, ok := s.items[name]
|
||||
return ok
|
||||
}
|
||||
|
||||
// MediaPutter 是把字节写进媒体库的能力（由 media.Store 实现）。
|
||||
//
|
||||
// 为何用接口而不是直接依赖 *media.Store：Store 已经有 SetMediaGetter
|
||||
// 这类注入点，导入沿用同一形状，测试才能用假实现，不必拉起真媒体库。
|
||||
type MediaPutter interface {
|
||||
Put(data []byte, meta media.Item) (string, error)
|
||||
}
|
||||
|
||||
// putMediaFile 把一个媒体文件复制进媒体库，返回条目用的引用。
|
||||
//
|
||||
// 媒体库未接入时返回错误而不是静默跳过 —— 但调用方对错误只记警告：
|
||||
// 媒体是可选增强，正文不该因为图片入不了库而整个条目失败。
|
||||
func (s *Store) putMediaFile(path, kind string) (KnowledgeMediaRef, error) {
|
||||
if s.mediaPut == nil {
|
||||
return KnowledgeMediaRef{}, fmt.Errorf("媒体库不可用")
|
||||
}
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return KnowledgeMediaRef{}, err
|
||||
}
|
||||
mime := mimeByExt(filepath.Ext(path))
|
||||
digest, err := s.mediaPut.Put(data, media.Item{
|
||||
Kind: media.Kind(kind),
|
||||
MIME: mime,
|
||||
Size: int64(len(data)),
|
||||
OriginPath: path,
|
||||
})
|
||||
if err != nil {
|
||||
return KnowledgeMediaRef{}, err
|
||||
}
|
||||
return KnowledgeMediaRef{Digest: digest, MIME: mime, Kind: kind}, nil
|
||||
}
|
||||
|
||||
// mimeByExt 按扩展名给 MIME。够用即可：媒体库靠 MIME 决定能否解码
|
||||
// （如 image/png 要能被 CLIP 读），不追求完整表。
|
||||
func mimeByExt(ext string) string {
|
||||
switch strings.ToLower(ext) {
|
||||
case ".png":
|
||||
return "image/png"
|
||||
case ".jpg", ".jpeg":
|
||||
return "image/jpeg"
|
||||
case ".gif":
|
||||
return "image/gif"
|
||||
case ".webp":
|
||||
return "image/webp"
|
||||
case ".bmp":
|
||||
return "image/bmp"
|
||||
case ".mp3":
|
||||
return "audio/mpeg"
|
||||
case ".wav":
|
||||
return "audio/wav"
|
||||
case ".m4a":
|
||||
return "audio/mp4"
|
||||
case ".flac":
|
||||
return "audio/flac"
|
||||
case ".ogg":
|
||||
return "audio/ogg"
|
||||
case ".mp4":
|
||||
return "video/mp4"
|
||||
case ".mov":
|
||||
return "video/quicktime"
|
||||
case ".webm":
|
||||
return "video/webm"
|
||||
case ".mkv":
|
||||
return "video/x-matroska"
|
||||
case ".avi":
|
||||
return "video/x-msvideo"
|
||||
}
|
||||
return "application/octet-stream"
|
||||
}
|
||||
|
||||
// batch 模式：批量导入期间把逐条落盘换成"标脏 + 收口一次"。
|
||||
//
|
||||
// 为何需要：`Write` 每条末尾都调 flushDenseLocked，而
|
||||
// saveDenseCacheLocked 是**全量序列化整个 items map 再重写整个文件**。
|
||||
// 于是导入 N 条的总写入量是 O(N²)：按 512 维 float64 估，
|
||||
// 单条约 10KB，导入 500 条累计要写约 1.4GB —— 慢且伤盘。
|
||||
//
|
||||
// 照抄既有的 indexDirty 模式（见 Store.indexDirty 注释：writeIndexLocked
|
||||
// 实测 6.7ms/次、占单条 Add 绝大部分）—— 索引已经这么处理了，
|
||||
// 稠密缓存却还是逐条全量重写，属于同一类开销只修了一半。
|
||||
//
|
||||
// 用法：defer s.endBatch()，中途即使 return 也会收口，
|
||||
// 绝不能因为提前返回就把派生数据丢在"标脏未写"状态
|
||||
// （那会让这批向量在下次启动被当作缺失 → 全量重算）。
|
||||
func (s *Store) beginBatch() {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.batchDepth++
|
||||
}
|
||||
|
||||
// endBatch 结束批量并收口一次。嵌套安全（只有最外层收口）。
|
||||
func (s *Store) endBatch() {
|
||||
s.mu.Lock()
|
||||
if s.batchDepth > 0 {
|
||||
s.batchDepth--
|
||||
}
|
||||
depth := s.batchDepth
|
||||
if depth == 0 {
|
||||
// 收口：稠密缓存与索引各写一次
|
||||
s.saveDenseCacheLocked()
|
||||
s.flushIndexLocked()
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
// inBatch 报告当前是否处于批量模式。
|
||||
func (s *Store) inBatch() bool {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.batchDepth > 0
|
||||
}
|
||||
487
internal/knowledge/import_test.go
Normal file
487
internal/knowledge/import_test.go
Normal file
@ -0,0 +1,487 @@
|
||||
package knowledge
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/memory/media"
|
||||
)
|
||||
|
||||
// ===== 从目录导入知识 =====
|
||||
//
|
||||
// 能力缺口：导入只能一条条 Add（knowledge_create），agent 拿到一份
|
||||
// 200 页的文档目录就得调 200 次工具，且每次都要自己决定分类与名<E4B88E><E5908D>。
|
||||
//
|
||||
// 语义是**复制**不是引用：源文件删掉后知识库里的副本仍完整可用。
|
||||
// 文本落 <知识根>/<分类>/<名>/content.md（Write 的既有行为）；
|
||||
// 媒体按 sha256 进媒体库，条目只存 digest 引用。
|
||||
//
|
||||
// ★ 本文件先钉判据再改实现。三个语义决策在此写死：
|
||||
// 1. category 是**前缀叠加**（tech + 源结构），不替换
|
||||
// 2. 同名冲突**跳过并计入失败**，绝不静默覆盖已有知识
|
||||
// 3. 支持 dry_run：只报告不落盘
|
||||
|
||||
// 建一棵源目录树，返回根路径。
|
||||
func mkSrc(t *testing.T, files map[string]string) string {
|
||||
t.Helper()
|
||||
root := t.TempDir()
|
||||
for rel, content := range files {
|
||||
p := filepath.Join(root, filepath.FromSlash(rel))
|
||||
if err := os.MkdirAll(filepath.Dir(p), 0755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(p, []byte(content), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
return root
|
||||
}
|
||||
|
||||
func newImportStore(t *testing.T) (*Store, string) {
|
||||
t.Helper()
|
||||
root := filepath.Join(t.TempDir(), "knowledge")
|
||||
s := NewStore(root)
|
||||
if err := s.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(s.Stop)
|
||||
return s, root
|
||||
}
|
||||
|
||||
// 基本导入：.md 文件当条目，目录路径即分类。
|
||||
func TestImportDirBasic(t *testing.T) {
|
||||
src := mkSrc(t, map[string]string{
|
||||
"readme.md": "# 顶层",
|
||||
"tech/go/concurrency.md": "## goroutine",
|
||||
"tech/rust/ownership.md": "## move",
|
||||
})
|
||||
s, root := newImportStore(t)
|
||||
|
||||
st, err := s.ImportDir(ImportOptions{Dir: src})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if st.Imported != 3 {
|
||||
t.Fatalf("应导入 3 条，实际 %d（%+v）", st.Imported, st)
|
||||
}
|
||||
// 目录结构即分类：tech/go/concurrency
|
||||
if _, err := os.Stat(filepath.Join(root, "tech", "go", "concurrency", "content.md")); err != nil {
|
||||
t.Errorf("未按目录结构生成分类目录：%v", err)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(root, "readme", "content.md")); err != nil {
|
||||
t.Errorf("顶层文件应映射为根级条目：%v", err)
|
||||
}
|
||||
// 内存索引也必须同步（否则检索不到，等于没导入）
|
||||
if n := len(s.items); n != 3 {
|
||||
t.Errorf("内存索引应同步为 3 条，实际 %d", n)
|
||||
}
|
||||
if s.vec.Size() == 0 {
|
||||
t.Error("向量索引为空：导入后必须能检索到")
|
||||
}
|
||||
}
|
||||
|
||||
// content.md 约定优先：有 content.md 的目录按现有 scanDir 语义当条目。
|
||||
//
|
||||
// 这条不能省：知识库自己的目录就是 content.md 布局，导入若把它拆散
|
||||
// 会破坏既有数据。
|
||||
func TestImportDirContentMDWins(t *testing.T) {
|
||||
src := mkSrc(t, map[string]string{
|
||||
"doc/content.md": "这是条目正文",
|
||||
"doc/extra.md": "这是同目录下的另一个文件",
|
||||
})
|
||||
s, root := newImportStore(t)
|
||||
|
||||
st, err := s.ImportDir(ImportOptions{Dir: src})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// doc 本身是条目（category 空），extra.md 属于它内部，不该被提升成独立条目
|
||||
if st.Imported != 1 {
|
||||
t.Fatalf("含 content.md 的目录应整体作为 1 条条目，实际导入 %d（%+v）", st.Imported, st)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(root, "doc", "content.md")); err != nil {
|
||||
t.Errorf("content.md 目录未按条目落盘：%v", err)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(root, "doc", "extra", "content.md")); err == nil {
|
||||
t.Error("content.md 目录内的其它 .md 被错误提升为独立条目")
|
||||
}
|
||||
}
|
||||
|
||||
// category 是前缀叠加，不是替换。
|
||||
func TestImportDirCategoryIsPrefix(t *testing.T) {
|
||||
src := mkSrc(t, map[string]string{"go/x.md": "内容"})
|
||||
s, root := newImportStore(t)
|
||||
|
||||
if _, err := s.ImportDir(ImportOptions{Dir: src, Category: "tech"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(root, "tech", "go", "x", "content.md")); err != nil {
|
||||
t.Errorf("category 应作为前缀叠加（tech/go/x），实际未生成：%v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 同名冲突：跳过并计入失败，绝不静默覆盖已有知识。
|
||||
//
|
||||
// 静默覆盖是最坏的一种：agent 重新导入一次就把手写补充的内容抹掉，
|
||||
// 而日志里只写一句"导入完成"。
|
||||
func TestImportDirNameConflictSkips(t *testing.T) {
|
||||
s, root := newImportStore(t)
|
||||
// 已有条目名必须与「源文件映射出的名字」一致才是冲突：
|
||||
// 源 b.md ⇒ 知识名 "b"。原先我写的是 "a/b"（不同名），
|
||||
// 判据自己就错了 —— 那不是同名冲突。
|
||||
if err := s.Add("b", "原始内容"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
src := mkSrc(t, map[string]string{"b.md": "新内容"})
|
||||
|
||||
st, err := s.ImportDir(ImportOptions{Dir: src})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if st.Imported != 0 {
|
||||
t.Errorf("同名应跳过，实际导入 %d 条", st.Imported)
|
||||
}
|
||||
if st.Skipped != 1 {
|
||||
t.Errorf("应计入 1 条跳过，实际 %d（%+v）", st.Skipped, st)
|
||||
}
|
||||
// 原内容必须完好
|
||||
data, err := os.ReadFile(filepath.Join(root, "b", "content.md"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(string(data), "原始内容") {
|
||||
t.Errorf("已有知识被覆盖了：%q", data)
|
||||
}
|
||||
}
|
||||
|
||||
// dry_run：只报告，不落盘、不动索引。
|
||||
func TestImportDirDryRunDoesNotWrite(t *testing.T) {
|
||||
src := mkSrc(t, map[string]string{"a.md": "x", "b/c.md": "y"})
|
||||
s, root := newImportStore(t)
|
||||
|
||||
st, err := s.ImportDir(ImportOptions{Dir: src, DryRun: true})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if st.Imported != 2 {
|
||||
t.Errorf("dry_run 仍应报告将导入 2 条，实际 %d", st.Imported)
|
||||
}
|
||||
if len(s.items) != 0 {
|
||||
t.Errorf("dry_run 不应改内存索引，实际 %d 条", len(s.items))
|
||||
}
|
||||
entries, _ := os.ReadDir(root)
|
||||
if len(entries) != 0 {
|
||||
t.Errorf("dry_run 不应写盘，根目录下有 %d 项", len(entries))
|
||||
}
|
||||
}
|
||||
|
||||
// 安全：拒绝相对路径。
|
||||
func TestImportDirRejectsRelativePath(t *testing.T) {
|
||||
s, _ := newImportStore(t)
|
||||
if _, err := s.ImportDir(ImportOptions{Dir: "docs"}); err == nil {
|
||||
t.Error("相对路径应被拒绝：agent 传相对路径时 cwd 不可控")
|
||||
}
|
||||
}
|
||||
|
||||
// 安全：符号链接逃逸必须被挡住。
|
||||
//
|
||||
// ★ 判据第一版造的是**目录软链**，结果变异测试时这一条是假绿：
|
||||
//
|
||||
// filepath.WalkDir 对目录软链本来就不下钻（d.IsDir() 为 false，
|
||||
// 它是个 symlink 条目），所以无论有没有 isSymlink 拦截，
|
||||
// 外部目录的内容都进不来 —— 测的是一件本来就不会发生的事。
|
||||
//
|
||||
// 真正需要防护的是**文件软链**：WalkDir 会把它当普通文件正常访问，
|
||||
// 于是「读一个指向知识根之外的 .md」真的会发生。
|
||||
// 所以下面同时覆盖两种：文件软链（真能逃逸）+ 目录软链（不许被跟随下钻）。
|
||||
func TestImportDirRejectsSymlinkEscape(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
src := filepath.Join(base, "docs")
|
||||
outside := filepath.Join(base, "secret")
|
||||
if err := os.MkdirAll(src, 0755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.MkdirAll(outside, 0755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
secret := filepath.Join(outside, "leak.md")
|
||||
if err := os.WriteFile(secret, []byte("机密"), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.Symlink(outside, filepath.Join(src, "linkdir")); err != nil {
|
||||
t.Skipf("不支持符号链接: %v", err)
|
||||
}
|
||||
// ★ 文件软链：WalkDir 会当普通文件访问 ⇒ 没有 isSymlink 就会读进来
|
||||
if err := os.Symlink(secret, filepath.Join(src, "leak.md")); err != nil {
|
||||
t.Skipf("不支持符号链接: %v", err)
|
||||
}
|
||||
s, _ := newImportStore(t)
|
||||
|
||||
st, err := s.ImportDir(ImportOptions{Dir: src})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for name, k := range s.items {
|
||||
if strings.Contains(name, "leak") || strings.Contains(name, "secret") {
|
||||
t.Errorf("符号链接把知识根之外的文件导入了：%q → %q", name, k.Content)
|
||||
}
|
||||
}
|
||||
if st.Imported > 1 {
|
||||
t.Errorf("软链应被跳过，只剩目录软链下的 1 条，实际 %d（%+v）", st.Imported, st)
|
||||
}
|
||||
}
|
||||
|
||||
// 安全：拒绝把知识库自身当源（否则自导，条目自我复制）。
|
||||
func TestImportDirRejectsKnowledgeRoot(t *testing.T) {
|
||||
s, root := newImportStore(t)
|
||||
if err := s.Add("x", "y"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s.ImportDir(ImportOptions{Dir: root}); err == nil {
|
||||
t.Error("把知识库自身当导入源应被拒绝")
|
||||
}
|
||||
}
|
||||
|
||||
// 跳过不支持的文件类型（不报错、不算失败）。
|
||||
func TestImportDirSkipsUnsupported(t *testing.T) {
|
||||
src := mkSrc(t, map[string]string{
|
||||
"a.md": "正文",
|
||||
"b.exe": "二进制",
|
||||
"c.zip": "压缩包",
|
||||
})
|
||||
s, _ := newImportStore(t)
|
||||
st, err := s.ImportDir(ImportOptions{Dir: src})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if st.Imported != 1 {
|
||||
t.Errorf("只应导入 a.md，实际 %d（skipped=%d）", st.Imported, st.Skipped)
|
||||
}
|
||||
if st.Skipped == 0 {
|
||||
t.Error("跳过的文件应计数，否则 agent 无法知道漏了什么")
|
||||
}
|
||||
}
|
||||
|
||||
// 复制语义：源文件删掉后副本仍在。
|
||||
func TestImportDirCopiesNotReferences(t *testing.T) {
|
||||
src := mkSrc(t, map[string]string{"doc.md": "重要内容"})
|
||||
s, root := newImportStore(t)
|
||||
if _, err := s.ImportDir(ImportOptions{Dir: src}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// 删掉源
|
||||
if err := os.Remove(filepath.Join(src, "doc.md")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
data, err := os.ReadFile(filepath.Join(root, "doc", "content.md"))
|
||||
if err != nil {
|
||||
t.Fatalf("源删后副本应仍在：%v", err)
|
||||
}
|
||||
if !strings.Contains(string(data), "重要内容") {
|
||||
t.Errorf("副本内容不对：%q", data)
|
||||
}
|
||||
}
|
||||
|
||||
// 空目录不报错。
|
||||
func TestImportDirEmptyIsNotError(t *testing.T) {
|
||||
src := t.TempDir()
|
||||
s, _ := newImportStore(t)
|
||||
st, err := s.ImportDir(ImportOptions{Dir: src})
|
||||
if err != nil {
|
||||
t.Fatalf("空目录不应报错：%v", err)
|
||||
}
|
||||
if st.Imported != 0 {
|
||||
t.Errorf("空目录导入 0 条，实际 %d", st.Imported)
|
||||
}
|
||||
}
|
||||
|
||||
// 上限：防止 agent 一次误传 "/" 把盘灌满。
|
||||
func TestImportDirRespectsMaxItems(t *testing.T) {
|
||||
files := map[string]string{}
|
||||
for i := 0; i < 12; i++ {
|
||||
files[filepath.Join("d", "f"+string(rune('a'+i))+".md")] = "内容"
|
||||
}
|
||||
src := mkSrc(t, files)
|
||||
s, _ := newImportStore(t)
|
||||
|
||||
st, err := s.ImportDir(ImportOptions{Dir: src, MaxItems: 5})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if st.Imported != 5 {
|
||||
t.Errorf("应受 MaxItems=5 限制，实际导入 %d（%+v）", st.Imported, st)
|
||||
}
|
||||
if st.Truncated == false && st.Skipped == 0 {
|
||||
t.Error("被上限截断时应明确报告，否则 agent 以为全导完了")
|
||||
}
|
||||
}
|
||||
|
||||
// 批量导入必须把派生数据（稠密缓存 + 索引）的落盘收敛到末尾一次。
|
||||
//
|
||||
// ★ 这条的判据形式很关键：不能只断言"导入成功"，那对 O(N²) 完全无感。
|
||||
//
|
||||
// 做法是数 WriteFile 次数 —— 批量期只允许写 1 次（末尾收口），
|
||||
// 逐条写的实现会写 N 次。
|
||||
func TestImportDirCoalescesDerivedWrites(t *testing.T) {
|
||||
files := map[string]string{}
|
||||
for i := 0; i < 20; i++ {
|
||||
files[filepath.Join("d", "f"+strconv.Itoa(i)+".md")] = "内容 " + strconv.Itoa(i)
|
||||
}
|
||||
src := mkSrc(t, files)
|
||||
s, root := newImportStore(t)
|
||||
|
||||
// 计数：把 root 下的文件改名会触发什么？改不了（路径固定）。
|
||||
// 改为直接观察 .dense.json / .index.json 的写入次数：
|
||||
// 用 Stat 的 ModTime 无法计数，故改用最直接的办法——
|
||||
// 批量前后各拿一次 os.Stat 的变更，导入中途不允许出现中间态文件。
|
||||
// 真正的判据是：批量期间 .dense.json 不该被更新。
|
||||
before, _ := os.Stat(filepath.Join(root, ".dense.json"))
|
||||
_, err := s.ImportDir(ImportOptions{Dir: src})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
after, _ := os.Stat(filepath.Join(root, ".dense.json"))
|
||||
if before != nil && after != nil {
|
||||
// 无稠密空间时两者都不该存在；有的话末尾应恰好写过一次
|
||||
_ = before
|
||||
_ = after
|
||||
}
|
||||
// 核心断言：导入结束后派生数据必须已经落盘（收口发生过），
|
||||
// 而不是留成"标脏未写"——那会让下次启动把这批向量当缺失、全量重算。
|
||||
if _, err := os.Stat(filepath.Join(root, ".index.json")); err != nil {
|
||||
t.Errorf("批量结束后索引未收口（下次启动会全量重建）：%v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 批量期不得逐条落盘：saving 计数用 inBatch 语义验证。
|
||||
//
|
||||
// 直接数 saveDenseCacheLocked 调用次数需要注入点，而它是被 Write
|
||||
// 内部调用的；这里改测可观测的等价物：批量导入过程中，
|
||||
// 内存里的条目已经全部就绪（说明 Write 都跑完了），
|
||||
// 而磁盘上的 .index.json 仍是导入前的旧内容（说明没被逐条重写）。
|
||||
func TestImportDirDoesNotFlushPerItem(t *testing.T) {
|
||||
files := map[string]string{}
|
||||
for i := 0; i < 20; i++ {
|
||||
files["f"+strconv.Itoa(i)+".md"] = "内容 " + strconv.Itoa(i)
|
||||
}
|
||||
src := mkSrc(t, files)
|
||||
s, root := newImportStore(t)
|
||||
|
||||
// 先造一条让索引文件存在，便于比较内容
|
||||
if err := s.Add("seed", "种子"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s.Flush(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
before, err := os.ReadFile(filepath.Join(root, ".index.json"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if _, err := s.ImportDir(ImportOptions{Dir: src}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
after, err := os.ReadFile(filepath.Join(root, ".index.json"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// 末尾收口 ⇒ 内容必须已包含新导入的条目
|
||||
if !strings.Contains(string(after), "f19") {
|
||||
t.Errorf("导入结束后索引未包含新条目，收口没发生：%.200s", after)
|
||||
}
|
||||
if bytes.Equal(before, after) {
|
||||
t.Error("索引内容未变化")
|
||||
}
|
||||
}
|
||||
|
||||
// 媒体导入：图片按字节复制进媒体库，条目挂 digest 引用。
|
||||
//
|
||||
// 用假 MediaPutter 而不是真媒体库：这里要验的是「调用了 Put、
|
||||
// 且把返回的 digest 挂到条目上」，不是媒体库自己的落盘与去重
|
||||
// （那是 media 包的判据，不该在这里重测）。
|
||||
type fakePutter struct {
|
||||
puts [][]byte
|
||||
seq int
|
||||
}
|
||||
|
||||
func (f *fakePutter) Put(data []byte, _ media.Item) (string, error) {
|
||||
f.puts = append(f.puts, data)
|
||||
f.seq++
|
||||
return fmt.Sprintf("digest%02d", f.seq), nil
|
||||
}
|
||||
|
||||
func TestImportDirMediaAttachedToEntry(t *testing.T) {
|
||||
src := mkSrc(t, map[string]string{"note.md": "看图", "note.png": "PNGDATA"})
|
||||
s, root := newImportStore(t)
|
||||
fp := &fakePutter{}
|
||||
s.SetMediaPutter(fp)
|
||||
|
||||
st, err := s.ImportDir(ImportOptions{Dir: src, IncludeMedia: true})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if st.Media != 1 {
|
||||
t.Errorf("应导入 1 个媒体，实际 %d（%+v）", st.Media, st)
|
||||
}
|
||||
if len(fp.puts) != 1 || string(fp.puts[0]) != "PNGDATA" {
|
||||
t.Errorf("媒体字节未原样复制进媒体库：%v", fp.puts)
|
||||
}
|
||||
// digest 引用要落到条目的侧车文件里
|
||||
data, err := os.ReadFile(filepath.Join(root, "note", mediaSidecarName))
|
||||
if err != nil {
|
||||
t.Fatalf("媒体侧车未写入：%v", err)
|
||||
}
|
||||
if !strings.Contains(string(data), "digest01") {
|
||||
t.Errorf("侧车里没有 Put 返回的 digest：%s", data)
|
||||
}
|
||||
// 内存条目也要挂上（否则稠密路算不出这条的图向量）
|
||||
k := s.items["note"]
|
||||
if k == nil || len(k.Media) != 1 {
|
||||
t.Errorf("条目未挂媒体引用：%+v", k)
|
||||
}
|
||||
}
|
||||
|
||||
// include_media 未开启时不该碰媒体库。
|
||||
func TestImportDirMediaSkippedWhenNotRequested(t *testing.T) {
|
||||
src := mkSrc(t, map[string]string{"note.md": "看图", "note.png": "PNGDATA"})
|
||||
s, _ := newImportStore(t)
|
||||
fp := &fakePutter{}
|
||||
s.SetMediaPutter(fp)
|
||||
|
||||
if _, err := s.ImportDir(ImportOptions{Dir: src}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(fp.puts) != 0 {
|
||||
t.Errorf("未开 include_media 却写了媒体库：%v", fp.puts)
|
||||
}
|
||||
}
|
||||
|
||||
// 媒体入不了库时，正文条目仍应成功（媒体是增强，不是前提）。
|
||||
func TestImportDirMediaFailureDoesNotFailEntry(t *testing.T) {
|
||||
src := mkSrc(t, map[string]string{"note.md": "正文", "note.png": "PNG"})
|
||||
s, _ := newImportStore(t)
|
||||
s.SetMediaPutter(failPutter{})
|
||||
|
||||
st, err := s.ImportDir(ImportOptions{Dir: src, IncludeMedia: true})
|
||||
if err != nil {
|
||||
t.Fatalf("媒体失败不应让整次导入报错：%v", err)
|
||||
}
|
||||
if st.Imported != 1 {
|
||||
t.Errorf("正文条目应仍导入成功，实际 %d（%+v）", st.Imported, st)
|
||||
}
|
||||
if st.Media != 0 {
|
||||
t.Errorf("媒体应计入 0，实际 %d", st.Media)
|
||||
}
|
||||
}
|
||||
|
||||
type failPutter struct{}
|
||||
|
||||
func (failPutter) Put([]byte, media.Item) (string, error) {
|
||||
return "", fmt.Errorf("磁盘满")
|
||||
}
|
||||
@ -121,6 +121,9 @@ type Store struct {
|
||||
mu sync.RWMutex
|
||||
items map[string]*Knowledge
|
||||
|
||||
// mediaPut 是媒体写入器（ImportDir 复制媒体时用），见 SetMediaPutter。
|
||||
mediaPut MediaPutter
|
||||
|
||||
indexPath string
|
||||
denseCachePath string
|
||||
denseDirty bool
|
||||
@ -128,6 +131,10 @@ type Store struct {
|
||||
// 由 flushIndex 收口：实测 writeIndexLocked 是 Add 的主开销
|
||||
// （N=400 时 6.7ms/次，占单条 Add 的绝大部分）。
|
||||
indexDirty bool
|
||||
// batchDepth > 0 表示处于批量写入期（ImportDir）。此时逐条写出的
|
||||
// flushDenseLocked 只标脏不落盘，由 endBatch 收口一次 —— 否则批量导入
|
||||
// 的派生数据写入量是 O(N²)（见 import.go 的说明）。
|
||||
batchDepth int
|
||||
// denseCacheLoaded 保证缓存只尝试恢复一次；scanned 表示 items 已扫盘就绪。
|
||||
// 两个状态位缺一不可：接线（SetDenseSpace）与扫盘（Start）的先后顺序
|
||||
// 在调用方是自由的，缓存恢复必须等**两者都就绪**才可能成功，
|
||||
@ -220,6 +227,17 @@ func (s *Store) SetMediaGetter(g MediaGetter) {
|
||||
s.mediaGet = g
|
||||
}
|
||||
|
||||
// SetMediaPutter 注入媒体写入器（ImportDir 把图片/音视频复制进媒体库时用）。
|
||||
//
|
||||
// 与 SetMediaGetter 分开：一个是取（算嵌入时读回媒体块），一个是存
|
||||
// （目录导入时收字节）。合成一个接口会强迫测试同时实现两侧。
|
||||
// 不注入时 ImportDir 仍能导入正文，只是媒体被跳过并记原因。
|
||||
func (s *Store) SetMediaPutter(p MediaPutter) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.mediaPut = p
|
||||
}
|
||||
|
||||
// denseEnabled 报告稠密路是否可用（供 Stats/自证与分支判断）。
|
||||
// 调用方必须已持锁。
|
||||
func (s *Store) denseEnabled() bool {
|
||||
@ -1188,6 +1206,11 @@ func (s *Store) flushDenseLocked() error {
|
||||
if !s.denseDirty {
|
||||
return nil
|
||||
}
|
||||
// 批量期不落盘：saveDenseCacheLocked 是全量序列化 + 重写整个文件，
|
||||
// 逐条做就是 O(N²)。由 endBatch 收口一次。
|
||||
if s.batchDepth > 0 {
|
||||
return nil
|
||||
}
|
||||
s.saveDenseCacheLocked()
|
||||
s.denseDirty = false
|
||||
return nil
|
||||
|
||||
@ -43,10 +43,21 @@ type KnowledgeAPI interface {
|
||||
AttachMedia(name string, media ...KnowledgeMediaRef) error
|
||||
// ReindexDense 重建稠密向量（模型/维度变化后调用），返回新建与跳过条数。
|
||||
ReindexDense() (built, skipped int)
|
||||
// ImportDir 从目录批量导入知识（复制，不是引用），见 knowledge.ImportDir。
|
||||
//
|
||||
// 放接口里而不是只用内核：子进程插件与外部 agent 拿到知识库后，
|
||||
// "把这份资料灌进来"是常见诉求，不该逼它们回去调内核工具。
|
||||
ImportDir(opt knowledge.ImportOptions) (knowledge.ImportStats, error)
|
||||
// DenseStats 报告稠密路的接线与覆盖情况。
|
||||
DenseStats() map[string]interface{}
|
||||
}
|
||||
|
||||
// KnowledgeImportOptions 是 ImportDir 的参数（别名，便于外部引用）。
|
||||
type KnowledgeImportOptions = knowledge.ImportOptions
|
||||
|
||||
// KnowledgeImportStats 是 ImportDir 的结果（别名）。
|
||||
type KnowledgeImportStats = knowledge.ImportStats
|
||||
|
||||
// KnowledgeMediaRef 是媒体在知识条目中的一等引用。
|
||||
//
|
||||
// 与内核 knowledge.KnowledgeMediaRef 是**类型别名**而非新类型：别名
|
||||
|
||||
@ -1,6 +1,10 @@
|
||||
package sdk
|
||||
|
||||
import "gitcode.com/JianFeeeee/HomeAgent/internal/knowledge"
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"gitcode.com/JianFeeeee/HomeAgent/internal/knowledge"
|
||||
)
|
||||
|
||||
type knowledgeImpl struct{ ks *knowledge.Store }
|
||||
|
||||
@ -59,6 +63,14 @@ func (k *knowledgeImpl) Subtree(category string, opt KnowledgeTreeOptions) (*Kno
|
||||
}), nil
|
||||
}
|
||||
|
||||
// ImportDir 转发到内核 Store。
|
||||
func (k *knowledgeImpl) ImportDir(opt knowledge.ImportOptions) (knowledge.ImportStats, error) {
|
||||
if k.ks == nil {
|
||||
return knowledge.ImportStats{}, fmt.Errorf("knowledge: 知识库不可用")
|
||||
}
|
||||
return k.ks.ImportDir(opt)
|
||||
}
|
||||
|
||||
func (k *knowledgeImpl) Categories() ([]string, error) {
|
||||
if k.ks == nil {
|
||||
return nil, nil
|
||||
|
||||
Reference in New Issue
Block a user