mirror of
https://gitcode.com/JianFeeeee/HomeAgent.git
synced 2026-09-21 17:38:10 +00:00
feat(memory): 内容寻址媒体存储（CAS）——图记忆支持二进制多媒体节点的底座
此前四层记忆全是纯文本载体，没有任何一层能存二进制：
L0 ContextEvent — Input/Response/ToolResults[].Output 全 string
L1 text.Event — 同上
L2 document.Doc — Summary/Content/Tags 全 string
L3 图库 — sentences.text TEXT UNIQUE，节点身份就是那串文本
于是 multimodal 插件注入的图只在本轮对话内可见（走 message 数组，不经
记忆），下一轮起只剩 ToolResultItem.Output 里那句
"[已将图片注入后续对话] /tmp/x.png"——一条路径字符串。那个文件被删或
被覆盖之后连线索都断了。
## 为何内容寻址而不是存路径
- 路径会失效。/tmp 下的探针图、下载缓存、别的进程的临时产物，记忆里
留个路径等于留个悬空指针。
- 同一内容常被反复注入（连问几轮同一张截图、see_video 相邻帧高度相似），
按 sha256 寻址天然去重。
- 内容即身份，与 L3 图库 sentences.text UNIQUE 思路一致：文本节点用文本
本身做身份，媒体节点用内容摘要做身份。
## 结构
元数据（SQLite media.db）与内容（磁盘 blobs/ 两级前缀分桶）分离，不把
blob 塞进库：单张图动辄几 MB，塞进去让每次 VACUUM/备份都拖着几百 MB 走，
WAL 也会迅速膨胀。
media(digest PK, kind, mime, size, width, height, origin_path, tool,
description, described_by, ref_count, first_seen, last_seen)
media_refs(digest, owner_kind, owner_id, created_at, PK 三列)
digest 既是主键也是文件名，所以没有 Path 字段——路径由 digest 推导，
不落库（落了就又是个会失效的引用）。origin_path 仅供人类溯源，注释里
明确标注不可用于读取。owner_kind 预留 context/document/graph_sentence。
## 几处刻意的决定
- Get 强制校验 digest：CAS 的全部保证建立在「文件名 == 内容摘要」上，
位翻转或外部误改必须被发现——把损坏的图喂给模型只会得到无从追溯的幻觉。
- 先写 .tmp 再 rename：中途崩溃不留半个 blob 被当成完整内容读走。
- AddRef 幂等：只有真插进 media_refs 才递增，否则计数虚高会让 GC 永远
不敢清。DropRef 用 MAX(0,...) 兜底防负数。
- Put 的空描述不冲掉已有描述（先到的可能来自更强的模型），但尺寸/工具名
这类前一次缺失的信息会被补写。Describe 是显式操作，允许覆盖。
- GC 两段 + minAge 保护：刚 Put 还没 AddRef 的项 refcount 也是 0，minAge
防「落地后还没挂上就被清掉」。有引用的项永不删除，即使超容量——宁可
超限也不断引用。
## 测试
21 例，覆盖去重 / MIME 归类 / 损坏检测 / 无残留临时文件 / AddRef 幂等 /
计数不为负 / DropOwner / GC 保留有引用项 / minAge 保护 / 容量淘汰 /
描述覆盖与补写策略 / Search 按描述与 kind 过滤 / Pending / data URL
往返 / Stats / 跨重启持久化。
本 commit 只加存储层，尚未接入 L0/L2/L3 与描述生成。
This commit is contained in:
JianFeeeee
14
internal/memory/media/base64.go
Normal file
14
internal/memory/media/base64.go
Normal file
@ -0,0 +1,14 @@
|
||||
package media
|
||||
|
||||
import "encoding/base64"
|
||||
|
||||
// base64 编解码单独抽出来，让 media.go 的 import 块只留业务依赖。
|
||||
// 用 StdEncoding：data URL 规范用的是标准表（含 + / =），不是 URL-safe 表。
|
||||
|
||||
func base64Decode(s string) ([]byte, error) {
|
||||
return base64.StdEncoding.DecodeString(s)
|
||||
}
|
||||
|
||||
func base64Encode(b []byte) string {
|
||||
return base64.StdEncoding.EncodeToString(b)
|
||||
}
|
||||
702
internal/memory/media/media.go
Normal file
702
internal/memory/media/media.go
Normal file
@ -0,0 +1,702 @@
|
||||
// Package media 是记忆系统的内容寻址媒体存储（CAS）。
|
||||
//
|
||||
// 为何需要它：此前四层记忆全是纯文本载体——L0 `ContextEvent`、L1 `text.Event`、
|
||||
// L2 `document.Doc`、L3 图库的 `sentences.text TEXT UNIQUE`——没有任何一层能
|
||||
// 存二进制。multimodal 插件注入的图片在本轮对话内可见（走 message 数组，不经
|
||||
// 记忆），下一轮起就只剩 `ToolResultItem.Output` 里那句
|
||||
// "[已将图片注入后续对话] /tmp/x.png"，即一条路径字符串。那个文件被删或被
|
||||
// 覆盖之后连线索都断了。
|
||||
//
|
||||
// 为何是内容寻址而不是存路径：
|
||||
// - 路径会失效。/tmp 下的探针图、下载缓存、其他进程的临时产物，记忆里留个
|
||||
// 路径等于留个悬空指针。
|
||||
// - 同一张图往往被反复注入（用户连问几轮同一张截图、see_video 相邻帧高度
|
||||
// 相似）。按 sha256 寻址天然去重，引用计数记住被引了几次。
|
||||
// - 内容即身份，跟 L3 图库 `sentences.text UNIQUE` 的思路一致：文本节点用
|
||||
// 文本本身做身份，媒体节点用内容摘要做身份。
|
||||
package media
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
_ "github.com/mattn/go-sqlite3"
|
||||
)
|
||||
|
||||
// Kind 是媒体大类。刻意只分三类而不细分具体格式：
|
||||
// 记忆检索关心的是"这是张图还是段音频"，具体编码交给 MIME 字段。
|
||||
type Kind string
|
||||
|
||||
const (
|
||||
KindImage Kind = "image"
|
||||
KindAudio Kind = "audio"
|
||||
KindVideo Kind = "video"
|
||||
KindOther Kind = "other"
|
||||
)
|
||||
|
||||
// Item 是一条媒体记录。
|
||||
//
|
||||
// Digest 既是主键也是磁盘文件名，所以没有单独的 Path 字段——路径可由
|
||||
// Store 根据 Digest 推导，不落库（落了就又是个会失效的引用）。
|
||||
type Item struct {
|
||||
// Digest 是内容 sha256 的十六进制串（64 字符），媒体的唯一身份。
|
||||
Digest string `json:"digest"`
|
||||
// Kind 是大类，供检索时按模态筛选。
|
||||
Kind Kind `json:"kind"`
|
||||
// MIME 是原始 MIME 类型，如 image/png。
|
||||
MIME string `json:"mime"`
|
||||
// Size 是字节数。
|
||||
Size int64 `json:"size"`
|
||||
// Width/Height 是像素尺寸，未知或不适用时为 0。
|
||||
Width int `json:"width,omitempty"`
|
||||
Height int `json:"height,omitempty"`
|
||||
// OriginPath 是首次入库时的来源路径，仅供人类溯源与调试。
|
||||
// **不可用于读取内容**——它随时可能失效，这正是本包存在的理由。
|
||||
OriginPath string `json:"origin_path,omitempty"`
|
||||
// Tool 是注入这条媒体的工具名（如 multimodal_see_picture）。
|
||||
Tool string `json:"tool,omitempty"`
|
||||
// Description 是视觉/音频模型生成的文字描述，供 L2/L3 检索。
|
||||
// 空表示未描述（未开启描述、模型不可用或描述失败）。
|
||||
Description string `json:"description,omitempty"`
|
||||
// DescribedBy 记录描述来自哪个源，让后续读者能判断可靠性。
|
||||
DescribedBy string `json:"described_by,omitempty"`
|
||||
// RefCount 是引用计数。GC 只清理归零的项。
|
||||
RefCount int `json:"ref_count"`
|
||||
// FirstSeen/LastSeen 是首末次入库时间。
|
||||
FirstSeen time.Time `json:"first_seen"`
|
||||
LastSeen time.Time `json:"last_seen"`
|
||||
}
|
||||
|
||||
// Store 管理媒体的元数据（SQLite）与内容（磁盘 CAS 目录）。
|
||||
//
|
||||
// 元数据与内容分离而不是把 blob 塞进 SQLite：单张图动辄几 MB，塞进库会让
|
||||
// 每次 VACUUM/备份都拖着几百 MB 走，也让 WAL 迅速膨胀。CAS 目录用两级
|
||||
// 前缀分桶（ab/cdef...）避免单目录几万文件。
|
||||
type Store struct {
|
||||
mu sync.RWMutex
|
||||
db *sql.DB
|
||||
blobDir string
|
||||
|
||||
// maxBytes 是内容目录的容量上限，0 表示不限。
|
||||
// 超限时 GC 按 LastSeen 从旧到新淘汰 RefCount=0 的项。
|
||||
maxBytes int64
|
||||
}
|
||||
|
||||
// New 打开（或初始化）媒体存储。
|
||||
// dir 下会建 media.db 与 blobs/ 两个条目。
|
||||
func New(dir string, maxBytes int64) (*Store, error) {
|
||||
if err := os.MkdirAll(filepath.Join(dir, "blobs"), 0755); err != nil {
|
||||
return nil, fmt.Errorf("media: create blob dir: %w", err)
|
||||
}
|
||||
dbPath := filepath.Join(dir, "media.db")
|
||||
db, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL&_busy_timeout=5000")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("media: open db: %w", err)
|
||||
}
|
||||
s := &Store{db: db, blobDir: filepath.Join(dir, "blobs"), maxBytes: maxBytes}
|
||||
if err := s.initSchema(); err != nil {
|
||||
db.Close()
|
||||
return nil, err
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func (s *Store) initSchema() error {
|
||||
stmts := []string{
|
||||
// digest 作主键：内容即身份，重复 Put 同一内容只递增 ref_count。
|
||||
`CREATE TABLE IF NOT EXISTS media (
|
||||
digest TEXT PRIMARY KEY,
|
||||
kind TEXT NOT NULL,
|
||||
mime TEXT NOT NULL,
|
||||
size INTEGER NOT NULL,
|
||||
width INTEGER DEFAULT 0,
|
||||
height INTEGER DEFAULT 0,
|
||||
origin_path TEXT,
|
||||
tool TEXT,
|
||||
description TEXT,
|
||||
described_by TEXT,
|
||||
ref_count INTEGER DEFAULT 0,
|
||||
first_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
last_seen TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_media_kind ON media(kind)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_media_refcount ON media(ref_count)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_media_last_seen ON media(last_seen)`,
|
||||
// 反向索引：哪条记忆引用了哪个媒体。
|
||||
// owner_kind 取 context / document / graph_sentence，owner_id 是各层自己的标识。
|
||||
// 主键含三列，同一 owner 重复挂同一媒体是幂等的。
|
||||
`CREATE TABLE IF NOT EXISTS media_refs (
|
||||
digest TEXT NOT NULL,
|
||||
owner_kind TEXT NOT NULL,
|
||||
owner_id TEXT NOT NULL,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
PRIMARY KEY (digest, owner_kind, owner_id)
|
||||
)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_refs_owner ON media_refs(owner_kind, owner_id)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_refs_digest ON media_refs(digest)`,
|
||||
}
|
||||
for _, q := range stmts {
|
||||
if _, err := s.db.Exec(q); err != nil {
|
||||
return fmt.Errorf("media: schema %q: %w", truncate(q, 60), err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// blobPath 按两级前缀分桶推导内容路径。
|
||||
func (s *Store) blobPath(digest string) string {
|
||||
if len(digest) < 4 {
|
||||
return filepath.Join(s.blobDir, digest)
|
||||
}
|
||||
return filepath.Join(s.blobDir, digest[:2], digest[2:])
|
||||
}
|
||||
|
||||
// Put 落盘并登记一段媒体内容，返回其 digest。
|
||||
//
|
||||
// 幂等：同一内容重复 Put 不重复落盘，只更新 last_seen 与可选的新元数据
|
||||
// （描述、尺寸等——后来者可能带着前一次没有的信息）。
|
||||
func (s *Store) Put(data []byte, meta Item) (string, error) {
|
||||
if len(data) == 0 {
|
||||
return "", fmt.Errorf("media: empty content")
|
||||
}
|
||||
sum := sha256.Sum256(data)
|
||||
digest := hex.EncodeToString(sum[:])
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
path := s.blobPath(digest)
|
||||
if _, err := os.Stat(path); os.IsNotExist(err) {
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
|
||||
return "", fmt.Errorf("media: mkdir: %w", err)
|
||||
}
|
||||
// 先写临时文件再 rename：中途崩溃不会留下半个 blob 被后续
|
||||
// 当成完整内容读走（digest 校验能发现，但那时已经把坏数据喂给模型了）。
|
||||
tmp := path + ".tmp"
|
||||
if err := os.WriteFile(tmp, data, 0644); err != nil {
|
||||
return "", fmt.Errorf("media: write blob: %w", err)
|
||||
}
|
||||
if err := os.Rename(tmp, path); err != nil {
|
||||
os.Remove(tmp)
|
||||
return "", fmt.Errorf("media: commit blob: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
if meta.Kind == "" {
|
||||
meta.Kind = KindFromMIME(meta.MIME)
|
||||
}
|
||||
_, err := s.db.Exec(`
|
||||
INSERT INTO media (digest, kind, mime, size, width, height,
|
||||
origin_path, tool, description, described_by,
|
||||
ref_count, first_seen, last_seen)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0, ?, ?)
|
||||
ON CONFLICT(digest) DO UPDATE SET
|
||||
last_seen = excluded.last_seen,
|
||||
-- 只在原值为空时补写：先到的描述可能来自更强的模型，
|
||||
-- 后到的空值不该把它冲掉。
|
||||
description = CASE WHEN COALESCE(media.description,'') = '' THEN excluded.description ELSE media.description END,
|
||||
described_by = CASE WHEN COALESCE(media.described_by,'') = '' THEN excluded.described_by ELSE media.described_by END,
|
||||
width = CASE WHEN media.width = 0 THEN excluded.width ELSE media.width END,
|
||||
height = CASE WHEN media.height = 0 THEN excluded.height ELSE media.height END,
|
||||
tool = CASE WHEN COALESCE(media.tool,'') = '' THEN excluded.tool ELSE media.tool END
|
||||
`, digest, string(meta.Kind), meta.MIME, int64(len(data)), meta.Width, meta.Height,
|
||||
meta.OriginPath, meta.Tool, meta.Description, meta.DescribedBy, now, now)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("media: upsert meta: %w", err)
|
||||
}
|
||||
return digest, nil
|
||||
}
|
||||
|
||||
// Get 读取内容并校验 digest。
|
||||
//
|
||||
// 校验不是多余的：CAS 的全部保证建立在"文件名 == 内容摘要"上，磁盘位翻转
|
||||
// 或外部误改会让这条保证失效，而把损坏的图喂给模型只会得到无从追溯的幻觉。
|
||||
func (s *Store) Get(digest string) ([]byte, error) {
|
||||
s.mu.RLock()
|
||||
path := s.blobPath(digest)
|
||||
s.mu.RUnlock()
|
||||
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("media: read %s: %w", shortDigest(digest), err)
|
||||
}
|
||||
sum := sha256.Sum256(data)
|
||||
if got := hex.EncodeToString(sum[:]); got != digest {
|
||||
return nil, fmt.Errorf("media: digest mismatch for %s (content corrupted)", shortDigest(digest))
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
|
||||
// Stat 返回元数据，不读内容。
|
||||
func (s *Store) Stat(digest string) (*Item, error) {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.scanOne(s.db.QueryRow(`
|
||||
SELECT digest, kind, mime, size, width, height, origin_path, tool,
|
||||
description, described_by, ref_count, first_seen, last_seen
|
||||
FROM media WHERE digest = ?`, digest))
|
||||
}
|
||||
|
||||
// Describe 写入（或覆盖）文字描述。
|
||||
//
|
||||
// 与 Put 的"只在空时补写"不同：Describe 是显式操作，调用方明确想要这份
|
||||
// 描述生效（例如换了更强的视觉模型重新描述）。
|
||||
func (s *Store) Describe(digest, description, describedBy string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
res, err := s.db.Exec(`UPDATE media SET description = ?, described_by = ? WHERE digest = ?`,
|
||||
description, describedBy, digest)
|
||||
if err != nil {
|
||||
return fmt.Errorf("media: describe: %w", err)
|
||||
}
|
||||
if n, _ := res.RowsAffected(); n == 0 {
|
||||
return fmt.Errorf("media: describe: unknown digest %s", shortDigest(digest))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// AddRef 登记一条引用并递增计数。幂等：同一 (digest, owner) 重复调用不重复计数。
|
||||
func (s *Store) AddRef(digest, ownerKind, ownerID string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
tx, err := s.db.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
res, err := tx.Exec(`INSERT OR IGNORE INTO media_refs (digest, owner_kind, owner_id) VALUES (?, ?, ?)`,
|
||||
digest, ownerKind, ownerID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("media: add ref: %w", err)
|
||||
}
|
||||
// 只有真的插进去才递增：否则重复调用会让计数虚高，GC 永远不敢清。
|
||||
if n, _ := res.RowsAffected(); n > 0 {
|
||||
if _, err := tx.Exec(`UPDATE media SET ref_count = ref_count + 1 WHERE digest = ?`, digest); err != nil {
|
||||
return fmt.Errorf("media: bump refcount: %w", err)
|
||||
}
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// DropRef 注销一条引用并递减计数。内容不立即删除，留给 GC。
|
||||
func (s *Store) DropRef(digest, ownerKind, ownerID string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
tx, err := s.db.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
res, err := tx.Exec(`DELETE FROM media_refs WHERE digest = ? AND owner_kind = ? AND owner_id = ?`,
|
||||
digest, ownerKind, ownerID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("media: drop ref: %w", err)
|
||||
}
|
||||
if n, _ := res.RowsAffected(); n > 0 {
|
||||
// MAX(0, ...) 兜底：历史数据或并发意外让计数与 refs 表不一致时，
|
||||
// 不让它掉成负数（负数会让容量 GC 的排序失去意义）。
|
||||
if _, err := tx.Exec(`UPDATE media SET ref_count = MAX(0, ref_count - 1) WHERE digest = ?`, digest); err != nil {
|
||||
return fmt.Errorf("media: lower refcount: %w", err)
|
||||
}
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
// DropOwner 注销某个 owner 的全部引用（该条记忆被删/被归档替换时用）。
|
||||
func (s *Store) DropOwner(ownerKind, ownerID string) (int, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
rows, err := s.db.Query(`SELECT digest FROM media_refs WHERE owner_kind = ? AND owner_id = ?`,
|
||||
ownerKind, ownerID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
var digests []string
|
||||
for rows.Next() {
|
||||
var d string
|
||||
if err := rows.Scan(&d); err == nil {
|
||||
digests = append(digests, d)
|
||||
}
|
||||
}
|
||||
rows.Close()
|
||||
if err := rows.Err(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if len(digests) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
tx, err := s.db.Begin()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
if _, err := tx.Exec(`DELETE FROM media_refs WHERE owner_kind = ? AND owner_id = ?`, ownerKind, ownerID); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
for _, d := range digests {
|
||||
if _, err := tx.Exec(`UPDATE media SET ref_count = MAX(0, ref_count - 1) WHERE digest = ?`, d); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return len(digests), nil
|
||||
}
|
||||
|
||||
// Refs 返回某个 owner 引用的全部 digest。
|
||||
func (s *Store) Refs(ownerKind, ownerID string) ([]string, error) {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
rows, err := s.db.Query(`SELECT digest FROM media_refs WHERE owner_kind = ? AND owner_id = ? ORDER BY created_at`,
|
||||
ownerKind, ownerID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []string
|
||||
for rows.Next() {
|
||||
var d string
|
||||
if err := rows.Scan(&d); err == nil {
|
||||
out = append(out, d)
|
||||
}
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// Search 按描述文本做 LIKE 匹配，返回最近的若干条。
|
||||
//
|
||||
// 刻意不在这里做向量检索：媒体的语义检索走 L2 文档层的既有索引
|
||||
// （描述文字随记忆条目一起进 Doc.Content，复用那套 TF-IDF/embedding），
|
||||
// 本方法只是"按关键词直接翻媒体库"的补充入口。
|
||||
func (s *Store) Search(query string, kind Kind, limit int) ([]*Item, error) {
|
||||
if limit <= 0 {
|
||||
limit = 20
|
||||
}
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
q := `SELECT digest, kind, mime, size, width, height, origin_path, tool,
|
||||
description, described_by, ref_count, first_seen, last_seen
|
||||
FROM media WHERE COALESCE(description,'') != ''`
|
||||
args := []interface{}{}
|
||||
if strings.TrimSpace(query) != "" {
|
||||
q += ` AND description LIKE ?`
|
||||
args = append(args, "%"+query+"%")
|
||||
}
|
||||
if kind != "" {
|
||||
q += ` AND kind = ?`
|
||||
args = append(args, string(kind))
|
||||
}
|
||||
q += ` ORDER BY last_seen DESC LIMIT ?`
|
||||
args = append(args, limit)
|
||||
|
||||
rows, err := s.db.Query(q, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []*Item
|
||||
for rows.Next() {
|
||||
it, err := s.scanRows(rows)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
out = append(out, it)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// Pending 返回尚无描述的媒体，供后台描述任务消费。
|
||||
func (s *Store) Pending(limit int) ([]*Item, error) {
|
||||
if limit <= 0 {
|
||||
limit = 10
|
||||
}
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
rows, err := s.db.Query(`
|
||||
SELECT digest, kind, mime, size, width, height, origin_path, tool,
|
||||
description, described_by, ref_count, first_seen, last_seen
|
||||
FROM media WHERE COALESCE(description,'') = ''
|
||||
ORDER BY last_seen DESC LIMIT ?`, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []*Item
|
||||
for rows.Next() {
|
||||
it, err := s.scanRows(rows)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
out = append(out, it)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// GC 清理无人引用的内容。
|
||||
//
|
||||
// 两段策略：
|
||||
// 1. ref_count=0 且 last_seen 早于 minAge 的一律清理。刚 Put 还没来得及
|
||||
// AddRef 的项 refcount 也是 0，minAge 保护它们不被立刻清掉。
|
||||
// 2. 清完仍超 maxBytes 时，继续按 last_seen 从旧到新淘汰 ref_count=0 的项。
|
||||
//
|
||||
// 有引用的项永不删除——那会让记忆里的 digest 变成悬空指针，正是本包要避免的。
|
||||
func (s *Store) GC(minAge time.Duration) (removed int, freed int64, err error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
cutoff := time.Now().Add(-minAge)
|
||||
rows, err := s.db.Query(`
|
||||
SELECT digest, size FROM media
|
||||
WHERE ref_count <= 0 AND last_seen < ?
|
||||
ORDER BY last_seen`, cutoff)
|
||||
if err != nil {
|
||||
return 0, 0, err
|
||||
}
|
||||
type cand struct {
|
||||
digest string
|
||||
size int64
|
||||
}
|
||||
var cands []cand
|
||||
for rows.Next() {
|
||||
var c cand
|
||||
if err := rows.Scan(&c.digest, &c.size); err == nil {
|
||||
cands = append(cands, c)
|
||||
}
|
||||
}
|
||||
rows.Close()
|
||||
|
||||
for _, c := range cands {
|
||||
if e := os.Remove(s.blobPath(c.digest)); e != nil && !os.IsNotExist(e) {
|
||||
continue // 删不掉就留着元数据，下轮再试；不制造"元数据没了文件还在"的孤儿
|
||||
}
|
||||
if _, e := s.db.Exec(`DELETE FROM media WHERE digest = ?`, c.digest); e != nil {
|
||||
continue
|
||||
}
|
||||
removed++
|
||||
freed += c.size
|
||||
}
|
||||
|
||||
if s.maxBytes > 0 {
|
||||
r2, f2 := s.enforceCapacityLocked()
|
||||
removed += r2
|
||||
freed += f2
|
||||
}
|
||||
return removed, freed, nil
|
||||
}
|
||||
|
||||
// enforceCapacityLocked 在超出 maxBytes 时继续淘汰无引用项（调用方已持锁）。
|
||||
func (s *Store) enforceCapacityLocked() (removed int, freed int64) {
|
||||
var total int64
|
||||
if err := s.db.QueryRow(`SELECT COALESCE(SUM(size), 0) FROM media`).Scan(&total); err != nil {
|
||||
return 0, 0
|
||||
}
|
||||
if total <= s.maxBytes {
|
||||
return 0, 0
|
||||
}
|
||||
need := total - s.maxBytes
|
||||
|
||||
rows, err := s.db.Query(`SELECT digest, size FROM media WHERE ref_count <= 0 ORDER BY last_seen`)
|
||||
if err != nil {
|
||||
return 0, 0
|
||||
}
|
||||
type cand struct {
|
||||
digest string
|
||||
size int64
|
||||
}
|
||||
var cands []cand
|
||||
for rows.Next() {
|
||||
var c cand
|
||||
if err := rows.Scan(&c.digest, &c.size); err == nil {
|
||||
cands = append(cands, c)
|
||||
}
|
||||
}
|
||||
rows.Close()
|
||||
|
||||
for _, c := range cands {
|
||||
if freed >= need {
|
||||
break
|
||||
}
|
||||
if e := os.Remove(s.blobPath(c.digest)); e != nil && !os.IsNotExist(e) {
|
||||
continue
|
||||
}
|
||||
if _, e := s.db.Exec(`DELETE FROM media WHERE digest = ?`, c.digest); e != nil {
|
||||
continue
|
||||
}
|
||||
removed++
|
||||
freed += c.size
|
||||
}
|
||||
return removed, freed
|
||||
}
|
||||
|
||||
// Stats 返回容量与条目统计，供 WebUI / healthcheck 展示。
|
||||
func (s *Store) Stats() map[string]interface{} {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
out := map[string]interface{}{"blob_dir": s.blobDir, "max_bytes": s.maxBytes}
|
||||
var count, described, orphan int
|
||||
var total int64
|
||||
s.db.QueryRow(`SELECT COUNT(*), COALESCE(SUM(size),0) FROM media`).Scan(&count, &total)
|
||||
s.db.QueryRow(`SELECT COUNT(*) FROM media WHERE COALESCE(description,'') != ''`).Scan(&described)
|
||||
s.db.QueryRow(`SELECT COUNT(*) FROM media WHERE ref_count <= 0`).Scan(&orphan)
|
||||
out["count"] = count
|
||||
out["total_bytes"] = total
|
||||
out["described"] = described
|
||||
out["unreferenced"] = orphan
|
||||
|
||||
byKind := map[string]int{}
|
||||
rows, err := s.db.Query(`SELECT kind, COUNT(*) FROM media GROUP BY kind`)
|
||||
if err == nil {
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var k string
|
||||
var n int
|
||||
if rows.Scan(&k, &n) == nil {
|
||||
byKind[k] = n
|
||||
}
|
||||
}
|
||||
}
|
||||
out["by_kind"] = byKind
|
||||
return out
|
||||
}
|
||||
|
||||
func (s *Store) Close() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.db.Close()
|
||||
}
|
||||
|
||||
// ---- 扫描辅助 ----
|
||||
|
||||
type rowScanner interface {
|
||||
Scan(dest ...interface{}) error
|
||||
}
|
||||
|
||||
func (s *Store) scanOne(r rowScanner) (*Item, error) {
|
||||
it, err := scanItem(r)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, fmt.Errorf("media: unknown digest")
|
||||
}
|
||||
return it, err
|
||||
}
|
||||
|
||||
func (s *Store) scanRows(r rowScanner) (*Item, error) { return scanItem(r) }
|
||||
|
||||
func scanItem(r rowScanner) (*Item, error) {
|
||||
var it Item
|
||||
var kind string
|
||||
var origin, tool, desc, by sql.NullString
|
||||
if err := r.Scan(&it.Digest, &kind, &it.MIME, &it.Size, &it.Width, &it.Height,
|
||||
&origin, &tool, &desc, &by, &it.RefCount, &it.FirstSeen, &it.LastSeen); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
it.Kind = Kind(kind)
|
||||
it.OriginPath = origin.String
|
||||
it.Tool = tool.String
|
||||
it.Description = desc.String
|
||||
it.DescribedBy = by.String
|
||||
return &it, nil
|
||||
}
|
||||
|
||||
// ---- 工具函数 ----
|
||||
|
||||
// KindFromMIME 把 MIME 归到大类。
|
||||
func KindFromMIME(mime string) Kind {
|
||||
m := strings.ToLower(strings.TrimSpace(mime))
|
||||
switch {
|
||||
case strings.HasPrefix(m, "image/"):
|
||||
return KindImage
|
||||
case strings.HasPrefix(m, "audio/"):
|
||||
return KindAudio
|
||||
case strings.HasPrefix(m, "video/"):
|
||||
return KindVideo
|
||||
default:
|
||||
return KindOther
|
||||
}
|
||||
}
|
||||
|
||||
// ParseDataURL 从 data:<mime>;base64,<data> 提取 MIME 与原始字节。
|
||||
//
|
||||
// 与 agent/api 里的 parseAudioDataURL 分开实现：那个只认音频且只回 base64
|
||||
// 串（它要把串塞回 OpenAI 的 input_audio 字段），这里要的是解码后的字节。
|
||||
func ParseDataURL(url string) (mime string, data []byte, ok bool) {
|
||||
const prefix = "data:"
|
||||
if !strings.HasPrefix(url, prefix) {
|
||||
return "", nil, false
|
||||
}
|
||||
rest := url[len(prefix):]
|
||||
comma := strings.IndexByte(rest, ',')
|
||||
if comma < 0 {
|
||||
return "", nil, false
|
||||
}
|
||||
head := rest[:comma]
|
||||
payload := rest[comma+1:]
|
||||
if !strings.HasSuffix(strings.ToLower(head), ";base64") {
|
||||
return "", nil, false
|
||||
}
|
||||
mime = head[:len(head)-len(";base64")]
|
||||
if mime == "" || payload == "" {
|
||||
return "", nil, false
|
||||
}
|
||||
decoded, err := base64Decode(payload)
|
||||
if err != nil {
|
||||
return "", nil, false
|
||||
}
|
||||
return mime, decoded, true
|
||||
}
|
||||
|
||||
// DataURL 把内容编回 data URL，供重新注入模型对话。
|
||||
func DataURL(mime string, data []byte) string {
|
||||
return "data:" + mime + ";base64," + base64Encode(data)
|
||||
}
|
||||
|
||||
// CopyFrom 从 reader 读全部内容后 Put，用于大文件不便一次性构造 []byte 的场合。
|
||||
func (s *Store) CopyFrom(r io.Reader, meta Item) (string, error) {
|
||||
data, err := io.ReadAll(r)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("media: read source: %w", err)
|
||||
}
|
||||
return s.Put(data, meta)
|
||||
}
|
||||
|
||||
// MarshalItems 序列化条目列表，供工具返回给模型。
|
||||
func MarshalItems(items []*Item) string {
|
||||
b, err := json.Marshal(items)
|
||||
if err != nil {
|
||||
return "[]"
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func shortDigest(d string) string {
|
||||
if len(d) > 12 {
|
||||
return d[:12]
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
func truncate(s string, n int) string {
|
||||
s = strings.Join(strings.Fields(s), " ")
|
||||
if len(s) <= n {
|
||||
return s
|
||||
}
|
||||
return s[:n] + "..."
|
||||
}
|
||||
477
internal/memory/media/media_test.go
Normal file
477
internal/memory/media/media_test.go
Normal file
@ -0,0 +1,477 @@
|
||||
package media
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func newTestStore(t *testing.T, maxBytes int64) *Store {
|
||||
t.Helper()
|
||||
s, err := New(t.TempDir(), maxBytes)
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { s.Close() })
|
||||
return s
|
||||
}
|
||||
|
||||
func TestPut_ContentAddressedDedup(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
data := []byte("fake-png-bytes")
|
||||
|
||||
d1, err := s.Put(data, Item{MIME: "image/png", OriginPath: "/tmp/a.png"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
d2, err := s.Put(data, Item{MIME: "image/png", OriginPath: "/tmp/b.png"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if d1 != d2 {
|
||||
t.Fatalf("同一内容应得同一 digest：%s vs %s", d1, d2)
|
||||
}
|
||||
|
||||
// 去重的意义：同一张图反复注入（连问几轮同一截图、see_video 相邻帧）
|
||||
// 只占一份磁盘。
|
||||
st := s.Stats()
|
||||
if st["count"].(int) != 1 {
|
||||
t.Fatalf("同一内容应只有 1 条记录，实际 %v", st["count"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestPut_KindInferredFromMIME(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
cases := map[string]Kind{
|
||||
"image/png": KindImage,
|
||||
"image/jpeg": KindImage,
|
||||
"audio/wav": KindAudio,
|
||||
"video/mp4": KindVideo,
|
||||
"text/plain": KindOther,
|
||||
}
|
||||
for mime, want := range cases {
|
||||
d, err := s.Put([]byte("payload-"+mime), Item{MIME: mime})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
it, err := s.Stat(d)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if it.Kind != want {
|
||||
t.Fatalf("%s 应归为 %s，实际 %s", mime, want, it.Kind)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestGet_DetectsCorruption(t *testing.T) {
|
||||
// CAS 的全部保证建立在「文件名 == 内容摘要」上。外部误改或位翻转必须
|
||||
// 被发现——把损坏的图喂给模型只会得到无从追溯的幻觉。
|
||||
s := newTestStore(t, 0)
|
||||
d, err := s.Put([]byte("original-content"), Item{MIME: "image/png"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s.Get(d); err != nil {
|
||||
t.Fatalf("正常读取应成功: %v", err)
|
||||
}
|
||||
|
||||
if err := os.WriteFile(s.blobPath(d), []byte("tampered!"), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = s.Get(d)
|
||||
if err == nil {
|
||||
t.Fatal("内容被改后应报 digest 不匹配，却读成功了")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "digest mismatch") {
|
||||
t.Fatalf("错误应指明 digest 不匹配，得到: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPut_NoPartialBlobOnDisk(t *testing.T) {
|
||||
// 先写 .tmp 再 rename：确认落地后目录里不留临时文件。
|
||||
s := newTestStore(t, 0)
|
||||
d, err := s.Put([]byte("some-bytes"), Item{MIME: "image/png"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
dir := filepath.Dir(s.blobPath(d))
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, e := range entries {
|
||||
if strings.HasSuffix(e.Name(), ".tmp") {
|
||||
t.Fatalf("落地后不该留临时文件: %s", e.Name())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefCount_AddIsIdempotent(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
d, _ := s.Put([]byte("img"), Item{MIME: "image/png"})
|
||||
|
||||
for i := 0; i < 3; i++ {
|
||||
if err := s.AddRef(d, "context", "evt-1"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
it, _ := s.Stat(d)
|
||||
// 重复 AddRef 若都递增，计数会虚高，GC 永远不敢清。
|
||||
if it.RefCount != 1 {
|
||||
t.Fatalf("同一 owner 重复 AddRef 应只计 1，实际 %d", it.RefCount)
|
||||
}
|
||||
|
||||
if err := s.AddRef(d, "document", "doc-9"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
it, _ = s.Stat(d)
|
||||
if it.RefCount != 2 {
|
||||
t.Fatalf("不同 owner 应各计一次，实际 %d", it.RefCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefCount_DropAndNeverNegative(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
d, _ := s.Put([]byte("img"), Item{MIME: "image/png"})
|
||||
s.AddRef(d, "context", "e1")
|
||||
|
||||
if err := s.DropRef(d, "context", "e1"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
it, _ := s.Stat(d)
|
||||
if it.RefCount != 0 {
|
||||
t.Fatalf("应归零，实际 %d", it.RefCount)
|
||||
}
|
||||
|
||||
// 多余的 DropRef 不该把计数压成负数（负数会让容量 GC 的排序失去意义）
|
||||
for i := 0; i < 3; i++ {
|
||||
s.DropRef(d, "context", "e1")
|
||||
}
|
||||
it, _ = s.Stat(d)
|
||||
if it.RefCount != 0 {
|
||||
t.Fatalf("重复 DropRef 后仍应为 0，实际 %d", it.RefCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDropOwner_RemovesAllItsRefs(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
d1, _ := s.Put([]byte("frame1"), Item{MIME: "image/jpeg"})
|
||||
d2, _ := s.Put([]byte("frame2"), Item{MIME: "image/jpeg"})
|
||||
s.AddRef(d1, "context", "evt-x")
|
||||
s.AddRef(d2, "context", "evt-x")
|
||||
s.AddRef(d1, "document", "doc-y") // 别的 owner 也引了 d1
|
||||
|
||||
n, err := s.DropOwner("context", "evt-x")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n != 2 {
|
||||
t.Fatalf("应注销 2 条引用，实际 %d", n)
|
||||
}
|
||||
|
||||
it1, _ := s.Stat(d1)
|
||||
it2, _ := s.Stat(d2)
|
||||
if it1.RefCount != 1 {
|
||||
t.Fatalf("d1 仍被 document 引用，应剩 1，实际 %d", it1.RefCount)
|
||||
}
|
||||
if it2.RefCount != 0 {
|
||||
t.Fatalf("d2 应归零，实际 %d", it2.RefCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefs_ListsOwnerDigests(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
d1, _ := s.Put([]byte("a"), Item{MIME: "image/png"})
|
||||
d2, _ := s.Put([]byte("b"), Item{MIME: "image/png"})
|
||||
s.AddRef(d1, "context", "e1")
|
||||
s.AddRef(d2, "context", "e1")
|
||||
|
||||
got, err := s.Refs("context", "e1")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(got) != 2 {
|
||||
t.Fatalf("应返回 2 个 digest，实际 %d", len(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestGC_KeepsReferencedContent(t *testing.T) {
|
||||
// 有引用的项永不删除——那会让记忆里的 digest 变成悬空指针，
|
||||
// 正是本包要避免的。
|
||||
s := newTestStore(t, 0)
|
||||
kept, _ := s.Put([]byte("referenced"), Item{MIME: "image/png"})
|
||||
orphan, _ := s.Put([]byte("orphaned"), Item{MIME: "image/png"})
|
||||
s.AddRef(kept, "context", "e1")
|
||||
|
||||
// minAge=0 让刚 Put 的都算超龄
|
||||
removed, _, err := s.GC(0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if removed != 1 {
|
||||
t.Fatalf("应只清 1 条无引用项，实际 %d", removed)
|
||||
}
|
||||
if _, err := s.Get(kept); err != nil {
|
||||
t.Fatalf("被引用的内容不该被清: %v", err)
|
||||
}
|
||||
if _, err := s.Stat(orphan); err == nil {
|
||||
t.Fatal("无引用项的元数据应已删除")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGC_MinAgeProtectsFreshUnreferenced(t *testing.T) {
|
||||
// 刚 Put 还没来得及 AddRef 的项 refcount 也是 0；
|
||||
// minAge 必须保护它们，否则「Put 完还没挂上就被 GC 清掉」。
|
||||
s := newTestStore(t, 0)
|
||||
d, _ := s.Put([]byte("just-arrived"), Item{MIME: "image/png"})
|
||||
|
||||
removed, _, err := s.GC(time.Hour)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if removed != 0 {
|
||||
t.Fatalf("新入库项应被 minAge 保护，却清掉了 %d 条", removed)
|
||||
}
|
||||
if _, err := s.Get(d); err != nil {
|
||||
t.Fatalf("内容应还在: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGC_EnforcesCapacity(t *testing.T) {
|
||||
// 容量上限：清完超龄项后仍超限，继续按 last_seen 从旧到新淘汰无引用项。
|
||||
blob := make([]byte, 1024)
|
||||
s := newTestStore(t, 2048) // 只容 2KB
|
||||
|
||||
var digests []string
|
||||
for i := 0; i < 4; i++ {
|
||||
b := append([]byte{byte(i)}, blob...) // 内容各异，避免去重
|
||||
d, err := s.Put(b, Item{MIME: "image/png"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
digests = append(digests, d)
|
||||
time.Sleep(2 * time.Millisecond) // 拉开 last_seen
|
||||
}
|
||||
|
||||
// 保护最后一个，确认容量 GC 也不碰有引用的
|
||||
s.AddRef(digests[3], "context", "e1")
|
||||
|
||||
removed, freed, err := s.GC(0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if removed == 0 {
|
||||
t.Fatal("超限应触发淘汰")
|
||||
}
|
||||
if _, err := s.Get(digests[3]); err != nil {
|
||||
t.Fatalf("有引用项即使超限也不该删: %v", err)
|
||||
}
|
||||
t.Logf("removed=%d freed=%d", removed, freed)
|
||||
|
||||
st := s.Stats()
|
||||
if total := st["total_bytes"].(int64); total > 2048 {
|
||||
// 有引用项可能让总量降不到线下，这是刻意的（宁可超限也不断引用）
|
||||
t.Logf("总量 %d 仍超 2048，因有引用项不可删（预期行为）", total)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDescribe_OverwritesExplicitly(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
d, _ := s.Put([]byte("img"), Item{MIME: "image/png"})
|
||||
|
||||
if err := s.Describe(d, "一只橘猫", "vis-a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
it, _ := s.Stat(d)
|
||||
if it.Description != "一只橘猫" || it.DescribedBy != "vis-a" {
|
||||
t.Fatalf("描述未写入: %+v", it)
|
||||
}
|
||||
|
||||
// Describe 是显式操作，允许覆盖（换更强模型重描述）
|
||||
if err := s.Describe(d, "一只橘色虎斑猫坐在窗台", "vis-b"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
it, _ = s.Stat(d)
|
||||
if !strings.Contains(it.Description, "虎斑") || it.DescribedBy != "vis-b" {
|
||||
t.Fatalf("Describe 应覆盖旧描述: %+v", it)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDescribe_UnknownDigestErrors(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
err := s.Describe("deadbeef", "x", "y")
|
||||
if err == nil {
|
||||
t.Fatal("未知 digest 应报错而非静默成功")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPut_DoesNotClobberExistingDescription(t *testing.T) {
|
||||
// 先到的描述可能来自更强的模型；后到的空值不该把它冲掉。
|
||||
s := newTestStore(t, 0)
|
||||
data := []byte("img")
|
||||
d, _ := s.Put(data, Item{MIME: "image/png", Description: "详细描述", DescribedBy: "strong-model"})
|
||||
|
||||
// 第二次 Put 同内容但不带描述
|
||||
if _, err := s.Put(data, Item{MIME: "image/png"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
it, _ := s.Stat(d)
|
||||
if it.Description != "详细描述" || it.DescribedBy != "strong-model" {
|
||||
t.Fatalf("重复 Put 的空描述不该冲掉已有描述: %+v", it)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPut_BackfillsMissingDimensions(t *testing.T) {
|
||||
// 后来者可能带着前一次没有的信息（尺寸、工具名）
|
||||
s := newTestStore(t, 0)
|
||||
data := []byte("img")
|
||||
d, _ := s.Put(data, Item{MIME: "image/png"})
|
||||
|
||||
if _, err := s.Put(data, Item{MIME: "image/png", Width: 640, Height: 480, Tool: "multimodal_see_picture"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
it, _ := s.Stat(d)
|
||||
if it.Width != 640 || it.Height != 480 {
|
||||
t.Fatalf("尺寸应被补写: %dx%d", it.Width, it.Height)
|
||||
}
|
||||
if it.Tool != "multimodal_see_picture" {
|
||||
t.Fatalf("工具名应被补写: %q", it.Tool)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearch_FiltersByDescriptionAndKind(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
di, _ := s.Put([]byte("chart-img"), Item{MIME: "image/png"})
|
||||
da, _ := s.Put([]byte("speech-aud"), Item{MIME: "audio/wav"})
|
||||
dn, _ := s.Put([]byte("no-desc"), Item{MIME: "image/png"})
|
||||
s.Describe(di, "一张蓝色的柱状图表", "vis")
|
||||
s.Describe(da, "一段关于图表的讲解录音", "aud")
|
||||
|
||||
all, err := s.Search("图表", "", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(all) != 2 {
|
||||
t.Fatalf("两条描述都含「图表」，应返回 2，实际 %d", len(all))
|
||||
}
|
||||
|
||||
imgs, _ := s.Search("图表", KindImage, 10)
|
||||
if len(imgs) != 1 || imgs[0].Digest != di {
|
||||
t.Fatalf("按 image 过滤应只剩图片，实际 %d 条", len(imgs))
|
||||
}
|
||||
|
||||
// 无描述的项不该出现在语义检索结果里
|
||||
for _, it := range all {
|
||||
if it.Digest == dn {
|
||||
t.Fatal("无描述的项不该被 Search 返回")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPending_ReturnsUndescribed(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
described, _ := s.Put([]byte("has-desc"), Item{MIME: "image/png"})
|
||||
undescribed, _ := s.Put([]byte("needs-desc"), Item{MIME: "image/png"})
|
||||
s.Describe(described, "已有描述", "vis")
|
||||
|
||||
pending, err := s.Pending(10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(pending) != 1 || pending[0].Digest != undescribed {
|
||||
t.Fatalf("应只返回未描述项，实际 %d 条", len(pending))
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseDataURL(t *testing.T) {
|
||||
// 与 agent/api 的 parseAudioDataURL 不同：这里要解码后的字节。
|
||||
raw := []byte{0x89, 'P', 'N', 'G'}
|
||||
url := DataURL("image/png", raw)
|
||||
|
||||
mime, data, ok := ParseDataURL(url)
|
||||
if !ok {
|
||||
t.Fatal("应解析成功")
|
||||
}
|
||||
if mime != "image/png" {
|
||||
t.Fatalf("MIME 应为 image/png，得到 %q", mime)
|
||||
}
|
||||
if string(data) != string(raw) {
|
||||
t.Fatalf("字节应还原，得到 %v", data)
|
||||
}
|
||||
|
||||
for _, bad := range []string{
|
||||
"http://example.com/x.png", // 不是 data URL
|
||||
"data:image/png,notbase64", // 缺 ;base64
|
||||
"data:;base64,", // 空 MIME 与空载荷
|
||||
"data:image/png;base64,!!!", // 非法 base64
|
||||
} {
|
||||
if _, _, ok := ParseDataURL(bad); ok {
|
||||
t.Fatalf("%q 应解析失败", bad)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestStats_CountsByKindAndDescription(t *testing.T) {
|
||||
s := newTestStore(t, 4096)
|
||||
d1, _ := s.Put([]byte("i1"), Item{MIME: "image/png"})
|
||||
s.Put([]byte("i2"), Item{MIME: "image/jpeg"})
|
||||
s.Put([]byte("a1"), Item{MIME: "audio/wav"})
|
||||
s.Describe(d1, "描述", "vis")
|
||||
s.AddRef(d1, "context", "e1")
|
||||
|
||||
st := s.Stats()
|
||||
if st["count"].(int) != 3 {
|
||||
t.Fatalf("count 应为 3，实际 %v", st["count"])
|
||||
}
|
||||
if st["described"].(int) != 1 {
|
||||
t.Fatalf("described 应为 1，实际 %v", st["described"])
|
||||
}
|
||||
if st["unreferenced"].(int) != 2 {
|
||||
t.Fatalf("unreferenced 应为 2，实际 %v", st["unreferenced"])
|
||||
}
|
||||
byKind := st["by_kind"].(map[string]int)
|
||||
if byKind["image"] != 2 || byKind["audio"] != 1 {
|
||||
t.Fatalf("by_kind 不对: %v", byKind)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPut_RejectsEmpty(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
if _, err := s.Put(nil, Item{MIME: "image/png"}); err == nil {
|
||||
t.Fatal("空内容应报错")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReopen_PersistsAcrossRestart(t *testing.T) {
|
||||
// 记忆的意义就在于跨重启还在。
|
||||
dir := t.TempDir()
|
||||
s1, err := New(dir, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
d, _ := s1.Put([]byte("persistent-img"), Item{MIME: "image/png", OriginPath: "/tmp/x.png"})
|
||||
s1.Describe(d, "跨重启的描述", "vis")
|
||||
s1.AddRef(d, "context", "e1")
|
||||
s1.Close()
|
||||
|
||||
s2, err := New(dir, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s2.Close()
|
||||
|
||||
it, err := s2.Stat(d)
|
||||
if err != nil {
|
||||
t.Fatalf("重开后应能查到: %v", err)
|
||||
}
|
||||
if it.Description != "跨重启的描述" || it.RefCount != 1 {
|
||||
t.Fatalf("元数据应持久化: %+v", it)
|
||||
}
|
||||
data, err := s2.Get(d)
|
||||
if err != nil || string(data) != "persistent-img" {
|
||||
t.Fatalf("内容应持久化: %v / %q", err, data)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user