MuyueWorkspace/internal/api/conversation.go

package api

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"
	"unicode/utf8"

	"github.com/muyue/muyue/internal/config"
)

const maxTokensApprox = 100000
const summarizeThreshold = 80000
const charsPerToken = 4

type FeedMessage struct {
	ID      string `json:"id"`
	Role    string `json:"role"`
	Content string `json:"content"`
	Time    string `json:"time"`
}

type Conversation struct {
	Messages  []FeedMessage `json:"messages"`
	Summary   string        `json:"summary,omitempty"`
	CreatedAt string        `json:"created_at"`
	UpdatedAt string        `json:"updated_at"`
}

type ConversationStore struct {
	mu   sync.RWMutex
	path string
	conv *Conversation
}

func NewConversationStore() *ConversationStore {
	dir, err := config.ConfigDir()
	if err != nil {
		dir = "/tmp/muyue"
	}
	path := filepath.Join(dir, "conversation.json")
	cs := &ConversationStore{path: path}
	cs.load()
	return cs
}

func (cs *ConversationStore) load() {
	data, err := os.ReadFile(cs.path)
	if err != nil {
		cs.conv = &Conversation{
			Messages:  []FeedMessage{},
			CreatedAt: time.Now().Format(time.RFC3339),
			UpdatedAt: time.Now().Format(time.RFC3339),
		}
		return
	}
	var conv Conversation
	if err := json.Unmarshal(data, &conv); err != nil {
		cs.conv = &Conversation{
			Messages:  []FeedMessage{},
			CreatedAt: time.Now().Format(time.RFC3339),
			UpdatedAt: time.Now().Format(time.RFC3339),
		}
		return
	}
	if conv.Messages == nil {
		conv.Messages = []FeedMessage{}
	}
	cs.conv = &conv
}

func (cs *ConversationStore) save() error {
	cs.conv.UpdatedAt = time.Now().Format(time.RFC3339)
	data, err := json.MarshalIndent(cs.conv, "", "  ")
	if err != nil {
		return err
	}
	dir := filepath.Dir(cs.path)
	os.MkdirAll(dir, 0755)
	return os.WriteFile(cs.path, data, 0600)
}

func (cs *ConversationStore) Get() []FeedMessage {
	cs.mu.RLock()
	defer cs.mu.RUnlock()
	out := make([]FeedMessage, len(cs.conv.Messages))
	copy(out, cs.conv.Messages)
	return out
}

func (cs *ConversationStore) GetSummary() string {
	cs.mu.RLock()
	defer cs.mu.RUnlock()
	return cs.conv.Summary
}

func (cs *ConversationStore) Add(role, content string) FeedMessage {
	cs.mu.Lock()
	defer cs.mu.Unlock()

	msg := FeedMessage{
		ID:      generateMsgID(),
		Role:    role,
		Content: content,
		Time:    time.Now().Format(time.RFC3339),
	}
	cs.conv.Messages = append(cs.conv.Messages, msg)
	cs.save()
	return msg
}

func (cs *ConversationStore) Clear() {
	cs.mu.Lock()
	defer cs.mu.Unlock()
	cs.conv.Messages = []FeedMessage{}
	cs.conv.Summary = ""
	cs.conv.CreatedAt = time.Now().Format(time.RFC3339)
	cs.conv.UpdatedAt = time.Now().Format(time.RFC3339)
	cs.save()
}

func (cs *ConversationStore) SetSummary(summary string) {
	cs.mu.Lock()
	defer cs.mu.Unlock()
	cs.conv.Summary = summary
	cs.save()
}

func (cs *ConversationStore) TrimOld(keepCount int) {
	cs.mu.Lock()
	defer cs.mu.Unlock()
	if len(cs.conv.Messages) <= keepCount {
		return
	}
	cs.conv.Messages = cs.conv.Messages[len(cs.conv.Messages)-keepCount:]
	cs.save()
}

func (cs *ConversationStore) ApproxTokenCount() int {
	cs.mu.RLock()
	defer cs.mu.RUnlock()
	total := utf8.RuneCountInString(cs.conv.Summary)
	for _, m := range cs.conv.Messages {
		total += utf8.RuneCountInString(m.Content)
	}
	return total / charsPerToken
}

func (cs *ConversationStore) NeedsSummarization() bool {
	return cs.ApproxTokenCount() > summarizeThreshold
}

func generateMsgID() string {
	return time.Now().Format("20060102150405.000") + "-" + fmt.Sprintf("%d", time.Now().UnixNano())
}