72 lines
1.4 KiB
Go
72 lines
1.4 KiB
Go
package memory
|
|
|
|
import (
|
|
"context"
|
|
"log/slog"
|
|
"sync"
|
|
"time"
|
|
|
|
"git.db123.ir/db123/divacode/internal/storage"
|
|
)
|
|
|
|
type Store struct {
|
|
mu sync.Mutex
|
|
db *storage.DB
|
|
}
|
|
|
|
func NewStore(db *storage.DB) *Store {
|
|
return &Store{db: db}
|
|
}
|
|
|
|
func (s *Store) Store(ctx context.Context, content string, source string) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
importance := 0.5
|
|
if len(content) > 100 {
|
|
importance = 0.7
|
|
}
|
|
|
|
id := time.Now().Format("150405.000000000")
|
|
return s.db.Exec(
|
|
"INSERT INTO memories (id, content, importance, source, created_at) VALUES (?, ?, ?, ?, datetime('now'))",
|
|
id, content, importance, source,
|
|
)
|
|
}
|
|
|
|
func (s *Store) Search(ctx context.Context, query string, topK int) ([]string, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
rows, err := s.db.Query(
|
|
`SELECT content FROM memories
|
|
WHERE content LIKE ?
|
|
ORDER BY importance DESC, created_at DESC
|
|
LIMIT ?`,
|
|
"%"+query+"%", topK,
|
|
)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
var results []string
|
|
for rows.Next() {
|
|
var content string
|
|
if err := rows.Scan(&content); err != nil {
|
|
continue
|
|
}
|
|
results = append(results, content)
|
|
}
|
|
|
|
slog.Debug("memory search", "query", query, "results", len(results))
|
|
return results, nil
|
|
}
|
|
|
|
func (s *Store) UpdateImportance(ctx context.Context, id string, importance float64) error {
|
|
return s.db.Exec(
|
|
"UPDATE memories SET importance = ?, last_accessed = datetime('now') WHERE id = ?",
|
|
importance, id,
|
|
)
|
|
}
|