153 lines
3.3 KiB
Go
153 lines
3.3 KiB
Go
package conversation
|
|
|
|
import (
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"aichat/message"
|
|
"aichat/utils"
|
|
)
|
|
|
|
type Store struct {
|
|
dir string
|
|
mu sync.Mutex
|
|
}
|
|
|
|
func NewStore(dir string) *Store {
|
|
os.MkdirAll(dir, 0755)
|
|
return &Store{dir: dir}
|
|
}
|
|
|
|
func (s *Store) path(id string) string {
|
|
return filepath.Join(s.dir, id+".json")
|
|
}
|
|
|
|
func (s *Store) Create() (*message.Conversation, error) {
|
|
conv := &message.Conversation{
|
|
ID: utils.NewUUID(),
|
|
Title: "新对话",
|
|
CreatedAt: time.Now(),
|
|
UpdatedAt: time.Now(),
|
|
}
|
|
if err := s.Save(conv); err != nil {
|
|
return nil, err
|
|
}
|
|
return conv, nil
|
|
}
|
|
|
|
func (s *Store) Save(conv *message.Conversation) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
conv.UpdatedAt = time.Now()
|
|
return atomicWriteJSON(s.path(conv.ID), conv)
|
|
}
|
|
|
|
func (s *Store) Get(id string) (*message.Conversation, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
data, err := os.ReadFile(s.path(id))
|
|
if err != nil {
|
|
if os.IsNotExist(err) {
|
|
return nil, errors.New("对话不存在")
|
|
}
|
|
return nil, fmt.Errorf("读取对话失败: %w", err)
|
|
}
|
|
var conv message.Conversation
|
|
if err := json.Unmarshal(data, &conv); err != nil {
|
|
return nil, fmt.Errorf("解析对话失败: %w", err)
|
|
}
|
|
return &conv, nil
|
|
}
|
|
|
|
func (s *Store) List() ([]message.Conversation, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
entries, err := os.ReadDir(s.dir)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("读取对话目录失败: %w", err)
|
|
}
|
|
|
|
var list []message.Conversation
|
|
for _, e := range entries {
|
|
if e.IsDir() || filepath.Ext(e.Name()) != ".json" {
|
|
continue
|
|
}
|
|
data, err := os.ReadFile(filepath.Join(s.dir, e.Name()))
|
|
if err != nil {
|
|
continue
|
|
}
|
|
var conv message.Conversation
|
|
if err := json.Unmarshal(data, &conv); err != nil {
|
|
continue
|
|
}
|
|
conv.Messages = nil // 列表不返回消息体
|
|
list = append(list, conv)
|
|
}
|
|
|
|
sort.Slice(list, func(i, j int) bool {
|
|
return list[i].UpdatedAt.After(list[j].UpdatedAt)
|
|
})
|
|
return list, nil
|
|
}
|
|
|
|
func (s *Store) Delete(id string) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if err := os.Remove(s.path(id)); err != nil && !os.IsNotExist(err) {
|
|
return fmt.Errorf("删除对话失败: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func atomicWriteJSON(path string, v any) error {
|
|
tmp := path + ".tmp"
|
|
data, err := json.Marshal(v)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := os.WriteFile(tmp, data, 0644); err != nil {
|
|
return err
|
|
}
|
|
return os.Rename(tmp, path)
|
|
}
|
|
|
|
func SaveMessages(store *Store, id string, messages []message.ChatMessage, assistantContent string) error {
|
|
conv, err := store.Get(id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
conv.Messages = append([]message.ChatMessage(nil), messages...)
|
|
conv.Messages = append(conv.Messages, message.ChatMessage{Role: "assistant", Content: assistantContent})
|
|
if conv.Title == "" || conv.Title == "新对话" {
|
|
conv.Title = GenTitle(conv.Messages)
|
|
}
|
|
return store.Save(conv)
|
|
}
|
|
|
|
func GenTitle(messages []message.ChatMessage) string {
|
|
for _, m := range messages {
|
|
if m.Hidden {
|
|
continue
|
|
}
|
|
if m.Role == "user" && strings.TrimSpace(m.Content) != "" {
|
|
title := strings.TrimSpace(m.Content)
|
|
title = strings.ReplaceAll(title, "\r\n", " ")
|
|
title = strings.ReplaceAll(title, "\n", " ")
|
|
runes := []rune(title)
|
|
if len(runes) > 30 {
|
|
return string(runes[:30]) + "..."
|
|
}
|
|
return title
|
|
}
|
|
}
|
|
return "新对话"
|
|
}
|