144 lines
2.6 KiB
Go
144 lines
2.6 KiB
Go
package session
|
|
|
|
import (
|
|
"encoding/json"
|
|
"os"
|
|
"path/filepath"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/sipeed/picoclaw/pkg/providers"
|
|
)
|
|
|
|
type Session struct {
|
|
Key string `json:"key"`
|
|
Messages []providers.Message `json:"messages"`
|
|
Created time.Time `json:"created"`
|
|
Updated time.Time `json:"updated"`
|
|
}
|
|
|
|
type SessionManager struct {
|
|
sessions map[string]*Session
|
|
mu sync.RWMutex
|
|
storage string
|
|
}
|
|
|
|
func NewSessionManager(storage string) *SessionManager {
|
|
sm := &SessionManager{
|
|
sessions: make(map[string]*Session),
|
|
storage: storage,
|
|
}
|
|
|
|
if storage != "" {
|
|
os.MkdirAll(storage, 0755)
|
|
sm.loadSessions()
|
|
}
|
|
|
|
return sm
|
|
}
|
|
|
|
func (sm *SessionManager) GetOrCreate(key string) *Session {
|
|
sm.mu.RLock()
|
|
session, ok := sm.sessions[key]
|
|
sm.mu.RUnlock()
|
|
|
|
if !ok {
|
|
sm.mu.Lock()
|
|
session = &Session{
|
|
Key: key,
|
|
Messages: []providers.Message{},
|
|
Created: time.Now(),
|
|
Updated: time.Now(),
|
|
}
|
|
sm.sessions[key] = session
|
|
sm.mu.Unlock()
|
|
}
|
|
|
|
return session
|
|
}
|
|
|
|
func (sm *SessionManager) AddMessage(sessionKey, role, content string) {
|
|
sm.mu.Lock()
|
|
defer sm.mu.Unlock()
|
|
|
|
session, ok := sm.sessions[sessionKey]
|
|
if !ok {
|
|
session = &Session{
|
|
Key: sessionKey,
|
|
Messages: []providers.Message{},
|
|
Created: time.Now(),
|
|
}
|
|
sm.sessions[sessionKey] = session
|
|
}
|
|
|
|
session.Messages = append(session.Messages, providers.Message{
|
|
Role: role,
|
|
Content: content,
|
|
})
|
|
session.Updated = time.Now()
|
|
}
|
|
|
|
func (sm *SessionManager) GetHistory(key string) []providers.Message {
|
|
sm.mu.RLock()
|
|
defer sm.mu.RUnlock()
|
|
|
|
session, ok := sm.sessions[key]
|
|
if !ok {
|
|
return []providers.Message{}
|
|
}
|
|
|
|
history := make([]providers.Message, len(session.Messages))
|
|
copy(history, session.Messages)
|
|
return history
|
|
}
|
|
|
|
func (sm *SessionManager) Save(session *Session) error {
|
|
if sm.storage == "" {
|
|
return nil
|
|
}
|
|
|
|
sm.mu.Lock()
|
|
defer sm.mu.Unlock()
|
|
|
|
sessionPath := filepath.Join(sm.storage, session.Key+".json")
|
|
|
|
data, err := json.MarshalIndent(session, "", " ")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return os.WriteFile(sessionPath, data, 0644)
|
|
}
|
|
|
|
func (sm *SessionManager) loadSessions() error {
|
|
files, err := os.ReadDir(sm.storage)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
for _, file := range files {
|
|
if file.IsDir() {
|
|
continue
|
|
}
|
|
|
|
if filepath.Ext(file.Name()) != ".json" {
|
|
continue
|
|
}
|
|
|
|
sessionPath := filepath.Join(sm.storage, file.Name())
|
|
data, err := os.ReadFile(sessionPath)
|
|
if err != nil {
|
|
continue
|
|
}
|
|
|
|
var session Session
|
|
if err := json.Unmarshal(data, &session); err != nil {
|
|
continue
|
|
}
|
|
|
|
sm.sessions[session.Key] = &session
|
|
}
|
|
|
|
return nil
|
|
}
|