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 }