mirror of
https://gitcode.com/JianFeeeee/HomeAgent.git
synced 2026-10-15 12:15:16 +00:00
refactor(memory)!: 删除 CAS 包与第二套召回（块 union 化第四阶段：收尾）
块 union 化前三阶段已让块自带负载、回填存量字节。本提交拆掉 CAS 本身： `internal/memory/media/`（2100 行）+ `QueryMediaScored`（第二套召回） + `media_refs`/`ref_count` 引用账本 + `internal/sdk/media.go`（MediaAPI） + `SDKConfig.Media` / `SetMediaStore` 接线。删完 CAS 才真正不存在。 ## 为何可删（不是"先删了再说"） - **MediaAPI 在生产早已是死的**：`SDKConfig.Media` 自86572d13起从未被设置， `sdk.Media()` 恒为 nil，插件 digest 路径一直走"内容寻址库不可用"分支。 删的是死代码，不是活路径。 - **第二套召回确无调用方**：`QueryMediaScored` 仅被 media 包自身的测试引用， 跨包引用为 0（块召回走 `RecallBlocks`）。 - **heldMediaDigests 补丁**（crossmodal 靠它防召回"无处可归的媒体"） 已在86572d13删除 —— 两套存储才需要的那种补丁，随存储一起没了。 ## 替代路径：块引用取代 digest 插件取媒体从 `sdk.Media().Get(digest)` 改为 `sdk.Memory().MediaBytes(ref)`： - **新增 `MemoryAPI.MediaBytes(ref) (data, mime, ok, err)`**（内置接口，非 pubsdk）： 按块引用（完整块 ID 或唯一前缀）取字节；`ok=false` 而非 error —— "这个引用没有媒体"是调用方能处理的常态。 - **seemedia / multimodal**：`digest` 参数语义改为**媒体块引用**，工具描述同步。 path/url 两条来源不变（生产实际只用这两条）。 - **webui 知识上传**：媒体字节不再 Put 进 CAS 再拿 digest，而是作为 `KnowledgeMediaRef.Data` 随引用携带，由 knowledge.Write 落进条目目录 （与 content.md 同生共灭，Remove 的 RemoveAll 天然清掉）。 - **回填命令去依赖**：cmd/homed-media-backfill 不再 import 待删的 media 包， 改为直接读**冻结的历史布局** blobs/<d[:2]>/<d[2:]> 并自校验 sha256。 依赖待删代码的迁移工具会在删包那刻失效 —— 迁移必须比被迁移的东西活得久。 ## 判据（全部做过变异测试） - `TestMemoryBytes_ByBlockRef`：完整 ID / 唯一前缀可取回且逐字节一致； 不存在→ok=false；过短前缀（歧义）→ok=false；无负载块→ok=false；空引用不 panic。 变异（不校验 ok，把不存在当有）当场被 3 条断言抓住。 - `TestE2EMediaBytesAgainstBackfilledDB`（-tags e2e，读真回填库）： 10/10 块经 MediaBytes 取回的字节与块 data 的 sha256 完全一致。 - `TestKnowledgeUploadedMediaRetrievable` 重写为**条目目录自带字节**判据 （不再是"向存储按 digest 读回"）；`TestKnowledgeUploadMediaNeedsNoStore` 改为"上传无外部存储依赖也必须成功"（旧的 503 语义已随 CAS 消失）。 ## 静态确认 `QueryMediaScored` / `internal/memory/media` import / `MediaAPI` / `sdk.Media()` / `SetMediaStore` / `heldMediaDigests` 全仓均为 0 处 （仅剩 medialoop.go 一句历史注释）。 回归：memory / sdk / agent-core / knowledge / plugin / webui / seemedia / multimodal 八包全绿；全仓生产编译 GREEN。
This commit is contained in:
JianFeeeee
@ -1,14 +0,0 @@
|
||||
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)
|
||||
}
|
||||
@ -1,662 +0,0 @@
|
||||
// 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"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
_ "github.com/mattn/go-sqlite3"
|
||||
)
|
||||
|
||||
// digestHexLen 是 sha256 的十六进制串长度。
|
||||
const digestHexLen = sha256.Size * 2
|
||||
|
||||
// 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"`
|
||||
// FirstSeen/LastSeen 是首末次入库时间。
|
||||
FirstSeen time.Time `json:"first_seen"`
|
||||
LastSeen time.Time `json:"last_seen"`
|
||||
// --- 多模态嵌入（v1.2.0）---
|
||||
// Vec 是视觉嵌入向量的序列化（JSON []float64），nil 表示未嵌入。
|
||||
Vec []float64 `json:"vec,omitempty"`
|
||||
// VecModel 是产生 Vec 的模型标识（如 "clip-vit-b32"），
|
||||
// 用于模型切换后判断是否需要重算。
|
||||
VecModel string `json:"vec_model,omitempty"`
|
||||
}
|
||||
|
||||
// Store 管理媒体的元数据（SQLite）与内容（磁盘 CAS 目录）。
|
||||
//
|
||||
// 元数据与内容分离而不是把 blob 塞进 SQLite：单张图动辄几 MB，塞进库会让
|
||||
// 每次 VACUUM/备份都拖着几百 MB 走，也让 WAL 迅速膨胀。CAS 目录用两级
|
||||
// 前缀分桶（ab/cdef...）避免单目录几万文件。
|
||||
type Store struct {
|
||||
mu sync.RWMutex
|
||||
db *sql.DB
|
||||
blobDir string
|
||||
}
|
||||
|
||||
// New 打开（或初始化）媒体存储。
|
||||
// dir 下会建 media.db 与 blobs/ 两个条目。
|
||||
func New(dir string) (*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")}
|
||||
if err := s.initSchema(); err != nil {
|
||||
db.Close()
|
||||
return nil, err
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func (s *Store) initSchema() error {
|
||||
stmts := []string{
|
||||
// digest 作主键：内容即身份，重复 Put 同一内容不重复落盘。
|
||||
`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,
|
||||
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_last_seen ON media(last_seen)`,
|
||||
}
|
||||
for _, q := range stmts {
|
||||
if _, err := s.db.Exec(q); err != nil {
|
||||
return fmt.Errorf("media: schema %q: %w", truncate(q, 60), err)
|
||||
}
|
||||
}
|
||||
// v1.2.0 迁移：给 media 表加 vec（视觉嵌入向量 JSON）和 vec_model（模型标识）。
|
||||
migrations := []string{
|
||||
`ALTER TABLE media ADD COLUMN vec TEXT`,
|
||||
`ALTER TABLE media ADD COLUMN vec_model TEXT`,
|
||||
}
|
||||
for _, q := range migrations {
|
||||
_, _ = s.db.Exec(q) // 列已存在时返回 "duplicate column name"，可忽略
|
||||
}
|
||||
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, first_seen, last_seen)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT(digest) DO UPDATE SET
|
||||
last_seen = excluded.last_seen,
|
||||
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, 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,
|
||||
first_seen, last_seen,
|
||||
vec, vec_model
|
||||
FROM media WHERE digest = ?`, digest))
|
||||
}
|
||||
|
||||
// Delete 删除媒体内容与元数据。
|
||||
//
|
||||
// 这不是 GC，也不看引用计数：调用方是记忆系统本身——当它把一个记忆块
|
||||
// 永久地从三层记忆中删掉（而非在层间迁移）时，媒体作为块的内容一并删除。
|
||||
// 文本块就是这么管理的：删除块即删除内容。
|
||||
func (s *Store) Delete(digest string) error {
|
||||
if digest == "" {
|
||||
return nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if err := os.Remove(s.blobPath(digest)); err != nil && !os.IsNotExist(err) {
|
||||
return fmt.Errorf("media: remove blob %s: %w", shortDigest(digest), err)
|
||||
}
|
||||
if _, err := s.db.Exec(`DELETE FROM media WHERE digest = ?`, digest); err != nil {
|
||||
return fmt.Errorf("media: delete meta %s: %w", shortDigest(digest), err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// 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}
|
||||
var count int
|
||||
var total int64
|
||||
s.db.QueryRow(`SELECT COUNT(*), COALESCE(SUM(size),0) FROM media`).Scan(&count, &total)
|
||||
out["count"] = count
|
||||
out["total_bytes"] = total
|
||||
|
||||
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()
|
||||
}
|
||||
|
||||
// ---- 多模态嵌入（v1.2.0） ----
|
||||
|
||||
// SetVec 给一条已入库的媒体设置视觉嵌入向量。
|
||||
//
|
||||
// 设计选择：vec 是 TEXT（JSON 序列化的 []float64）而非 BLOB，
|
||||
// 因为 Go 的 json.Marshal/Unmarshal 对 []float64 是自然的，
|
||||
// 而 SQLite 的 BLOB 是 []byte，序列化多一层反而复杂。
|
||||
// 量级：一条 vec 最多 1536 维 × ~15 字节 ≈ 23KB，TEXT 合适。
|
||||
func (s *Store) SetVec(digest string, vec []float64, model string) error {
|
||||
vecJSON, err := json.Marshal(vec)
|
||||
if err != nil {
|
||||
return fmt.Errorf("media: marshal vec: %w", err)
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
_, err = s.db.Exec(`UPDATE media SET vec = ?, vec_model = ? WHERE digest = ?`,
|
||||
string(vecJSON), model, digest)
|
||||
return err
|
||||
}
|
||||
|
||||
// StaleVecDigests 返回所有需要重新嵌入的图片 digest：
|
||||
// vec_model 不等于 currentModel（模型切换）或 vec_model 为空（从未嵌入）。
|
||||
// 调用方使用返回的 digest 列表调用 Get/EmbedImage/SetVec 完成重算。
|
||||
func (s *Store) StaleVecDigests(currentModel string) ([]string, error) {
|
||||
return s.staleVecDigests(currentModel, "image")
|
||||
}
|
||||
|
||||
// StaleVecDigestsAll 返回所有需要重新嵌入的媒体 digest（不限 kind），
|
||||
// 供模型切换后全量迁移向量空间（image + audio + video 等）。
|
||||
func (s *Store) StaleVecDigestsAll(currentModel string) ([]string, error) {
|
||||
return s.staleVecDigests(currentModel, "")
|
||||
}
|
||||
|
||||
// staleVecDigests 是 StaleVecDigests 的核心实现，kind=” 时不按 kind 过滤。
|
||||
// 废弃了"只迁移图片"的限定：模型切换后所有模态都应迁移到新向量空间。
|
||||
func (s *Store) staleVecDigests(currentModel string, kind string) ([]string, error) {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
query := `
|
||||
SELECT digest FROM media
|
||||
WHERE (COALESCE(vec_model,'') = '' OR vec_model != ?)`
|
||||
if kind != "" {
|
||||
query += ` AND kind = ?`
|
||||
}
|
||||
query += ` ORDER BY last_seen`
|
||||
|
||||
var args []interface{}
|
||||
args = append(args, currentModel)
|
||||
if kind != "" {
|
||||
args = append(args, kind)
|
||||
}
|
||||
|
||||
rows, err := s.db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var digests []string
|
||||
for rows.Next() {
|
||||
var d string
|
||||
if err := rows.Scan(&d); err == nil {
|
||||
digests = append(digests, d)
|
||||
}
|
||||
}
|
||||
return digests, rows.Err()
|
||||
}
|
||||
|
||||
// QueryMedia 用查询向量对所有已嵌入媒体做余弦相似度检索，返回 topK 个最相似的 Item。
|
||||
//
|
||||
// 这是跨模态检索的关键：查询可以是图片也可以是文本（经文本向量化后调用此方法），
|
||||
// 被查的媒体库里的每个 item 也有一个视觉向量。两者在同一空间比对，
|
||||
// 谁的相似度更高就召回谁——不再区分「这是一张图的查询」还是「这是一段文字的查询」，
|
||||
// 由向量空间的相似度自动判断。
|
||||
func (s *Store) QueryMedia(queryVec []float64, model string, topK int) ([]*Item, error) {
|
||||
hits, err := s.QueryMediaScored(queryVec, model, topK)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if hits == nil {
|
||||
return nil, nil
|
||||
}
|
||||
out := make([]*Item, len(hits))
|
||||
for i, h := range hits {
|
||||
out[i] = h.Item
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// MediaHit 是一条媒体相似度候选及其分数。
|
||||
// 跨模态融合需要原始分数做归一化，仅返回 Item 会丢掉尺度信息。
|
||||
type MediaHit struct {
|
||||
Item *Item
|
||||
Score float64
|
||||
}
|
||||
|
||||
// QueryMediaScored 用查询向量对所有已嵌入媒体做余弦相似度检索，
|
||||
// 返回 topK 个最相似的候选及其原始 cosine 分数（供跨模态归一化）。
|
||||
//
|
||||
// 分数只做排序，不在存储层设绝对阈值：多模态文本→图像的绝对 cosine 随模型、
|
||||
// 语言与数据域漂移，真实标定中有效命中可以低至 0.015。相关性门控在融合器中
|
||||
// 使用当前候选集合的相对分布完成。
|
||||
func (s *Store) QueryMediaScored(queryVec []float64, model string, topK int) ([]MediaHit, error) {
|
||||
return s.queryMediaScored(queryVec, model, topK)
|
||||
}
|
||||
|
||||
func (s *Store) queryMediaScored(queryVec []float64, model string, topK int) ([]MediaHit, error) {
|
||||
if topK <= 0 {
|
||||
topK = 20
|
||||
}
|
||||
if len(queryVec) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
|
||||
query := `SELECT digest, kind, mime, size, width, height,
|
||||
origin_path, tool, first_seen, last_seen,
|
||||
vec, vec_model
|
||||
FROM media WHERE vec IS NOT NULL AND vec != ''`
|
||||
var args []interface{}
|
||||
if model != "" {
|
||||
query += ` AND vec_model = ?`
|
||||
args = append(args, model)
|
||||
}
|
||||
rows, err := s.db.Query(query, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
type scored struct {
|
||||
item *Item
|
||||
score float64
|
||||
}
|
||||
var candidates []scored
|
||||
for rows.Next() {
|
||||
var it Item
|
||||
var kind string
|
||||
var origin, tool, vecJSON, vecModel sql.NullString
|
||||
if err := rows.Scan(&it.Digest, &kind, &it.MIME, &it.Size, &it.Width, &it.Height,
|
||||
&origin, &tool, &it.FirstSeen, &it.LastSeen,
|
||||
&vecJSON, &vecModel); err != nil {
|
||||
continue
|
||||
}
|
||||
it.Kind = Kind(kind)
|
||||
it.OriginPath = origin.String
|
||||
it.Tool = tool.String
|
||||
if !vecJSON.Valid || vecJSON.String == "" {
|
||||
continue
|
||||
}
|
||||
var itemVec []float64
|
||||
if err := json.Unmarshal([]byte(vecJSON.String), &itemVec); err != nil || len(itemVec) == 0 {
|
||||
continue
|
||||
}
|
||||
if len(itemVec) != len(queryVec) {
|
||||
continue // 维度不一致，跳过
|
||||
}
|
||||
score := cosineSimilaritySlice(queryVec, itemVec)
|
||||
if score > 0.05 {
|
||||
candidates = append(candidates, scored{&it, score})
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 按分数降序排序
|
||||
sort.Slice(candidates, func(i, j int) bool {
|
||||
return candidates[i].score > candidates[j].score
|
||||
})
|
||||
if len(candidates) > topK {
|
||||
candidates = candidates[:topK]
|
||||
}
|
||||
out := make([]MediaHit, len(candidates))
|
||||
for i, c := range candidates {
|
||||
out[i] = MediaHit{Item: c.item, Score: c.score}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// cosineSimilaritySlice 计算两个 []float64 向量的余弦相似度。
|
||||
func cosineSimilaritySlice(a, b []float64) float64 {
|
||||
var dot, normA, normB float64
|
||||
for i := range a {
|
||||
dot += a[i] * b[i]
|
||||
normA += a[i] * a[i]
|
||||
normB += b[i] * b[i]
|
||||
}
|
||||
if normA == 0 || normB == 0 {
|
||||
return 0
|
||||
}
|
||||
return dot / (math.Sqrt(normA) * math.Sqrt(normB))
|
||||
}
|
||||
|
||||
// ---- 扫描辅助 ----
|
||||
|
||||
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, vecJSON, vecModel sql.NullString
|
||||
if err := r.Scan(&it.Digest, &kind, &it.MIME, &it.Size, &it.Width, &it.Height,
|
||||
&origin, &tool, &it.FirstSeen, &it.LastSeen,
|
||||
&vecJSON, &vecModel); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
it.Kind = Kind(kind)
|
||||
it.OriginPath = origin.String
|
||||
it.Tool = tool.String
|
||||
if vecJSON.Valid && vecJSON.String != "" {
|
||||
var v []float64
|
||||
if err := json.Unmarshal([]byte(vecJSON.String), &v); err == nil {
|
||||
it.Vec = v
|
||||
}
|
||||
}
|
||||
it.VecModel = vecModel.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] + "..."
|
||||
}
|
||||
|
||||
// ResolvePrefix 把 digest 前缀补全为完整 digest。
|
||||
//
|
||||
// 日志、事件摘要与图库句子里出现的都是 shortDigest（前 12 位），
|
||||
// 因为完整的 64 位 sha256 会把一行文字撑爆、也无助于人眼辨认。
|
||||
// 反查时需要这个补全，否则那些短标记只能看不能用。
|
||||
//
|
||||
// 前缀歧义视为错误而非"取第一个"：挂错引用会让 GC 删掉仍被引用的内容，
|
||||
// 宁可这次绑定失败。12 位十六进制的碰撞概率极低，真撞上说明该用更长前缀。
|
||||
func (s *Store) ResolvePrefix(prefix string) (string, error) {
|
||||
prefix = strings.ToLower(strings.TrimSpace(prefix))
|
||||
if len(prefix) < 8 {
|
||||
return "", fmt.Errorf("digest 前缀过短（至少 8 位）: %q", prefix)
|
||||
}
|
||||
if len(prefix) == digestHexLen {
|
||||
// 已是完整 digest：仍要确认存在，否则调用方会挂一条孤儿引用
|
||||
if _, err := s.Stat(prefix); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return prefix, nil
|
||||
}
|
||||
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
rows, err := s.db.Query(
|
||||
`SELECT digest FROM media WHERE digest LIKE ? || '%' LIMIT 2`, prefix)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var found []string
|
||||
for rows.Next() {
|
||||
var d string
|
||||
if err := rows.Scan(&d); err != nil {
|
||||
return "", err
|
||||
}
|
||||
found = append(found, d)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
switch len(found) {
|
||||
case 0:
|
||||
return "", fmt.Errorf("digest 前缀 %q 未匹配到媒体", prefix)
|
||||
case 1:
|
||||
return found[0], nil
|
||||
default:
|
||||
return "", fmt.Errorf("digest 前缀 %q 有歧义（至少匹配 %s 和 %s）",
|
||||
prefix, found[0][:16], found[1][:16])
|
||||
}
|
||||
}
|
||||
@ -1,278 +0,0 @@
|
||||
package media
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// newTestStore 建一个临时媒体存储。
|
||||
// 参数保留只为兼容旧调用点；媒体不再有容量上限（生命周期由记忆块决定）。
|
||||
func newTestStore(t *testing.T, _ ...int64) *Store {
|
||||
t.Helper()
|
||||
s, err := New(t.TempDir())
|
||||
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 TestDelete_RemovesContentAndMetadata(t *testing.T) {
|
||||
// 删除块即删除内容：Delete 同时清掉 blob 与元数据。
|
||||
// 这不是 GC，也不看引用计数——调用方是记忆系统本身。
|
||||
s := newTestStore(t)
|
||||
d, _ := s.Put([]byte("held"), Item{MIME: "image/png"})
|
||||
other, _ := s.Put([]byte("orphaned"), Item{MIME: "image/png"})
|
||||
|
||||
if err := s.Delete(other); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s.Stat(other); err == nil {
|
||||
t.Fatal("删除后元数据应已移除")
|
||||
}
|
||||
if _, err := s.Get(other); err == nil {
|
||||
t.Fatal("删除后内容应已移除")
|
||||
}
|
||||
// 未被删除的项不受影响
|
||||
if _, err := s.Get(d); err != nil {
|
||||
t.Fatalf("未删除的内容不该受影响: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDelete_UnknownDigestIsNoop(t *testing.T) {
|
||||
s := newTestStore(t)
|
||||
if err := s.Delete(""); err != nil {
|
||||
t.Fatalf("空 digest 应为无操作: %v", err)
|
||||
}
|
||||
if err := s.Delete("ffffffffffffffff"); err != nil {
|
||||
t.Fatalf("不存在的 digest 应为无操作: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDescribe_Removed(t *testing.T) {
|
||||
// 媒体不再有文字描述：描述式索引是废弃的就机制。
|
||||
// 这里只保留一个编译期断言，确保 API 不会静默回归。
|
||||
s := newTestStore(t)
|
||||
if _, err := s.Put([]byte("img"), Item{MIME: "image/png"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
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 TestStats_CountsByKind(t *testing.T) {
|
||||
s := newTestStore(t)
|
||||
s.Put([]byte("i1"), Item{MIME: "image/png"})
|
||||
s.Put([]byte("i2"), Item{MIME: "image/jpeg"})
|
||||
s.Put([]byte("a1"), Item{MIME: "audio/wav"})
|
||||
|
||||
st := s.Stats()
|
||||
if st["count"].(int) != 3 {
|
||||
t.Fatalf("count 应为 3，实际 %v", st["count"])
|
||||
}
|
||||
if _, ok := st["described"]; ok {
|
||||
t.Fatal("媒体已不再有描述计数")
|
||||
}
|
||||
byKind := st["by_kind"].(map[string]int)
|
||||
if byKind["image"] != 2 || byKind["audio"] != 1 {
|
||||
t.Fatalf("by_kind 不对: %v", byKind)
|
||||
}
|
||||
}
|
||||
|
||||
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 TestPut_RejectsEmpty(t *testing.T) {
|
||||
s := newTestStore(t)
|
||||
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)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
d, _ := s1.Put([]byte("persistent-img"), Item{MIME: "image/png", OriginPath: "/tmp/x.png"})
|
||||
s1.SetVec(d, []float64{0.1, 0.2}, "test-space")
|
||||
s1.Close()
|
||||
|
||||
s2, err := New(dir)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s2.Close()
|
||||
|
||||
it, err := s2.Stat(d)
|
||||
if err != nil {
|
||||
t.Fatalf("重开后应能查到: %v", err)
|
||||
}
|
||||
if it.OriginPath != "/tmp/x.png" || len(it.Vec) != 2 || it.VecModel != "test-space" {
|
||||
t.Fatalf("元数据与向量应持久化: %+v", it)
|
||||
}
|
||||
data, err := s2.Get(d)
|
||||
if err != nil || string(data) != "persistent-img" {
|
||||
t.Fatalf("内容应持久化: %v / %q", err, data)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStaleVecDigests_TracksModelSwitch(t *testing.T) {
|
||||
// 模型切换后旧向量必须被重算：StaleVecDigests 是启动迁移的入口。
|
||||
s := newTestStore(t)
|
||||
d1, _ := s.Put([]byte("a"), Item{MIME: "image/png"})
|
||||
d2, _ := s.Put([]byte("b"), Item{MIME: "image/png"})
|
||||
s.SetVec(d1, []float64{0.1}, "space-a")
|
||||
|
||||
stale, err := s.StaleVecDigestsAll("space-a")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(stale) != 1 || stale[0] != d2 {
|
||||
t.Fatalf("只有未嵌入的 d2 需重算，实际 %v", stale)
|
||||
}
|
||||
|
||||
stale, err = s.StaleVecDigestsAll("space-b")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(stale) != 2 {
|
||||
t.Fatalf("换空间后两条都需重算，实际 %v", stale)
|
||||
}
|
||||
}
|
||||
@ -1,176 +0,0 @@
|
||||
package media
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestQueryMedia_BasicSimilarity(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
defer s.Close()
|
||||
|
||||
// 入库三张带向量的媒体：两张图、一段音频
|
||||
d1, _ := s.Put([]byte("img1"), Item{MIME: "image/png"})
|
||||
d2, _ := s.Put([]byte("img2"), Item{MIME: "image/jpeg"})
|
||||
d3, _ := s.Put([]byte("aud1"), Item{MIME: "audio/wav"})
|
||||
|
||||
// 模拟视觉嵌入：img1 和 img2 向量接近，aud1 远离
|
||||
vec1 := []float64{0.9, 0.1, 0.0, 0.0}
|
||||
vec2 := []float64{0.8, 0.2, 0.0, 0.0} // 与 vec1 相似
|
||||
vec3 := []float64{0.0, 0.0, 0.9, 0.1} // 与前两个完全不同
|
||||
|
||||
s.SetVec(d1, vec1, "test-clip")
|
||||
s.SetVec(d2, vec2, "test-clip")
|
||||
s.SetVec(d3, vec3, "test-clip")
|
||||
|
||||
// 用 vec1 作为查询：vec2 最相似，vec3 与 vec1 正交（相似度 0，被阈值过滤）
|
||||
results, err := s.QueryMedia(vec1, "test-clip", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// vec3 与 vec1 正交（余弦相似度 0），被 0.05 阈值正确剔除 → 只召回 2 个
|
||||
if len(results) != 2 {
|
||||
t.Fatalf("expected 2 results (正交的 aud1 被阈值过滤), got %d", len(results))
|
||||
}
|
||||
// 第一个应该是 img2（0.9 vs d1 的 1.0？不，这里算清楚）
|
||||
// vec1·vec2 与 vec1·vec1 比较：
|
||||
// sim(vec1,vec1) = 1.0（img1 与自身），sim(vec1,vec2) = 0.9*0.8+0.1*0.2 = 0.74
|
||||
// 所以 img1（自相似 1.0）排第一，img2 排第二
|
||||
if results[0].Digest != d1 {
|
||||
t.Errorf("expected d1 (自相似 1.0) as first, got %s", results[0].Digest)
|
||||
}
|
||||
if results[1].Digest != d2 {
|
||||
t.Errorf("expected d2 as second, got %s", results[1].Digest)
|
||||
}
|
||||
|
||||
// 验证分数：img1 与自身是 1.0
|
||||
selfScore := cosineSimilaritySlice(vec1, vec1)
|
||||
if math.Abs(selfScore-1.0) > 1e-10 {
|
||||
t.Errorf("self-similarity should be 1.0, got %f", selfScore)
|
||||
}
|
||||
|
||||
// img1 与 aud1 的相似度应该很低
|
||||
crossScore := cosineSimilaritySlice(vec1, vec3)
|
||||
if crossScore > 0.1 {
|
||||
t.Errorf("cross-modality similarity should be low, got %f", crossScore)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueryMedia_EmptyVecSkipped(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
defer s.Close()
|
||||
|
||||
d1, _ := s.Put([]byte("img1"), Item{MIME: "image/png"})
|
||||
_, _ = s.Put([]byte("img2"), Item{MIME: "image/png"})
|
||||
|
||||
// d1 有向量，d2 没有
|
||||
s.SetVec(d1, []float64{0.5, 0.5}, "test")
|
||||
// d2 留空
|
||||
|
||||
results, err := s.QueryMedia([]float64{0.5, 0.5}, "test", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(results) != 1 {
|
||||
t.Fatalf("expected 1 result (d2 has no vec), got %d", len(results))
|
||||
}
|
||||
if results[0].Digest != d1 {
|
||||
t.Errorf("expected d1, got %s", results[0].Digest)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueryMedia_DimensionMismatchSkipped(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
defer s.Close()
|
||||
|
||||
d1, _ := s.Put([]byte("img1"), Item{MIME: "image/png"})
|
||||
s.SetVec(d1, []float64{0.5, 0.5}, "model-A") // 2 维
|
||||
|
||||
// 查询用 3 维向量：维度不匹配，应该返回空
|
||||
results, err := s.QueryMedia([]float64{0.3, 0.3, 0.3}, "model-A", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(results) != 0 {
|
||||
t.Fatalf("expected 0 results (dim mismatch), got %d", len(results))
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueryMedia_EmptyQueryReturnsNil(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
defer s.Close()
|
||||
|
||||
results, err := s.QueryMedia(nil, "", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if results != nil {
|
||||
t.Fatalf("expected nil, got %d results", len(results))
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetVec_PersistsCorrectly(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
defer s.Close()
|
||||
|
||||
d, _ := s.Put([]byte("hello"), Item{MIME: "image/png"})
|
||||
vec := []float64{0.1, 0.2, 0.3, 0.4}
|
||||
s.SetVec(d, vec, "clip-vit-b32")
|
||||
|
||||
it, err := s.Stat(d)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if it.VecModel != "clip-vit-b32" {
|
||||
t.Errorf("VecModel = %q, want clip-vit-b32", it.VecModel)
|
||||
}
|
||||
if len(it.Vec) != 4 {
|
||||
t.Fatalf("Vec len = %d, want 4", len(it.Vec))
|
||||
}
|
||||
for i, v := range vec {
|
||||
if math.Abs(it.Vec[i]-v) > 1e-10 {
|
||||
t.Errorf("Vec[%d] = %f, want %f", i, it.Vec[i], v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestStaleVecDigests(t *testing.T) {
|
||||
s := newTestStore(t, 0)
|
||||
defer s.Close()
|
||||
|
||||
// 有向量且 vec_model 匹配 → 非 stale
|
||||
d1, _ := s.Put([]byte("img1"), Item{MIME: "image/png"})
|
||||
s.SetVec(d1, []float64{0.1}, "space-a")
|
||||
|
||||
// 有向量但 vec_model 旧 → stale
|
||||
d2, _ := s.Put([]byte("img2"), Item{MIME: "image/png"})
|
||||
s.SetVec(d2, []float64{0.2}, "space-old")
|
||||
|
||||
// 从未嵌入（vec_model 空）→ stale
|
||||
d3, _ := s.Put([]byte("img3"), Item{MIME: "image/png"})
|
||||
|
||||
// 与向量/描述无关的图片同样应被迁移：图片独立参与向量空间
|
||||
d4, _ := s.Put([]byte("img4"), Item{MIME: "image/png"})
|
||||
|
||||
// 音频不参与图片迁移（StaleVecDigests 只查 kind='image'）
|
||||
s.Put([]byte("aud1"), Item{MIME: "audio/wav"})
|
||||
|
||||
stale, err := s.StaleVecDigests("space-a")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// d1 匹配模型 → 非 stale；d2 旧模型 + d3 未嵌入 + d4 无描述图片 = 3 stale；aud1 不算
|
||||
if len(stale) != 3 {
|
||||
t.Fatalf("expected 3 stale digests (d2 旧模型 + d3 未嵌入 + d4 无描述), got %d: %v", len(stale), stale)
|
||||
}
|
||||
got := map[string]bool{}
|
||||
for _, d := range stale {
|
||||
got[d] = true
|
||||
}
|
||||
if !got[d2] || !got[d3] || !got[d4] {
|
||||
t.Errorf("expected d2, d3, d4 stale, got %v", stale)
|
||||
}
|
||||
if got[d1] {
|
||||
t.Errorf("d1 (匹配模型) 不应 stale")
|
||||
}
|
||||
}
|
||||
@ -1,277 +0,0 @@
|
||||
package media
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"compress/zlib"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"hash/crc32"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// 冒烟测试：走真实数据路径的端到端场景，而非孤立的 API 单测。
|
||||
//
|
||||
// 媒体存储现在只做内容寻址（CAS）：字节 + 元数据 + 向量。
|
||||
// “哪些字节还活着”由三层记忆持有的一等记忆块决定，调用方把该集合传给
|
||||
// GC/检索，本层不维护 media_refs/ref_count 这类平行账本。
|
||||
|
||||
// makePNG 生成一张 w×h 的条带 PNG，用真 PNG 而不是随机字节，
|
||||
// 让入库/回读/digest 走的是与生产一致的数据形态。
|
||||
func makePNG(w, h, variant int) []byte {
|
||||
palette := [][3]byte{
|
||||
{255, 0, 0}, {0, 192, 0}, {0, 0, 255}, {255, 220, 0}, {160, 0, 200},
|
||||
}
|
||||
var raw bytes.Buffer
|
||||
for y := 0; y < h; y++ {
|
||||
raw.WriteByte(0) // 每行的滤波器字节
|
||||
c := palette[(variant+y*3/h)%len(palette)]
|
||||
for x := 0; x < w; x++ {
|
||||
raw.Write(c[:])
|
||||
}
|
||||
}
|
||||
// 把 variant 写进首行头几个像素，确保不同 variant 字节必然不同。
|
||||
b := raw.Bytes()
|
||||
if len(b) > 8 {
|
||||
b[1] = byte(variant)
|
||||
b[2] = byte(variant >> 8)
|
||||
}
|
||||
|
||||
var z bytes.Buffer
|
||||
zw := zlib.NewWriter(&z)
|
||||
zw.Write(b)
|
||||
zw.Close()
|
||||
|
||||
chunk := func(typ string, data []byte) []byte {
|
||||
var out bytes.Buffer
|
||||
binary.Write(&out, binary.BigEndian, uint32(len(data)))
|
||||
out.WriteString(typ)
|
||||
out.Write(data)
|
||||
binary.Write(&out, binary.BigEndian, crc32.ChecksumIEEE(append([]byte(typ), data...)))
|
||||
return out.Bytes()
|
||||
}
|
||||
var ihdr bytes.Buffer
|
||||
binary.Write(&ihdr, binary.BigEndian, uint32(w))
|
||||
binary.Write(&ihdr, binary.BigEndian, uint32(h))
|
||||
ihdr.Write([]byte{8, 2, 0, 0, 0}) // 8bit 深度、truecolor
|
||||
|
||||
var out bytes.Buffer
|
||||
out.Write([]byte{0x89, 'P', 'N', 'G', '\r', '\n', 0x1a, '\n'})
|
||||
out.Write(chunk("IHDR", ihdr.Bytes()))
|
||||
out.Write(chunk("IDAT", z.Bytes()))
|
||||
out.Write(chunk("IEND", nil))
|
||||
return out.Bytes()
|
||||
}
|
||||
|
||||
func TestSmoke_SamePictureAcrossTurns(t *testing.T) {
|
||||
// 场景：用户连问几轮同一张截图。multimodal 每轮都会重新注入，
|
||||
// 内容寻址天然去重，磁盘上只应有一份。
|
||||
s := newTestStore(t, 50*1024*1024)
|
||||
png := makePNG(400, 400, 0)
|
||||
|
||||
var d0 string
|
||||
for turn := 1; turn <= 5; turn++ {
|
||||
// 走 data URL：这是 SetToolBlocks 实际给出的形态
|
||||
url := DataURL("image/png", png)
|
||||
mime, data, ok := ParseDataURL(url)
|
||||
if !ok {
|
||||
t.Fatalf("第 %d 轮 data URL 解析失败", turn)
|
||||
}
|
||||
d, err := s.Put(data, Item{
|
||||
MIME: mime, Width: 400, Height: 400,
|
||||
OriginPath: fmt.Sprintf("/tmp/probe_%d.png", turn),
|
||||
Tool: "multimodal_see_picture",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("第 %d 轮 Put: %v", turn, err)
|
||||
}
|
||||
if d0 == "" {
|
||||
d0 = d
|
||||
} else if d != d0 {
|
||||
t.Fatalf("同一张图第 %d 轮 digest 变了", turn)
|
||||
}
|
||||
}
|
||||
|
||||
st := s.Stats()
|
||||
if st["count"].(int) != 1 {
|
||||
t.Fatalf("5 轮同图应只存 1 份，实际 %v 条", st["count"])
|
||||
}
|
||||
if total := st["total_bytes"].(int64); total != int64(len(png)) {
|
||||
t.Fatalf("字节数应等于单张原图 %d，实际 %d", len(png), total)
|
||||
}
|
||||
// 三层记忆持有它；内容应仍可读
|
||||
if _, err := s.Get(d0); err != nil {
|
||||
t.Fatalf("内容应仍可读: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSmoke_VideoFramesDistinct(t *testing.T) {
|
||||
// 场景：see_video 抽 6 帧，帧间内容不同，应各存一份。
|
||||
s := newTestStore(t, 50*1024*1024)
|
||||
keep := map[string]bool{}
|
||||
for i := 0; i < 6; i++ {
|
||||
d, err := s.Put(makePNG(320, 240, i), Item{
|
||||
MIME: "image/jpeg", Width: 320, Height: 240, Tool: "multimodal_see_video",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("第 %d 帧: %v", i, err)
|
||||
}
|
||||
keep[d] = true
|
||||
}
|
||||
|
||||
if st := s.Stats(); st["count"].(int) != 6 {
|
||||
t.Fatalf("6 帧应各存一份，实际 %v 条", st["count"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestSmoke_NearestNeighborVectorRetrieve(t *testing.T) {
|
||||
// 场景 C：图片只按自己的原生向量被检索。
|
||||
// 没有描述文本参与——描述式索引是废弃的就机制。
|
||||
s := newTestStore(t)
|
||||
|
||||
pic, _ := s.Put(makePNG(400, 400, 0), Item{MIME: "image/png", Tool: "multimodal_see_picture"})
|
||||
s.SetVec(pic, []float64{1, 0, 0, 0}, "space")
|
||||
var frames []string
|
||||
for i := 0; i < 6; i++ {
|
||||
d, _ := s.Put(makePNG(320, 240, i), Item{MIME: "image/jpeg", Tool: "multimodal_see_video"})
|
||||
s.SetVec(d, []float64{1, 1, float64(i) / 10, 0}, "space")
|
||||
frames = append(frames, d)
|
||||
}
|
||||
|
||||
hits, err := s.QueryMediaScored([]float64{1, 0, 0, 0}, "space", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(hits) != 7 {
|
||||
t.Fatalf("7 份媒体都有同空间向量，应全部可召，实际 %d", len(hits))
|
||||
}
|
||||
if hits[0].Item.Digest != pic {
|
||||
t.Fatalf("与查询同向的应是第一命中，实际 %s", shortDigest(hits[0].Item.Digest))
|
||||
}
|
||||
|
||||
// 不同向量空间/模型的条目不得参与：坐标系不同，余弦无意义。
|
||||
foreign := frames[0]
|
||||
if err := s.SetVec(foreign, []float64{1, 0, 0, 0}, "other-space"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
hits, err = s.QueryMediaScored([]float64{1, 0, 0, 0}, "space", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, h := range hits {
|
||||
if h.Item.Digest == foreign {
|
||||
t.Fatal("另一套空间（other-space）的向量不该被 space 查询召回")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSmoke_ContentSurvivesLayerMigration(t *testing.T) {
|
||||
// 场景：同一份媒体随记忆块从 Context 迁移到 Document 再到 Graph。
|
||||
// 迁移的是块本身，digest 不变，因此内容在整条链路上始终可读。
|
||||
s := newTestStore(t, 50*1024*1024)
|
||||
png := makePNG(400, 400, 0)
|
||||
d, _ := s.Put(png, Item{MIME: "image/png", Tool: "multimodal_see_picture"})
|
||||
|
||||
// 迁移过程中该 digest 始终可读
|
||||
for _, layer := range []string{"context", "document", "graph"} {
|
||||
if got, err := s.Get(d); err != nil || !bytes.Equal(got, png) {
|
||||
t.Fatalf("迁移到 %s 时内容应完好: %v", layer, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSmoke_DeleteRemovesOnlyThatContent(t *testing.T) {
|
||||
// 场景：某个工具产出的一次性图片所在的记忆块被删除时，
|
||||
// 只有它自己的内容被删；其他块的内容一个都不能少。
|
||||
s := newTestStore(t)
|
||||
|
||||
held, _ := s.Put(makePNG(400, 400, 0), Item{MIME: "image/png"})
|
||||
var frames []string
|
||||
for i := 0; i < 6; i++ {
|
||||
d, _ := s.Put(makePNG(320, 240, i), Item{MIME: "image/jpeg"})
|
||||
frames = append(frames, d)
|
||||
}
|
||||
var ephemeral []string
|
||||
for i := 0; i < 20; i++ {
|
||||
d, _ := s.Put(makePNG(100, 100, 1000+i), Item{MIME: "image/png", Tool: "cmd_run"})
|
||||
ephemeral = append(ephemeral, d)
|
||||
}
|
||||
|
||||
before := s.Stats()["count"].(int)
|
||||
for _, d := range ephemeral {
|
||||
if err := s.Delete(d); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
after := s.Stats()["count"].(int)
|
||||
if after != before-20 {
|
||||
t.Fatalf("条目数应从 %d 降到 %d，实际 %d", before, before-20, after)
|
||||
}
|
||||
if _, err := s.Get(held); err != nil {
|
||||
t.Fatalf("被保留的内容被误删: %v", err)
|
||||
}
|
||||
for i, f := range frames {
|
||||
if _, err := s.Get(f); err != nil {
|
||||
t.Fatalf("第 %d 帧被误删: %v", i, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSmoke_FullLifecycleAcrossRestart(t *testing.T) {
|
||||
// 端到端：入库 → 嵌入 → 删除一些内容 → 重启 → 向量检索，
|
||||
// 并确认磁盘与元数据不出现双向孤儿。记忆的意义就在于跨重启还在。
|
||||
dir := t.TempDir()
|
||||
s, err := New(dir)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
png := makePNG(400, 400, 0)
|
||||
pic, _ := s.Put(png, Item{MIME: "image/png", Width: 400, Height: 400, Tool: "multimodal_see_picture"})
|
||||
s.SetVec(pic, []float64{1, 0, 0}, "space")
|
||||
for i := 0; i < 6; i++ {
|
||||
d, _ := s.Put(makePNG(320, 240, i), Item{MIME: "image/jpeg", Tool: "multimodal_see_video"})
|
||||
s.SetVec(d, []float64{0, 1, float64(i)}, "space")
|
||||
}
|
||||
for i := 0; i < 10; i++ {
|
||||
d, _ := s.Put(makePNG(64, 64, 2000+i), Item{MIME: "image/png", Tool: "cmd_run"})
|
||||
if err := s.Delete(d); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
beforeCount := s.Stats()["count"].(int)
|
||||
s.Close()
|
||||
|
||||
s2, err := New(dir)
|
||||
if err != nil {
|
||||
t.Fatalf("重开失败: %v", err)
|
||||
}
|
||||
defer s2.Close()
|
||||
|
||||
if got := s2.Stats()["count"].(int); got != beforeCount {
|
||||
t.Fatalf("重开后条目数变了: %d → %d", beforeCount, got)
|
||||
}
|
||||
it, err := s2.Stat(pic)
|
||||
if err != nil {
|
||||
t.Fatalf("重开后查不到: %v", err)
|
||||
}
|
||||
if len(it.Vec) != 3 || it.VecModel != "space" {
|
||||
t.Fatalf("向量未持久化: %+v", it)
|
||||
}
|
||||
data, err := s2.Get(pic)
|
||||
if err != nil || !bytes.Equal(data, png) {
|
||||
t.Fatalf("重开后内容不一致: %v", err)
|
||||
}
|
||||
hits, err := s2.QueryMediaScored([]float64{1, 0, 0}, "space", 10)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(hits) == 0 || hits[0].Item.Digest != pic {
|
||||
t.Fatal("重开后向量检索应仍能命中")
|
||||
}
|
||||
|
||||
// 磁盘文件数 == 元数据条数：无「元数据在文件没了」也无「文件在元数据没了」
|
||||
if n := blobFileCount(t, s2); n != beforeCount {
|
||||
t.Fatalf("磁盘 blob=%d 与元数据=%d 不一致", n, beforeCount)
|
||||
}
|
||||
t.Logf("跨重启：%d 条目、向量与内容全部完好", beforeCount)
|
||||
}
|
||||
@ -1,166 +0,0 @@
|
||||
package media
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/rand"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TestSoak_SustainedMixedLoad 长稳测试：持续混合负载下不变量不破。
|
||||
// 用 -run TestSoak -timeout 300s 单独跑，默认 short 模式跳过。
|
||||
//
|
||||
// 媒体没有独立生命周期管理：blob 是记忆块的内容，块被删除时内容随之删除。
|
||||
func TestSoak_SustainedMixedLoad(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("long soak test; run with -run TestSoak")
|
||||
}
|
||||
dur := 60 * time.Second
|
||||
s := newTestStore(t)
|
||||
|
||||
// 常驻受保护区：全程被记忆块持有，模拟 Graph L3 中的块
|
||||
const keepN = 20
|
||||
keep := make([]string, keepN)
|
||||
keepData := make([][]byte, keepN)
|
||||
for i := range keep {
|
||||
d := make([]byte, 4096)
|
||||
rand.Read(d)
|
||||
d = append([]byte(fmt.Sprintf("keep-%d-", i)), d...)
|
||||
dg, err := s.Put(d, Item{MIME: "image/png"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
keep[i] = dg
|
||||
keepData[i] = d
|
||||
}
|
||||
|
||||
stop := make(chan struct{})
|
||||
var wg sync.WaitGroup
|
||||
var puts, gets, deletes, embeds, searches atomic.Int64
|
||||
var fatal atomic.Int64
|
||||
|
||||
worker := func(name string, fn func(iter int) error) {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for i := 0; ; i++ {
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
default:
|
||||
}
|
||||
if err := fn(i); err != nil {
|
||||
fatal.Add(1)
|
||||
t.Errorf("%s 第 %d 次失败: %v", name, i, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// 写入者 ×3：持续写入一次性内容（无人持有）
|
||||
for w := 0; w < 3; w++ {
|
||||
wid := w
|
||||
worker(fmt.Sprintf("put-%d", wid), func(i int) error {
|
||||
b := make([]byte, 2048)
|
||||
rand.Read(b)
|
||||
b = append([]byte(fmt.Sprintf("eph-%d-%d-", wid, i)), b...)
|
||||
if _, err := s.Put(b, Item{MIME: "image/png", Tool: "cmd_run"}); err != nil {
|
||||
return err
|
||||
}
|
||||
puts.Add(1)
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
// 读取者 ×3：受保护集必须始终完好
|
||||
for r := 0; r < 3; r++ {
|
||||
worker("get", func(i int) error {
|
||||
idx := i % keepN
|
||||
got, err := s.Get(keep[idx])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !bytes.Equal(got, keepData[idx]) {
|
||||
return fmt.Errorf("内容被改 %s", shortDigest(keep[idx]))
|
||||
}
|
||||
gets.Add(1)
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
// 删除者：持续删除一次性内容（模拟块创建后又被遗忘）
|
||||
worker("delete", func(i int) error {
|
||||
b := make([]byte, 2048)
|
||||
rand.Read(b)
|
||||
b = append([]byte(fmt.Sprintf("del-%d-", i)), b...)
|
||||
d, err := s.Put(b, Item{MIME: "image/png", Tool: "cmd_run"})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := s.Delete(d); err != nil {
|
||||
return err
|
||||
}
|
||||
deletes.Add(1)
|
||||
time.Sleep(time.Millisecond)
|
||||
return nil
|
||||
})
|
||||
|
||||
// 向量写入者：持续给新内容嵌入并删除（模拟启动迁移/短命媒体）
|
||||
worker("embed", func(i int) error {
|
||||
b := make([]byte, 1024)
|
||||
rand.Read(b)
|
||||
b = append([]byte(fmt.Sprintf("emb-%d-", i)), b...)
|
||||
d, err := s.Put(b, Item{MIME: "image/png", Tool: "cmd_run"})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := s.SetVec(d, []float64{1, float64(i % 7)}, "soak-space"); err != nil {
|
||||
return err
|
||||
}
|
||||
embeds.Add(1)
|
||||
time.Sleep(time.Millisecond)
|
||||
return nil
|
||||
})
|
||||
|
||||
// 检索者
|
||||
worker("search", func(i int) error {
|
||||
if _, err := s.QueryMediaScored([]float64{1, 0}, "soak-space", 20); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := s.Stat(keep[i%keepN]); err != nil {
|
||||
return err
|
||||
}
|
||||
searches.Add(1)
|
||||
time.Sleep(2 * time.Millisecond)
|
||||
return nil
|
||||
})
|
||||
|
||||
time.Sleep(dur)
|
||||
close(stop)
|
||||
wg.Wait()
|
||||
|
||||
if n := fatal.Load(); n > 0 {
|
||||
t.Fatalf("%d 个 worker 报致命错误", n)
|
||||
}
|
||||
|
||||
t.Logf("%v 内: put=%d get=%d delete=%d embed=%d search=%d",
|
||||
dur, puts.Load(), gets.Load(), deletes.Load(), embeds.Load(), searches.Load())
|
||||
|
||||
// 收尾断言
|
||||
for i, d := range keep {
|
||||
got, err := s.Get(d)
|
||||
if err != nil {
|
||||
t.Fatalf("受保护项丢失 %s: %v", shortDigest(d), err)
|
||||
}
|
||||
if !bytes.Equal(got, keepData[i]) {
|
||||
t.Fatalf("受保护项内容变了 %s", shortDigest(d))
|
||||
}
|
||||
}
|
||||
|
||||
st := s.Stats()
|
||||
t.Logf("收尾: 条目=%v 字节=%v 类型=%v", st["count"], st["total_bytes"], st["by_kind"])
|
||||
}
|
||||
@ -1,527 +0,0 @@
|
||||
package media
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/rand"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// 压力测试与冒烟测试。
|
||||
//
|
||||
// 关注点不是吞吐数字，而是并发下的不变量是否被破坏：
|
||||
// 1. GC 与读写并发时，被记忆块持有的内容绝不能被删
|
||||
// 2. 同内容并发 Put 只落一份磁盘、digest 一致
|
||||
// 3. SQLite 在多 goroutine 下不出现 "database is locked"
|
||||
//
|
||||
// 存活判定不再依赖 media_refs/ref_count：调用方把「三层记忆当前持有的
|
||||
// digest 集合」传给 GC，本层只做 CAS。
|
||||
|
||||
func randBytes(t *testing.T, n int) []byte {
|
||||
t.Helper()
|
||||
b := make([]byte, n)
|
||||
if _, err := rand.Read(b); err != nil {
|
||||
t.Fatalf("rand: %v", err)
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// blobFileCount 统计 CAS 目录下的实际文件数（不含 .tmp）。
|
||||
func blobFileCount(t *testing.T, s *Store) int {
|
||||
t.Helper()
|
||||
n := 0
|
||||
filepath.Walk(s.blobDir, func(path string, info os.FileInfo, err error) error {
|
||||
if err != nil || info == nil || info.IsDir() {
|
||||
return nil
|
||||
}
|
||||
if strings.HasSuffix(path, ".tmp") {
|
||||
t.Errorf("残留临时文件: %s", path)
|
||||
return nil
|
||||
}
|
||||
n++
|
||||
return nil
|
||||
})
|
||||
return n
|
||||
}
|
||||
|
||||
func TestStress_ConcurrentPutSameContent(t *testing.T) {
|
||||
// 同一内容被 N 个 goroutine 同时 Put：digest 必须一致，磁盘只一份。
|
||||
// 现实对应：see_video 抽出的相邻帧、用户连发同一张图。
|
||||
s := newTestStore(t, 0)
|
||||
data := randBytes(t, 64*1024)
|
||||
|
||||
const workers = 32
|
||||
var wg sync.WaitGroup
|
||||
digests := make([]string, workers)
|
||||
errs := make([]error, workers)
|
||||
for i := 0; i < workers; i++ {
|
||||
wg.Add(1)
|
||||
go func(idx int) {
|
||||
defer wg.Done()
|
||||
d, err := s.Put(data, Item{MIME: "image/png", OriginPath: fmt.Sprintf("/tmp/%d.png", idx)})
|
||||
digests[idx] = d
|
||||
errs[idx] = err
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
for i, err := range errs {
|
||||
if err != nil {
|
||||
t.Fatalf("worker %d Put 失败: %v", i, err)
|
||||
}
|
||||
}
|
||||
first := digests[0]
|
||||
for i, d := range digests {
|
||||
if d != first {
|
||||
t.Fatalf("worker %d digest 不一致: %s vs %s", i, d, first)
|
||||
}
|
||||
}
|
||||
if n := blobFileCount(t, s); n != 1 {
|
||||
t.Fatalf("同一内容应只落一份 blob，实际 %d 个文件", n)
|
||||
}
|
||||
if got, err := s.Get(first); err != nil || !bytes.Equal(got, data) {
|
||||
t.Fatalf("内容应可完整读回: err=%v len=%d", err, len(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestStress_ConcurrentPutDistinctContent(t *testing.T) {
|
||||
// 大量不同内容并发入库：不丢条目、不串内容。
|
||||
s := newTestStore(t, 0)
|
||||
const workers = 16
|
||||
const perWorker = 25
|
||||
|
||||
var wg sync.WaitGroup
|
||||
var failed atomic.Int64
|
||||
type rec struct {
|
||||
digest string
|
||||
data []byte
|
||||
}
|
||||
recCh := make(chan rec, workers*perWorker)
|
||||
|
||||
for w := 0; w < workers; w++ {
|
||||
wg.Add(1)
|
||||
go func(wid int) {
|
||||
defer wg.Done()
|
||||
for i := 0; i < perWorker; i++ {
|
||||
data := []byte(fmt.Sprintf("w%d-i%d-", wid, i))
|
||||
data = append(data, randBytes(t, 512)...)
|
||||
d, err := s.Put(data, Item{MIME: "image/png"})
|
||||
if err != nil {
|
||||
failed.Add(1)
|
||||
continue
|
||||
}
|
||||
recCh <- rec{digest: d, data: data}
|
||||
}
|
||||
}(w)
|
||||
}
|
||||
wg.Wait()
|
||||
close(recCh)
|
||||
|
||||
if n := failed.Load(); n > 0 {
|
||||
t.Fatalf("%d 次 Put 失败", n)
|
||||
}
|
||||
|
||||
var records []rec
|
||||
for r := range recCh {
|
||||
records = append(records, r)
|
||||
}
|
||||
if len(records) != workers*perWorker {
|
||||
t.Fatalf("应有 %d 条记录，实际 %d", workers*perWorker, len(records))
|
||||
}
|
||||
|
||||
// 逐条回读校验内容没串
|
||||
for _, r := range records {
|
||||
got, err := s.Get(r.digest)
|
||||
if err != nil {
|
||||
t.Fatalf("读 %s 失败: %v", shortDigest(r.digest), err)
|
||||
}
|
||||
if !bytes.Equal(got, r.data) {
|
||||
t.Fatalf("内容串了: %s", shortDigest(r.digest))
|
||||
}
|
||||
}
|
||||
|
||||
st := s.Stats()
|
||||
if st["count"].(int) != len(records) {
|
||||
t.Fatalf("库内条目应为 %d，实际 %v", len(records), st["count"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestStress_ConcurrentDeleteAndPut(t *testing.T) {
|
||||
// 删除与写入并发：核心断言是被保留的内容永远可读，
|
||||
// 删除只影响目标 digest，不误伤其他内容。
|
||||
s := newTestStore(t)
|
||||
|
||||
const heldCount = 8
|
||||
held := make([]string, heldCount)
|
||||
for i := range held {
|
||||
d, err := s.Put([]byte(fmt.Sprintf("payload-%d", i)), Item{MIME: "image/png"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
held[i] = d
|
||||
}
|
||||
|
||||
const workers = 16
|
||||
const rounds = 30
|
||||
var wg sync.WaitGroup
|
||||
for w := 0; w < workers; w++ {
|
||||
wg.Add(1)
|
||||
go func(wid int) {
|
||||
defer wg.Done()
|
||||
for r := 0; r < rounds; r++ {
|
||||
d, err := s.Put([]byte(fmt.Sprintf("tmp-%d-%d", wid, r)), Item{MIME: "image/png"})
|
||||
if err != nil {
|
||||
t.Errorf("Put: %v", err)
|
||||
return
|
||||
}
|
||||
if err := s.Delete(d); err != nil {
|
||||
t.Errorf("Delete: %v", err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}(w)
|
||||
}
|
||||
// 并发读取被保留内容
|
||||
for rdr := 0; rdr < 4; rdr++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for r := 0; r < rounds; r++ {
|
||||
for i, d := range held {
|
||||
if _, err := s.Get(d); err != nil {
|
||||
t.Errorf("内容 %d 被误删: %v", i, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
for i, d := range held {
|
||||
if _, err := s.Get(d); err != nil {
|
||||
t.Fatalf("仍被保留的第 %d 项不可读: %v", i, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestStress_DeleteConcurrentWithReads(t *testing.T) {
|
||||
// 删除与读取并发。最重要的断言：被保留的内容在整个过程中始终可读。
|
||||
s := newTestStore(t)
|
||||
|
||||
const protectedCount = 10
|
||||
protected := make([]string, protectedCount)
|
||||
protectedData := make([][]byte, protectedCount)
|
||||
for i := range protected {
|
||||
data := append([]byte(fmt.Sprintf("protected-%d-", i)), randBytes(t, 256)...)
|
||||
d, err := s.Put(data, Item{MIME: "image/png"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
protected[i] = d
|
||||
protectedData[i] = data
|
||||
}
|
||||
|
||||
stop := make(chan struct{})
|
||||
var wg sync.WaitGroup
|
||||
var readErr atomic.Int64
|
||||
var deleteCount atomic.Int64
|
||||
var putCount atomic.Int64
|
||||
|
||||
// 写入者：持续 Put 一次性内容再删除（模拟块创建后又被遗忘）
|
||||
for w := 0; w < 4; w++ {
|
||||
wg.Add(1)
|
||||
go func(wid int) {
|
||||
defer wg.Done()
|
||||
i := 0
|
||||
for {
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
default:
|
||||
}
|
||||
data := append([]byte(fmt.Sprintf("ephemeral-%d-%d-", wid, i)), randBytes(t, 128)...)
|
||||
d, err := s.Put(data, Item{MIME: "image/png"})
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
putCount.Add(1)
|
||||
if err := s.Delete(d); err == nil {
|
||||
deleteCount.Add(1)
|
||||
}
|
||||
i++
|
||||
}
|
||||
}(w)
|
||||
}
|
||||
|
||||
// 读取者：反复读受保护内容，任何一次失败都是致命的
|
||||
for r := 0; r < 4; r++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for {
|
||||
select {
|
||||
case <-stop:
|
||||
return
|
||||
default:
|
||||
}
|
||||
for i, d := range protected {
|
||||
got, err := s.Get(d)
|
||||
if err != nil {
|
||||
readErr.Add(1)
|
||||
t.Errorf("受保护内容读失败 %s: %v", shortDigest(d), err)
|
||||
return
|
||||
}
|
||||
if !bytes.Equal(got, protectedData[i]) {
|
||||
readErr.Add(1)
|
||||
t.Errorf("受保护内容被改 %s", shortDigest(d))
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
time.Sleep(1500 * time.Millisecond)
|
||||
close(stop)
|
||||
wg.Wait()
|
||||
|
||||
if n := readErr.Load(); n > 0 {
|
||||
t.Fatalf("受保护内容读取失败 %d 次", n)
|
||||
}
|
||||
t.Logf("并发窗口内: Put=%d Delete=%d", putCount.Load(), deleteCount.Load())
|
||||
|
||||
// 收尾确认：受保护的一个都没少
|
||||
for i, d := range protected {
|
||||
got, err := s.Get(d)
|
||||
if err != nil || !bytes.Equal(got, protectedData[i]) {
|
||||
t.Fatalf("收尾检查失败 %s: %v", shortDigest(d), err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestStress_SetVecConcurrentWithQuery(t *testing.T) {
|
||||
// 嵌入写入与向量检索并发（启动时的向量迁移就会长期这样跑）。
|
||||
s := newTestStore(t)
|
||||
const n = 60
|
||||
digests := make([]string, n)
|
||||
for i := range digests {
|
||||
d, err := s.Put([]byte(fmt.Sprintf("img-%d", i)), Item{MIME: "image/png"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
digests[i] = d
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
var writeErr, queryErr atomic.Int64
|
||||
|
||||
// 向量写入者
|
||||
for w := 0; w < 4; w++ {
|
||||
wg.Add(1)
|
||||
go func(wid int) {
|
||||
defer wg.Done()
|
||||
for i := wid; i < n; i += 4 {
|
||||
vec := []float64{1, float64(i) / 100, 0}
|
||||
if err := s.SetVec(digests[i], vec, "space"); err != nil {
|
||||
writeErr.Add(1)
|
||||
}
|
||||
}
|
||||
}(w)
|
||||
}
|
||||
|
||||
// 检索者
|
||||
for r := 0; r < 3; r++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for i := 0; i < 50; i++ {
|
||||
if _, err := s.QueryMediaScored([]float64{1, 0, 0}, "space", 20); err != nil {
|
||||
queryErr.Add(1)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if v := writeErr.Load(); v > 0 {
|
||||
t.Fatalf("SetVec 失败 %d 次", v)
|
||||
}
|
||||
if v := queryErr.Load(); v > 0 {
|
||||
t.Fatalf("QueryMediaScored 失败 %d 次", v)
|
||||
}
|
||||
|
||||
// 全部应已嵌入，且都在同一空间
|
||||
stale, err := s.StaleVecDigestsAll("space")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(stale) != 0 {
|
||||
t.Fatalf("应全部已嵌入，仍有 %d 条未嵌入", len(stale))
|
||||
}
|
||||
got, err := s.QueryMediaScored([]float64{1, 0, 0}, "space", 1000)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(got) != n {
|
||||
t.Fatalf("应检索到 %d 条，实际 %d", n, len(got))
|
||||
}
|
||||
}
|
||||
|
||||
func TestStress_DeleteUnderLoad(t *testing.T) {
|
||||
// 持续写入 + 删除下，被保留的项必须始终可读。
|
||||
s := newTestStore(t)
|
||||
|
||||
const keepN = 3
|
||||
keepList := make([]string, keepN)
|
||||
for i := range keepList {
|
||||
data := append([]byte(fmt.Sprintf("keep-%d-", i)), randBytes(t, 32*1024)...)
|
||||
d, err := s.Put(data, Item{MIME: "image/png"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
keepList[i] = d
|
||||
}
|
||||
|
||||
for round := 0; round < 30; round++ {
|
||||
for i := 0; i < 3; i++ {
|
||||
data := append([]byte(fmt.Sprintf("tmp-%d-%d-", round, i)), randBytes(t, 16*1024)...)
|
||||
d, err := s.Put(data, Item{MIME: "image/png"})
|
||||
if err != nil {
|
||||
t.Fatalf("round %d Put: %v", round, err)
|
||||
}
|
||||
if err := s.Delete(d); err != nil {
|
||||
t.Fatalf("round %d Delete: %v", round, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 被保留的项必须都在
|
||||
for _, d := range keepList {
|
||||
if _, err := s.Get(d); err != nil {
|
||||
t.Fatalf("被保留项被误删 %s: %v", shortDigest(d), err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestStress_ReopenAfterHeavyChurn(t *testing.T) {
|
||||
// 大量写入 + 删除之后重开：元数据与磁盘不该出现互相不认的孤儿。
|
||||
dir := t.TempDir()
|
||||
s1, err := New(dir)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var kept []string
|
||||
for i := 0; i < 100; i++ {
|
||||
data := append([]byte(fmt.Sprintf("churn-%d-", i)), randBytes(t, 256)...)
|
||||
d, err := s1.Put(data, Item{MIME: "image/png"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if i%5 == 0 {
|
||||
kept = append(kept, d)
|
||||
} else if err := s1.Delete(d); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
beforeStats := s1.Stats()
|
||||
s1.Close()
|
||||
|
||||
s2, err := New(dir)
|
||||
if err != nil {
|
||||
t.Fatalf("重开失败: %v", err)
|
||||
}
|
||||
defer s2.Close()
|
||||
|
||||
afterStats := s2.Stats()
|
||||
if beforeStats["count"] != afterStats["count"] {
|
||||
t.Fatalf("重开后条目数变了: %v → %v", beforeStats["count"], afterStats["count"])
|
||||
}
|
||||
|
||||
// 每条元数据都应有对应磁盘文件（无「元数据在文件没了」的孤儿）
|
||||
rows, err := s2.db.Query(`SELECT digest FROM media`)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var missing int
|
||||
for rows.Next() {
|
||||
var d string
|
||||
if rows.Scan(&d) != nil {
|
||||
continue
|
||||
}
|
||||
if _, err := os.Stat(s2.blobPath(d)); err != nil {
|
||||
missing++
|
||||
if missing <= 3 {
|
||||
t.Errorf("元数据存在但 blob 缺失: %s", shortDigest(d))
|
||||
}
|
||||
}
|
||||
}
|
||||
rows.Close()
|
||||
if missing > 0 {
|
||||
t.Fatalf("%d 条元数据没有对应文件", missing)
|
||||
}
|
||||
|
||||
// 磁盘文件数应等于元数据条数（无「文件在元数据没了」的孤儿）
|
||||
if n := blobFileCount(t, s2); n != afterStats["count"].(int) {
|
||||
t.Fatalf("磁盘文件 %d 与元数据 %v 不一致", n, afterStats["count"])
|
||||
}
|
||||
|
||||
for _, d := range kept {
|
||||
if _, err := s2.Get(d); err != nil {
|
||||
t.Fatalf("被持有项重开后读不到 %s: %v", shortDigest(d), err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestStress_LargeBlob(t *testing.T) {
|
||||
// 单个大文件：see_video 10 帧 × 2MB 是现实上限附近。
|
||||
s := newTestStore(t, 0)
|
||||
data := randBytes(t, 4*1024*1024) // 4MB
|
||||
|
||||
d, err := s.Put(data, Item{MIME: "image/jpeg", Width: 1920, Height: 1080})
|
||||
if err != nil {
|
||||
t.Fatalf("4MB Put 失败: %v", err)
|
||||
}
|
||||
got, err := s.Get(d)
|
||||
if err != nil {
|
||||
t.Fatalf("4MB Get 失败: %v", err)
|
||||
}
|
||||
if !bytes.Equal(got, data) {
|
||||
t.Fatal("4MB 内容回读不一致")
|
||||
}
|
||||
it, _ := s.Stat(d)
|
||||
if it.Size != int64(len(data)) {
|
||||
t.Fatalf("Size 记录错: %d vs %d", it.Size, len(data))
|
||||
}
|
||||
}
|
||||
|
||||
func TestStress_DataURLRoundTripAtScale(t *testing.T) {
|
||||
// data URL 往返是插件注入的实际路径（SetToolBlocks 给的就是 data URL）。
|
||||
s := newTestStore(t, 0)
|
||||
for i := 0; i < 50; i++ {
|
||||
raw := randBytes(t, 2048)
|
||||
url := DataURL("image/png", raw)
|
||||
mime, decoded, ok := ParseDataURL(url)
|
||||
if !ok {
|
||||
t.Fatalf("第 %d 次解析失败", i)
|
||||
}
|
||||
if mime != "image/png" || !bytes.Equal(decoded, raw) {
|
||||
t.Fatalf("第 %d 次往返不一致", i)
|
||||
}
|
||||
d, err := s.Put(decoded, Item{MIME: mime})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
back, err := s.Get(d)
|
||||
if err != nil || !bytes.Equal(back, raw) {
|
||||
t.Fatalf("第 %d 次入库回读不一致: %v", i, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user