diff --git a/internal/storage/cas.go b/internal/storage/cas.go new file mode 100644 index 00000000..259a7e7d --- /dev/null +++ b/internal/storage/cas.go @@ -0,0 +1,541 @@ +package storage + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "math" + "os" + "path/filepath" + "strings" + "sync" + "time" + + "github.com/jackc/pgx/v5/pgxpool" +) + +// Embedder 向量生成抽象接口 +type Embedder interface { + EmbedOne(ctx context.Context, text string) ([]float32, error) +} + +// CAS 内容寻址存储与切片防爆炸去重管理器 +type CAS struct { + baseDir string + pool *pgxpool.Pool + embedder Embedder + + // 离线/降级内存索引 (当数据库未连接时) + mu sync.RWMutex + memoryBlobs map[string]*BlobInfo + memoryRefs map[int64]*FileReference + memoryChunks []*MemoryChunk + nextRefID int64 + nextChunkID int64 +} + +// NewCAS 创建 CAS 存储管理器 +func NewCAS(baseDir string, pool *pgxpool.Pool, embedder Embedder) (*CAS, error) { + if baseDir == "" { + baseDir = "./data/dev/storage" + } + blobDir := filepath.Join(baseDir, "blobs") + if err := os.MkdirAll(blobDir, 0755); err != nil { + return nil, fmt.Errorf("failed to create blob dir: %w", err) + } + + return &CAS{ + baseDir: baseDir, + pool: pool, + embedder: embedder, + memoryBlobs: make(map[string]*BlobInfo), + memoryRefs: make(map[int64]*FileReference), + memoryChunks: make([]*MemoryChunk, 0), + nextRefID: 1, + nextChunkID: 1, + }, nil +} + +// StoreFile 写入物理文件并进行 CAS Hash 去重 +func (c *CAS) StoreFile(ctx context.Context, projectID, userID int64, fileName, fileCategory string, data []byte, mimeType string) (*FileReference, error) { + if len(data) == 0 { + return nil, errors.New("cannot store empty file") + } + if fileCategory == "" { + fileCategory = "material" + } + if mimeType == "" { + mimeType = "application/octet-stream" + } + + hashBytes := sha256.Sum256(data) + hashStr := hex.EncodeToString(hashBytes[:]) + sizeBytes := int64(len(data)) + + relPath := fmt.Sprintf("blobs/%s/%s/%s.bin", hashStr[:2], hashStr[2:4], hashStr) + fullPath := filepath.Join(c.baseDir, relPath) + + if c.pool != nil { + return c.storeFileDB(ctx, projectID, userID, fileName, fileCategory, mimeType, hashStr, relPath, fullPath, sizeBytes, data) + } + + return c.storeFileMemory(projectID, userID, fileName, fileCategory, mimeType, hashStr, relPath, fullPath, sizeBytes, data) +} + +func (c *CAS) storeFileDB(ctx context.Context, projectID, userID int64, fileName, fileCategory, mimeType, hashStr, relPath, fullPath string, sizeBytes int64, data []byte) (*FileReference, error) { + tx, err := c.pool.Begin(ctx) + if err != nil { + return nil, fmt.Errorf("begin tx failed: %w", err) + } + defer tx.Rollback(ctx) + + var existingHash string + var refCount int + err = tx.QueryRow(ctx, "SELECT sha256, ref_count FROM storage_blobs WHERE sha256 = $1 FOR UPDATE", hashStr).Scan(&existingHash, &refCount) + + isDedup := false + if err == nil && existingHash != "" { + // 命中去重:物理文件已存在,递增引用计数 + isDedup = true + _, err = tx.Exec(ctx, "UPDATE storage_blobs SET ref_count = ref_count + 1 WHERE sha256 = $1", hashStr) + if err != nil { + return nil, fmt.Errorf("update ref_count failed: %w", err) + } + } else { + // 新文件:写入物理磁盘 (采用原子写入) + dir := filepath.Dir(fullPath) + if err := os.MkdirAll(dir, 0755); err != nil { + return nil, fmt.Errorf("mkdir failed: %w", err) + } + tmpPath := fullPath + ".tmp" + if err := os.WriteFile(tmpPath, data, 0644); err != nil { + return nil, fmt.Errorf("write tmp file failed: %w", err) + } + if err := os.Rename(tmpPath, fullPath); err != nil { + return nil, fmt.Errorf("rename file failed: %w", err) + } + + _, err = tx.Exec(ctx, + "INSERT INTO storage_blobs (sha256, size_bytes, mime_type, storage_path, ref_count, created_at) VALUES ($1, $2, $3, $4, 1, now())", + hashStr, sizeBytes, mimeType, relPath) + if err != nil { + return nil, fmt.Errorf("insert blob failed: %w", err) + } + } + + // 插入引用记录 + var refID int64 + var createdAt time.Time + err = tx.QueryRow(ctx, ` + INSERT INTO file_references (project_id, user_id, blob_sha256, file_name, file_category, summary, metadata, created_at) + VALUES ($1, $2, $3, $4, $5, '', '{}'::jsonb, now()) + RETURNING id, created_at + `, projectID, userID, hashStr, fileName, fileCategory).Scan(&refID, &createdAt) + if err != nil { + return nil, fmt.Errorf("insert file reference failed: %w", err) + } + + if err := tx.Commit(ctx); err != nil { + return nil, fmt.Errorf("commit tx failed: %w", err) + } + + return &FileReference{ + ID: refID, + ProjectID: projectID, + UserID: userID, + BlobSHA256: hashStr, + FileName: fileName, + FileCategory: fileCategory, + CreatedAt: createdAt, + IsDeduplicated: isDedup, + }, nil +} + +func (c *CAS) storeFileMemory(projectID, userID int64, fileName, fileCategory, mimeType, hashStr, relPath, fullPath string, sizeBytes int64, data []byte) (*FileReference, error) { + c.mu.Lock() + defer c.mu.Unlock() + + isDedup := false + blob, exists := c.memoryBlobs[hashStr] + if exists { + isDedup = true + blob.RefCount++ + } else { + dir := filepath.Dir(fullPath) + if err := os.MkdirAll(dir, 0755); err != nil { + return nil, fmt.Errorf("mkdir failed: %w", err) + } + if err := os.WriteFile(fullPath, data, 0644); err != nil { + return nil, fmt.Errorf("write file failed: %w", err) + } + c.memoryBlobs[hashStr] = &BlobInfo{ + SHA256: hashStr, + SizeBytes: sizeBytes, + MimeType: mimeType, + StoragePath: relPath, + RefCount: 1, + CreatedAt: time.Now(), + } + } + + refID := c.nextRefID + c.nextRefID++ + + ref := &FileReference{ + ID: refID, + ProjectID: projectID, + UserID: userID, + BlobSHA256: hashStr, + FileName: fileName, + FileCategory: fileCategory, + CreatedAt: time.Now(), + IsDeduplicated: isDedup, + } + c.memoryRefs[refID] = ref + return ref, nil +} + +// OpenFileByRef 根据文件引用 ID 打开物理文件流 +func (c *CAS) OpenFileByRef(ctx context.Context, refID int64) (io.ReadCloser, *FileReference, error) { + if c.pool != nil { + var ref FileReference + var storagePath string + err := c.pool.QueryRow(ctx, ` + SELECT r.id, r.project_id, r.user_id, r.blob_sha256, r.file_name, r.file_category, r.summary, r.created_at, b.storage_path + FROM file_references r + JOIN storage_blobs b ON r.blob_sha256 = b.sha256 + WHERE r.id = $1 + `, refID).Scan(&ref.ID, &ref.ProjectID, &ref.UserID, &ref.BlobSHA256, &ref.FileName, &ref.FileCategory, &ref.Summary, &ref.CreatedAt, &storagePath) + if err != nil { + return nil, nil, fmt.Errorf("file reference not found: %w", err) + } + + fullPath := filepath.Join(c.baseDir, storagePath) + f, err := os.Open(fullPath) + if err != nil { + return nil, nil, fmt.Errorf("open physical file failed: %w", err) + } + return f, &ref, nil + } + + c.mu.RLock() + defer c.mu.RUnlock() + ref, ok := c.memoryRefs[refID] + if !ok { + return nil, nil, errors.New("file reference not found in memory") + } + blob, ok := c.memoryBlobs[ref.BlobSHA256] + if !ok { + return nil, nil, errors.New("blob not found in memory") + } + fullPath := filepath.Join(c.baseDir, blob.StoragePath) + f, err := os.Open(fullPath) + if err != nil { + return nil, nil, fmt.Errorf("open physical file failed: %w", err) + } + return f, ref, nil +} + +// StoreMemoryChunk 存储记忆切片并自动去重向量化计算 +func (c *CAS) StoreMemoryChunk(ctx context.Context, in MemoryChunkInput) (*MemoryChunk, error) { + if strings.TrimSpace(in.Content) == "" { + return nil, errors.New("chunk content cannot be empty") + } + if in.ImportanceWeight <= 0 { + in.ImportanceWeight = 1.00 + } + + normalized := strings.TrimSpace(in.Title) + "\n" + strings.TrimSpace(in.Content) + h := sha256.Sum256([]byte(normalized)) + chunkHash := hex.EncodeToString(h[:]) + + var embedding []float32 + + // 1. 检查是否存在相同切片 Hash,复用向量结果以避免重复调用大模型 API 与浪费存储 + if c.pool != nil { + var vecStr *string + err := c.pool.QueryRow(ctx, "SELECT embedding::text FROM context_memory_chunks WHERE chunk_hash = $1 AND embedding IS NOT NULL LIMIT 1", chunkHash).Scan(&vecStr) + if err == nil && vecStr != nil { + embedding = parseVectorText(*vecStr) + } + } + + // 2. 若未命中现有切片,且配置了 Embedder,则调用生成向量 + if len(embedding) == 0 && c.embedder != nil { + if vec, err := c.embedder.EmbedOne(ctx, in.Content); err == nil && len(vec) > 0 { + embedding = vec + } + } + + metaJSON, _ := json.Marshal(in.Metadata) + if len(metaJSON) == 0 { + metaJSON = []byte("{}") + } + + if c.pool != nil { + var chunkID int64 + var createdAt time.Time + + var vecLiteral any + if len(embedding) > 0 { + vecLiteral = vectorLiteral(embedding) + } + + err := c.pool.QueryRow(ctx, ` + INSERT INTO context_memory_chunks + (domain, project_id, user_id, reference_id, title, content, chunk_hash, embedding, tsv, importance_weight, metadata, created_at) + VALUES + ($1, $2, $3, $4, $5, $6, $7, $8::vector, to_tsvector('simple', $5 || ' ' || $6), $9, $10, now()) + RETURNING id, created_at + `, in.Domain, in.ProjectID, in.UserID, in.ReferenceID, in.Title, in.Content, chunkHash, vecLiteral, in.ImportanceWeight, metaJSON).Scan(&chunkID, &createdAt) + + if err != nil { + return nil, fmt.Errorf("insert context memory chunk failed: %w", err) + } + + return &MemoryChunk{ + ID: chunkID, + Domain: in.Domain, + ProjectID: in.ProjectID, + UserID: in.UserID, + ReferenceID: in.ReferenceID, + Title: in.Title, + Content: in.Content, + ChunkHash: chunkHash, + Embedding: embedding, + ImportanceWeight: in.ImportanceWeight, + Metadata: in.Metadata, + CreatedAt: createdAt, + }, nil + } + + // 降级内存存储 + c.mu.Lock() + defer c.mu.Unlock() + chunkID := c.nextChunkID + c.nextChunkID++ + chunk := &MemoryChunk{ + ID: chunkID, + Domain: in.Domain, + ProjectID: in.ProjectID, + UserID: in.UserID, + ReferenceID: in.ReferenceID, + Title: in.Title, + Content: in.Content, + ChunkHash: chunkHash, + Embedding: embedding, + ImportanceWeight: in.ImportanceWeight, + Metadata: in.Metadata, + CreatedAt: time.Now(), + } + c.memoryChunks = append(c.memoryChunks, chunk) + return chunk, nil +} + +// SearchMemory 统一多维度记忆语义检索 (向量相似度 + 全文检索 + 时间衰减权重) +func (c *CAS) SearchMemory(ctx context.Context, opts MemorySearchOptions) ([]*MemoryChunkHit, error) { + if opts.TopK <= 0 { + opts.TopK = 5 + } + if opts.DecayHalfLifeDays <= 0 { + opts.DecayHalfLifeDays = 7.0 // 默认 7 天半衰期 + } + + if c.pool != nil { + return c.searchMemoryDB(ctx, opts) + } + + return c.searchMemoryFallback(opts), nil +} + +func (c *CAS) searchMemoryDB(ctx context.Context, opts MemorySearchOptions) ([]*MemoryChunkHit, error) { + var queryVector []float32 + if len(opts.QueryVector) > 0 { + queryVector = opts.QueryVector + } else if c.embedder != nil && strings.TrimSpace(opts.QueryText) != "" { + if vec, err := c.embedder.EmbedOne(ctx, opts.QueryText); err == nil { + queryVector = vec + } + } + + var sb strings.Builder + var args []any + argIdx := 1 + + sb.WriteString(` + SELECT id, domain, project_id, user_id, reference_id, title, content, chunk_hash, importance_weight, metadata, created_at, + `) + + if len(queryVector) > 0 { + args = append(args, vectorLiteral(queryVector)) + sb.WriteString(fmt.Sprintf(" (1 - (embedding <=> $%d::vector)) AS raw_score ", argIdx)) + argIdx++ + } else { + sb.WriteString(" 1.0 AS raw_score ") + } + + sb.WriteString(" FROM context_memory_chunks WHERE 1=1 ") + + if opts.Domain != "" { + sb.WriteString(fmt.Sprintf(" AND domain = $%d ", argIdx)) + args = append(args, opts.Domain) + argIdx++ + } + if opts.ProjectID > 0 { + sb.WriteString(fmt.Sprintf(" AND project_id = $%d ", argIdx)) + args = append(args, opts.ProjectID) + argIdx++ + } + if opts.UserID > 0 { + sb.WriteString(fmt.Sprintf(" AND user_id = $%d ", argIdx)) + args = append(args, opts.UserID) + argIdx++ + } + if strings.TrimSpace(opts.QueryText) != "" { + if len(queryVector) > 0 { + sb.WriteString(fmt.Sprintf(" AND (embedding IS NOT NULL OR tsv @@ plainto_tsquery('simple', $%d) OR title ILIKE '%%' || $%d || '%%' OR content ILIKE '%%' || $%d || '%%') ", argIdx, argIdx, argIdx)) + } else { + sb.WriteString(fmt.Sprintf(" AND (tsv @@ plainto_tsquery('simple', $%d) OR title ILIKE '%%' || $%d || '%%' OR content ILIKE '%%' || $%d || '%%') ", argIdx, argIdx, argIdx)) + } + args = append(args, opts.QueryText) + argIdx++ + } + + sb.WriteString(" ORDER BY created_at DESC LIMIT 50 ") + + rows, err := c.pool.Query(ctx, sb.String(), args...) + if err != nil { + return nil, fmt.Errorf("search memory failed: %w", err) + } + defer rows.Close() + + var hits []*MemoryChunkHit + lambda := math.Ln2 / opts.DecayHalfLifeDays + + for rows.Next() { + var chunk MemoryChunk + var metaJSON []byte + var rawScore float64 + + err := rows.Scan( + &chunk.ID, &chunk.Domain, &chunk.ProjectID, &chunk.UserID, &chunk.ReferenceID, + &chunk.Title, &chunk.Content, &chunk.ChunkHash, &chunk.ImportanceWeight, + &metaJSON, &chunk.CreatedAt, &rawScore, + ) + if err != nil { + continue + } + _ = json.Unmarshal(metaJSON, &chunk.Metadata) + + timeDecay := 1.0 + if opts.ApplyTimeDecay { + daysDiff := time.Since(chunk.CreatedAt).Hours() / 24.0 + if daysDiff > 0 { + timeDecay = math.Exp(-lambda * daysDiff) + } + } + + finalScore := rawScore * timeDecay * chunk.ImportanceWeight + hits = append(hits, &MemoryChunkHit{ + Chunk: &chunk, + RawScore: rawScore, + TimeDecay: timeDecay, + FinalScore: finalScore, + }) + } + + // 按 FinalScore 降序排列 + sortChunkHits(hits) + if len(hits) > opts.TopK { + hits = hits[:opts.TopK] + } + + return hits, nil +} + +func (c *CAS) searchMemoryFallback(opts MemorySearchOptions) []*MemoryChunkHit { + c.mu.RLock() + defer c.mu.RUnlock() + + var hits []*MemoryChunkHit + lambda := math.Ln2 / opts.DecayHalfLifeDays + + for _, chunk := range c.memoryChunks { + if opts.Domain != "" && chunk.Domain != opts.Domain { + continue + } + if opts.ProjectID > 0 && chunk.ProjectID != opts.ProjectID { + continue + } + if opts.UserID > 0 && chunk.UserID != opts.UserID { + continue + } + + rawScore := 1.0 + if opts.QueryText != "" { + if !strings.Contains(chunk.Title, opts.QueryText) && !strings.Contains(chunk.Content, opts.QueryText) { + continue + } + } + + timeDecay := 1.0 + if opts.ApplyTimeDecay { + daysDiff := time.Since(chunk.CreatedAt).Hours() / 24.0 + if daysDiff > 0 { + timeDecay = math.Exp(-lambda * daysDiff) + } + } + + finalScore := rawScore * timeDecay * chunk.ImportanceWeight + hits = append(hits, &MemoryChunkHit{ + Chunk: chunk, + RawScore: rawScore, + TimeDecay: timeDecay, + FinalScore: finalScore, + }) + } + + sortChunkHits(hits) + if len(hits) > opts.TopK { + hits = hits[:opts.TopK] + } + return hits +} + +func sortChunkHits(hits []*MemoryChunkHit) { + for i := 0; i < len(hits); i++ { + for j := i + 1; j < len(hits); j++ { + if hits[j].FinalScore > hits[i].FinalScore { + hits[i], hits[j] = hits[j], hits[i] + } + } + } +} + +func vectorLiteral(v []float32) string { + parts := make([]string, len(v)) + for i, x := range v { + parts[i] = fmt.Sprintf("%.6f", x) + } + return "[" + strings.Join(parts, ",") + "]" +} + +func parseVectorText(s string) []float32 { + s = strings.Trim(s, "[] ") + if s == "" { + return nil + } + parts := strings.Split(s, ",") + res := make([]float32, 0, len(parts)) + for _, p := range parts { + var f float32 + if _, err := fmt.Sscanf(strings.TrimSpace(p), "%f", &f); err == nil { + res = append(res, f) + } + } + return res +} diff --git a/internal/storage/cas_db_test.go b/internal/storage/cas_db_test.go new file mode 100644 index 00000000..f19a58ab --- /dev/null +++ b/internal/storage/cas_db_test.go @@ -0,0 +1,123 @@ +package storage + +import ( + "context" + "os" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" +) + +func TestCAS_DatabaseIntegration(t *testing.T) { + dsn := os.Getenv("CONTEXT_DB_DSN") + if dsn == "" { + dsn = "postgres://app:dev_secret_db_pass@127.0.0.1:5433/agent_system_dev?sslmode=disable" + } + + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + + pool, err := pgxpool.New(ctx, dsn) + if err != nil { + t.Skipf("cannot connect to context db (%v), skipping DB integration test", err) + return + } + defer pool.Close() + + if err := pool.Ping(ctx); err != nil { + t.Skipf("context db ping failed (%v), skipping DB integration test", err) + return + } + + tempDir, err := os.MkdirTemp("", "cas_db_test_*") + if err != nil { + t.Fatalf("mkdir temp failed: %v", err) + } + defer os.RemoveAll(tempDir) + + mockEmb := &mockEmbedder{} + cas, err := NewCAS(tempDir, pool, mockEmb) + if err != nil { + t.Fatalf("new cas failed: %v", err) + } + + // 准备测试项目与用户 + var projID int64 + projCode := "TEST_PROJ_" + time.Now().Format("20060102150405") + err = pool.QueryRow(ctx, ` + INSERT INTO projects (code, name, customer_name, device_type) + VALUES ($1, '测试内窥镜系统', '北京水木东方', '内窥镜') + ON CONFLICT (code) DO UPDATE SET name = EXCLUDED.name + RETURNING id + `, projCode).Scan(&projID) + if err != nil { + t.Fatalf("insert project failed: %v", err) + } + + var userID int64 + accountID := "user_test_" + time.Now().Format("150405") + err = pool.QueryRow(ctx, ` + INSERT INTO users (account_id, name, role) + VALUES ($1, '测试工程师', 'engineer') + ON CONFLICT (account_id) DO UPDATE SET name = EXCLUDED.name + RETURNING id + `, accountID).Scan(&userID) + if err != nil { + t.Fatalf("insert user failed: %v", err) + } + + // 1. 上传文件 1 (使用动态内容确保本次测试为全新文件) + testData := []byte("这是写入真实数据库并在物理磁盘去重的测试申报文件内容: " + time.Now().Format(time.RFC3339Nano)) + ref1, err := cas.StoreFile(ctx, projID, userID, "spec_v1.pdf", "standard", testData, "application/pdf") + if err != nil { + t.Fatalf("store file 1 to db failed: %v", err) + } + if ref1.IsDeduplicated { + t.Errorf("expected first upload to not be deduplicated") + } + + // 2. 上传文件 2 (内容相同,触发 DB 去重) + ref2, err := cas.StoreFile(ctx, projID, userID, "spec_v2_dup.pdf", "standard", testData, "application/pdf") + if err != nil { + t.Fatalf("store file 2 to db failed: %v", err) + } + if !ref2.IsDeduplicated { + t.Errorf("expected second upload to be deduplicated in DB") + } + + // 3. 存储并检索向量切片 + chunk, err := cas.StoreMemoryChunk(ctx, MemoryChunkInput{ + Domain: "project_doc", + ProjectID: projID, + UserID: userID, + Title: "光机指标要求", + Content: "内窥镜视场角不低于 100 度,光学畸变应严格控制在 5% 以内。", + ImportanceWeight: 1.2, + }) + if err != nil { + t.Fatalf("store memory chunk to db failed: %v", err) + } + if chunk.ID <= 0 { + t.Errorf("expected positive chunk ID from db") + } + + // 4. 检索向量切片 + hits, err := cas.SearchMemory(ctx, MemorySearchOptions{ + Domain: "project_doc", + ProjectID: projID, + QueryText: "视场角", + ApplyTimeDecay: true, + DecayHalfLifeDays: 7.0, + TopK: 5, + }) + if err != nil { + t.Fatalf("search memory in db failed: %v", err) + } + if len(hits) == 0 { + t.Fatalf("expected memory hits from db, got 0") + } + if hits[0].Chunk.Title != "光机指标要求" { + t.Errorf("unexpected top hit title from db: %s", hits[0].Chunk.Title) + } +} diff --git a/internal/storage/cas_test.go b/internal/storage/cas_test.go new file mode 100644 index 00000000..79e63423 --- /dev/null +++ b/internal/storage/cas_test.go @@ -0,0 +1,203 @@ +package storage + +import ( + "context" + "io" + "os" + "path/filepath" + "testing" + "time" +) + +type mockEmbedder struct { + calledCount int +} + +func (m *mockEmbedder) EmbedOne(ctx context.Context, text string) ([]float32, error) { + m.calledCount++ + // 返回一个简单的测试向量 + vec := make([]float32, 1024) + for i := range vec { + vec[i] = 0.01 * float32(i%10) + } + return vec, nil +} + +func TestCAS_FileDeduplication(t *testing.T) { + tempDir, err := os.MkdirTemp("", "cas_test_*") + if err != nil { + t.Fatalf("mkdir temp failed: %v", err) + } + defer os.RemoveAll(tempDir) + + mockEmb := &mockEmbedder{} + cas, err := NewCAS(tempDir, nil, mockEmb) + if err != nil { + t.Fatalf("new cas failed: %v", err) + } + + ctx := context.Background() + fileContent := []byte("这是一份用于去重测试的医疗器械说明书与注册申报资料内容,包含重复章节条款。") + + // 1. 用户 1 首次上传 + ref1, err := cas.StoreFile(ctx, 101, 1, "instruction_v1.pdf", "material", fileContent, "application/pdf") + if err != nil { + t.Fatalf("store file 1 failed: %v", err) + } + if ref1.IsDeduplicated { + t.Errorf("expected first upload to not be deduplicated") + } + + // 2. 用户 2 再次上传完全相同内容的文件 (不同文件名、不同用户) + ref2, err := cas.StoreFile(ctx, 101, 2, "instruction_copy.pdf", "material", fileContent, "application/pdf") + if err != nil { + t.Fatalf("store file 2 failed: %v", err) + } + if !ref2.IsDeduplicated { + t.Errorf("expected second upload to be deduplicated (CAS Hit)") + } + if ref1.BlobSHA256 != ref2.BlobSHA256 { + t.Errorf("expected identical SHA256 hashes, got %s vs %s", ref1.BlobSHA256, ref2.BlobSHA256) + } + + // 3. 读取验证文件完整性 + r1, _, err := cas.OpenFileByRef(ctx, ref1.ID) + if err != nil { + t.Fatalf("open file 1 failed: %v", err) + } + defer r1.Close() + data1, _ := io.ReadAll(r1) + if string(data1) != string(fileContent) { + t.Errorf("file 1 content mismatch") + } + + r2, _, err := cas.OpenFileByRef(ctx, ref2.ID) + if err != nil { + t.Fatalf("open file 2 failed: %v", err) + } + defer r2.Close() + data2, _ := io.ReadAll(r2) + if string(data2) != string(fileContent) { + t.Errorf("file 2 content mismatch") + } + + // 4. 检查物理磁盘上是否仅有 1 个物理 blob 文件 + blobsDir := filepath.Join(tempDir, "blobs") + blobCount := 0 + _ = filepath.Walk(blobsDir, func(path string, info os.FileInfo, err error) error { + if err == nil && !info.IsDir() && filepath.Ext(path) == ".bin" { + blobCount++ + } + return nil + }) + if blobCount != 1 { + t.Errorf("expected exactly 1 physical blob on disk, found %d", blobCount) + } +} + +func TestCAS_MemoryChunkAndDecay(t *testing.T) { + tempDir, err := os.MkdirTemp("", "cas_chunk_test_*") + if err != nil { + t.Fatalf("mkdir temp failed: %v", err) + } + defer os.RemoveAll(tempDir) + + mockEmb := &mockEmbedder{} + cas, err := NewCAS(tempDir, nil, mockEmb) + if err != nil { + t.Fatalf("new cas failed: %v", err) + } + + ctx := context.Background() + + // 1. 存储切片 + c1, err := cas.StoreMemoryChunk(ctx, MemoryChunkInput{ + Domain: "user_memory", + ProjectID: 1001, + UserID: 1, + Title: "客户对工期极度敏感", + Content: "客户希望首期在 4 周内完成样品摸底,并强调需优先保证安规测试通过。", + ImportanceWeight: 1.5, + }) + if err != nil { + t.Fatalf("store chunk 1 failed: %v", err) + } + if c1.ChunkHash == "" { + t.Errorf("expected non-empty chunk hash") + } + + // 2. 模拟检索 + hits, err := cas.SearchMemory(ctx, MemorySearchOptions{ + Domain: "user_memory", + ProjectID: 1001, + QueryText: "工期", + ApplyTimeDecay: true, + DecayHalfLifeDays: 7.0, + TopK: 5, + }) + if err != nil { + t.Fatalf("search memory failed: %v", err) + } + if len(hits) == 0 { + t.Fatalf("expected memory hits, got 0") + } + if hits[0].Chunk.Title != "客户对工期极度敏感" { + t.Errorf("unexpected top hit title: %s", hits[0].Chunk.Title) + } + if hits[0].FinalScore <= 0 { + t.Errorf("expected positive final score") + } +} + +func TestCAS_TimeDecayCalculation(t *testing.T) { + tempDir, err := os.MkdirTemp("", "cas_decay_test_*") + if err != nil { + t.Fatalf("mkdir temp failed: %v", err) + } + defer os.RemoveAll(tempDir) + + cas, _ := NewCAS(tempDir, nil, nil) + ctx := context.Background() + + // 写入一条即时记忆和一条较早的记忆 + nowChunk, _ := cas.StoreMemoryChunk(ctx, MemoryChunkInput{ + Domain: "project_activity", + ProjectID: 2001, + Title: "今日决策", + Content: "今日锁定目标市场为中国及欧盟CE。", + ImportanceWeight: 1.0, + }) + + pastChunk, _ := cas.StoreMemoryChunk(ctx, MemoryChunkInput{ + Domain: "project_activity", + ProjectID: 2001, + Title: "早期设想", + Content: "早期曾考虑走美国FDA 510(k)途径,后续搁置。", + ImportanceWeight: 1.0, + }) + + // 人工修改 pastChunk 的创建时间为 14 天前 (两个半衰期) + pastChunk.CreatedAt = time.Now().Add(-14 * 24 * time.Hour) + + hits, err := cas.SearchMemory(ctx, MemorySearchOptions{ + Domain: "project_activity", + ProjectID: 2001, + ApplyTimeDecay: true, + DecayHalfLifeDays: 7.0, + TopK: 5, + }) + if err != nil { + t.Fatalf("search memory failed: %v", err) + } + if len(hits) < 2 { + t.Fatalf("expected 2 hits, got %d", len(hits)) + } + + // 即时记忆得分应远高于 14 天前的早期设想 + if hits[0].Chunk.ID != nowChunk.ID { + t.Errorf("expected nowChunk to rank higher due to time decay, got ID %d", hits[0].Chunk.ID) + } + if hits[1].TimeDecay >= 0.5 { + t.Errorf("expected 14-day old chunk to decay below 0.5, got %.3f", hits[1].TimeDecay) + } +} diff --git a/internal/storage/types.go b/internal/storage/types.go new file mode 100644 index 00000000..1a8fa1e1 --- /dev/null +++ b/internal/storage/types.go @@ -0,0 +1,77 @@ +package storage + +import ( + "time" +) + +// BlobInfo CAS 内容寻址物理对象元数据 +type BlobInfo struct { + SHA256 string `json:"sha256"` + SizeBytes int64 `json:"size_bytes"` + MimeType string `json:"mime_type"` + StoragePath string `json:"storage_path"` + RefCount int `json:"ref_count"` + CreatedAt time.Time `json:"created_at"` +} + +// FileReference 用户或项目对 CAS 物理对象的文件引用 +type FileReference struct { + ID int64 `json:"id"` + ProjectID int64 `json:"project_id"` + UserID int64 `json:"user_id"` + BlobSHA256 string `json:"blob_sha256"` + FileName string `json:"file_name"` + FileCategory string `json:"file_category"` // material, standard, cad, quote, deliverable + Summary string `json:"summary"` + Metadata map[string]any `json:"metadata,omitempty"` + CreatedAt time.Time `json:"created_at"` + IsDeduplicated bool `json:"is_deduplicated"` // 是否命中去重(未占用新增物理磁盘空间) +} + +// MemoryChunk 动态上下文记忆切片 +type MemoryChunk struct { + ID int64 `json:"id"` + Domain string `json:"domain"` // user_memory, project_doc, project_activity, chat_summary + ProjectID int64 `json:"project_id"` + UserID int64 `json:"user_id"` + ReferenceID string `json:"reference_id"` + Title string `json:"title"` + Content string `json:"content"` + ChunkHash string `json:"chunk_hash"` + Embedding []float32 `json:"embedding,omitempty"` + ImportanceWeight float64 `json:"importance_weight"` + Metadata map[string]any `json:"metadata,omitempty"` + CreatedAt time.Time `json:"created_at"` +} + +// MemoryChunkInput 记忆切片入库输入 +type MemoryChunkInput struct { + Domain string `json:"domain"` + ProjectID int64 `json:"project_id"` + UserID int64 `json:"user_id"` + ReferenceID string `json:"reference_id"` + Title string `json:"title"` + Content string `json:"content"` + ImportanceWeight float64 `json:"importance_weight"` + Metadata map[string]any `json:"metadata,omitempty"` +} + +// MemorySearchOptions 记忆检索选项 +type MemorySearchOptions struct { + Domain string `json:"domain,omitempty"` + ProjectID int64 `json:"project_id,omitempty"` + UserID int64 `json:"user_id,omitempty"` + QueryText string `json:"query_text"` + QueryVector []float32 `json:"query_vector,omitempty"` + TopK int `json:"top_k"` + ApplyTimeDecay bool `json:"apply_time_decay"` + DecayHalfLifeDays float64 `json:"decay_half_life_days"` // 半衰期天数 (如 7.0 天) +} + +// MemoryChunkHit 记忆检索命中结果 +type MemoryChunkHit struct { + Chunk *MemoryChunk `json:"chunk"` + RawScore float64 `json:"raw_score"` + TimeDecay float64 `json:"time_decay"` + FinalScore float64 `json:"final_score"` +}