323 lines
8.8 KiB
Go
323 lines
8.8 KiB
Go
//go:build !cgo
|
||
|
||
// 织忆 MemoryWeave — 纯 Go SQLite + 内存向量混合存储（Windows 持久化方案）
|
||
// 使用 modernc.org/sqlite（纯 Go 实现，无需 CGO）
|
||
// 向量搜索走内存（与 MemLanceClient 相同），元数据持久化到 SQLite
|
||
package storage
|
||
|
||
import (
|
||
"database/sql"
|
||
"encoding/json"
|
||
"fmt"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/xiaoxue/memoryweave/internal/models"
|
||
|
||
_ "modernc.org/sqlite"
|
||
)
|
||
|
||
// SQLiteMemClient 纯 Go SQLite + 内存向量混合存储
|
||
// 适用于 Windows 或无 CGO 环境
|
||
type SQLiteMemClient struct {
|
||
*MemLanceClient // 向量搜索走内存实现
|
||
db *sql.DB
|
||
dbPath string
|
||
pendingDirty bool
|
||
dirtyMu sync.Mutex
|
||
}
|
||
|
||
// NewSQLiteMemClient 创建持久化存储客户端
|
||
func NewSQLiteMemClient(dbPath string, embedder *Embedder) (*SQLiteMemClient, error) {
|
||
if dbPath == "" {
|
||
dbPath = "C:\\Users\\Administrator\\.zhiyi\\zhiyi.db"
|
||
}
|
||
|
||
db, err := sql.Open("sqlite3", dbPath+"?_journal_mode=WAL&_busy_timeout=5000")
|
||
if err != nil {
|
||
return nil, fmt.Errorf("sqlite open: %w", err)
|
||
}
|
||
|
||
sm := &SQLiteMemClient{
|
||
MemLanceClient: NewMemLanceClient(embedder),
|
||
db: db,
|
||
dbPath: dbPath,
|
||
}
|
||
if err := sm.migrate(); err != nil {
|
||
db.Close()
|
||
return nil, fmt.Errorf("migrate: %w", err)
|
||
}
|
||
if err := sm.load(); err != nil {
|
||
// 只警告，不失败——内存空也能跑
|
||
fmt.Printf("[zhiyi] warn: load from sqlite failed: %v (starting fresh)\n", err)
|
||
}
|
||
return sm, nil
|
||
}
|
||
|
||
func (sm *SQLiteMemClient) migrate() error {
|
||
sqls := []string{
|
||
`CREATE TABLE IF NOT EXISTS memories (
|
||
id TEXT PRIMARY KEY,
|
||
content TEXT NOT NULL,
|
||
agent_id TEXT NOT NULL DEFAULT '',
|
||
namespace TEXT NOT NULL DEFAULT '',
|
||
category TEXT NOT NULL DEFAULT '',
|
||
tier TEXT NOT NULL DEFAULT '',
|
||
quality_score REAL NOT NULL DEFAULT 0,
|
||
vector BLOB,
|
||
is_deleted INTEGER NOT NULL DEFAULT 0,
|
||
version INTEGER NOT NULL DEFAULT 1,
|
||
version_history TEXT,
|
||
source TEXT NOT NULL DEFAULT '',
|
||
created_at TEXT NOT NULL,
|
||
last_recalled_at TEXT,
|
||
recall_count INTEGER NOT NULL DEFAULT 0,
|
||
useful_count INTEGER NOT NULL DEFAULT 0,
|
||
not_useful_count INTEGER NOT NULL DEFAULT 0
|
||
)`,
|
||
`CREATE TABLE IF NOT EXISTS episodes (
|
||
id TEXT PRIMARY KEY,
|
||
agent_id TEXT NOT NULL,
|
||
namespace TEXT NOT NULL DEFAULT '',
|
||
content TEXT NOT NULL,
|
||
category TEXT NOT NULL DEFAULT '',
|
||
created_at TEXT NOT NULL
|
||
)`,
|
||
`CREATE INDEX IF NOT EXISTS idx_memories_agent ON memories(agent_id)`,
|
||
`CREATE INDEX IF NOT EXISTS idx_memories_namespace ON memories(namespace)`,
|
||
`CREATE INDEX IF NOT EXISTS idx_memories_tier ON memories(tier)`,
|
||
}
|
||
for _, s := range sqls {
|
||
if _, err := sm.db.Exec(s); err != nil {
|
||
return fmt.Errorf("exec(%s): %w", s[:30], err)
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// load 从 SQLite 恢复到内存
|
||
func (sm *SQLiteMemClient) load() error {
|
||
// 加载 memories
|
||
rows, err := sm.db.Query(`SELECT id,content,agent_id,namespace,category,tier,
|
||
quality_score,vector,is_deleted,version,version_history,source,
|
||
created_at,last_recalled_at,recall_count,useful_count,not_useful_count
|
||
FROM memories`)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
defer rows.Close()
|
||
|
||
sm.mu.Lock()
|
||
defer sm.mu.Unlock()
|
||
|
||
for rows.Next() {
|
||
var id, content, agentID, namespace, category, tier, source string
|
||
var qualityScore float64
|
||
var vectorJSON []byte
|
||
var isDeleted bool
|
||
var version, recallCount, usefulCount, notUsefulCount int
|
||
var versionHistoryJSON sql.NullString
|
||
var createdAtStr, lastRecalledAtStr sql.NullString
|
||
|
||
if err := rows.Scan(&id, &content, &agentID, &namespace, &category, &tier,
|
||
&qualityScore, &vectorJSON, &isDeleted, &version, &versionHistoryJSON, &source,
|
||
&createdAtStr, &lastRecalledAtStr, &recallCount, &usefulCount, ¬UsefulCount); err != nil {
|
||
continue
|
||
}
|
||
|
||
vec := make([]float32, 1024)
|
||
if len(vectorJSON) > 0 {
|
||
vec64 := make([]float64, 1024)
|
||
json.Unmarshal(vectorJSON, &vec64)
|
||
for i := range vec64 {
|
||
vec[i] = float32(vec64[i])
|
||
}
|
||
}
|
||
|
||
createdAt := time.Now()
|
||
if createdAtStr.Valid {
|
||
createdAt, _ = time.Parse(time.RFC3339, createdAtStr.String)
|
||
}
|
||
|
||
var vh []map[string]interface{}
|
||
if versionHistoryJSON.Valid {
|
||
json.Unmarshal([]byte(versionHistoryJSON.String), &vh)
|
||
}
|
||
|
||
sm.memories[id] = &memEntry{
|
||
ID: id,
|
||
Content: content,
|
||
AgentID: agentID,
|
||
Namespace: namespace,
|
||
Category: category,
|
||
Tier: tier,
|
||
QualityScore: qualityScore,
|
||
Vector: vec,
|
||
IsDeleted: isDeleted,
|
||
Version: version,
|
||
VersionHistory: vh,
|
||
Source: source,
|
||
CreatedAt: createdAt,
|
||
RecallCount: recallCount,
|
||
UsefulCount: usefulCount,
|
||
NotUsefulCount: notUsefulCount,
|
||
}
|
||
}
|
||
|
||
// 加载 episodes
|
||
epRows, err := sm.db.Query(`SELECT id,agent_id,namespace,content,category,created_at FROM episodes`)
|
||
if err != nil {
|
||
return nil
|
||
}
|
||
defer epRows.Close()
|
||
|
||
for epRows.Next() {
|
||
var id, agentID, namespace, content, category, createdAtStr string
|
||
if err := epRows.Scan(&id, &agentID, &namespace, &content, &category, &createdAtStr); err != nil {
|
||
continue
|
||
}
|
||
createdAt, _ := time.Parse(time.RFC3339, createdAtStr)
|
||
sm.episodes = append(sm.episodes, models.EpisodeRecord{
|
||
ID: id,
|
||
AgentID: agentID,
|
||
Namespace: namespace,
|
||
Content: content,
|
||
Category: category,
|
||
CreatedAt: createdAt,
|
||
})
|
||
}
|
||
|
||
return nil
|
||
}
|
||
|
||
func float32sToFloat64s(f32 []float32) []float64 {
|
||
f64 := make([]float64, len(f32))
|
||
for i := range f32 {
|
||
f64[i] = float64(f32[i])
|
||
}
|
||
return f64
|
||
}
|
||
|
||
func (sm *SQLiteMemClient) persistMemory(e *memEntry) error {
|
||
vecJSON, _ := json.Marshal(float32sToFloat64s(e.Vector))
|
||
vhJSON, _ := json.Marshal(e.VersionHistory)
|
||
createdAt := e.CreatedAt.Format(time.RFC3339)
|
||
_, err := sm.db.Exec(`INSERT OR REPLACE INTO memories
|
||
(id,content,agent_id,namespace,category,tier,quality_score,vector,
|
||
is_deleted,version,version_history,source,created_at,last_recalled_at,
|
||
recall_count,useful_count,not_useful_count)
|
||
VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`,
|
||
e.ID, e.Content, e.AgentID, e.Namespace, e.Category, e.Tier,
|
||
e.QualityScore, vecJSON, e.IsDeleted, e.Version, vhJSON, e.Source,
|
||
createdAt, e.LastRecalledAt, e.RecallCount, e.UsefulCount, e.NotUsefulCount)
|
||
return err
|
||
}
|
||
|
||
func (sm *SQLiteMemClient) persistEpisode(r models.EpisodeRecord) error {
|
||
_, err := sm.db.Exec(`INSERT OR REPLACE INTO episodes (id,agent_id,namespace,content,category,created_at)
|
||
VALUES (?,?,?,?,?,?)`,
|
||
r.ID, r.AgentID, r.Namespace, r.Content, r.Category, r.CreatedAt.Format(time.RFC3339))
|
||
return err
|
||
}
|
||
|
||
// ─── Override MemLanceClient 方法，追加 SQLite 持久化 ───
|
||
|
||
func (sm *SQLiteMemClient) InsertMemory(m models.MemoryRecord) error {
|
||
// 先走内存逻辑（计算向量、写入内存 map）
|
||
if err := sm.MemLanceClient.InsertMemory(m); err != nil {
|
||
return err
|
||
}
|
||
// 同步写 SQLite
|
||
sm.mu.RLock()
|
||
e := sm.memories[m.ID]
|
||
sm.mu.RUnlock()
|
||
if e != nil {
|
||
return sm.persistMemory(e)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (sm *SQLiteMemClient) InsertEpisode(agentID, namespace, content, category string) (string, error) {
|
||
id, err := sm.MemLanceClient.InsertEpisode(agentID, namespace, content, category)
|
||
if err != nil {
|
||
return id, err
|
||
}
|
||
sm.mu.RLock()
|
||
for i := len(sm.episodes) - 1; i >= 0; i-- {
|
||
if sm.episodes[i].ID == id {
|
||
sm.persistEpisode(sm.episodes[i])
|
||
break
|
||
}
|
||
}
|
||
sm.mu.RUnlock()
|
||
return id, nil
|
||
}
|
||
|
||
func (sm *SQLiteMemClient) SoftDelete(id, reason string) error {
|
||
if err := sm.MemLanceClient.SoftDelete(id, reason); err != nil {
|
||
return err
|
||
}
|
||
sm.mu.RLock()
|
||
e := sm.memories[id]
|
||
sm.mu.RUnlock()
|
||
if e != nil {
|
||
return sm.persistMemory(e)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (sm *SQLiteMemClient) UpdateMemoryContent(id, newContent, source string) error {
|
||
if err := sm.MemLanceClient.UpdateMemoryContent(id, newContent, source); err != nil {
|
||
return err
|
||
}
|
||
sm.mu.RLock()
|
||
e := sm.memories[id]
|
||
sm.mu.RUnlock()
|
||
if e != nil {
|
||
return sm.persistMemory(e)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (sm *SQLiteMemClient) IncrementUseful(id string) {
|
||
sm.MemLanceClient.IncrementUseful(id)
|
||
sm.mu.RLock()
|
||
e := sm.memories[id]
|
||
sm.mu.RUnlock()
|
||
if e != nil {
|
||
sm.persistMemory(e)
|
||
}
|
||
}
|
||
|
||
func (sm *SQLiteMemClient) IncrementNotUseful(id string) {
|
||
sm.MemLanceClient.IncrementNotUseful(id)
|
||
sm.mu.RLock()
|
||
e := sm.memories[id]
|
||
sm.mu.RUnlock()
|
||
if e != nil {
|
||
sm.persistMemory(e)
|
||
}
|
||
}
|
||
|
||
func (sm *SQLiteMemClient) Update(table, id string, fields map[string]any) error {
|
||
if err := sm.MemLanceClient.Update(table, id, fields); err != nil {
|
||
return err
|
||
}
|
||
sm.mu.RLock()
|
||
e := sm.memories[id]
|
||
sm.mu.RUnlock()
|
||
if e != nil {
|
||
return sm.persistMemory(e)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (sm *SQLiteMemClient) Stats() (map[string]interface{}, error) {
|
||
sm.mu.RLock()
|
||
defer sm.mu.RUnlock()
|
||
return map[string]interface{}{
|
||
"memory_count": len(sm.memories),
|
||
"episode_count": len(sm.episodes),
|
||
"backend": "sqlite+memvector (pure Go)",
|
||
}, nil
|
||
} |