Cyrene_v2/backend/services/ai-core/context/context.go

// Package context manages conversation context with sliding window and token budgeting.
package context

import (
	"fmt"
	"strings"
	"sync"
	"unicode/utf8"

	"git.yeij.top/AskaEth/Cyrene_v2/services/ai-core/model"
)

// Persister saves conversation messages durably and can restore them
// (e.g. from PostgreSQL) when a session is not in memory.
type Persister interface {
	// SaveMessage persists one message. Called synchronously by Add —
	// implementations should be fast or dispatch async work themselves.
	SaveMessage(sessionID, userID string, msg model.Message)
	// LoadRecent returns the most recent messages of a session, ascending.
	LoadRecent(sessionID string, limit int) []model.Message
}

// Store holds conversation history with a sliding window.
type Store struct {
	mu       sync.RWMutex
	sessions map[string]*SessionHistory
	maxTokens int // Maximum tokens to include in context
	maxTurns  int // Maximum message turns to keep
	persister Persister
}

// SessionHistory holds messages for a single session.
type SessionHistory struct {
	Messages []model.Message
	mu       sync.RWMutex
}

// NewStore creates a new conversation store.
func NewStore(maxTokens, maxTurns int) *Store {
	if maxTokens <= 0 {
		maxTokens = 4000
	}
	if maxTurns <= 0 {
		maxTurns = 50
	}
	return &Store{
		sessions:  make(map[string]*SessionHistory),
		maxTokens: maxTokens,
		maxTurns:  maxTurns,
	}
}

// SetPersister injects durable persistence for conversation messages.
func (s *Store) SetPersister(p Persister) {
	s.mu.Lock()
	s.persister = p
	s.mu.Unlock()
}

// Add appends a message to the session history.
func (s *Store) Add(sessionID string, msg model.Message) {
	s.add(sessionID, "", msg)
}

// AddWithUser appends a message and records the owning user for persistence.
func (s *Store) AddWithUser(sessionID, userID string, msg model.Message) {
	s.add(sessionID, userID, msg)
}

func (s *Store) add(sessionID, userID string, msg model.Message) {
	s.mu.Lock()
	sh, ok := s.sessions[sessionID]
	if !ok {
		sh = &SessionHistory{}
		s.sessions[sessionID] = sh
	}
	persister := s.persister
	s.mu.Unlock()

	sh.mu.Lock()
	sh.Messages = append(sh.Messages, msg)
	// Trim to max turns
	if len(sh.Messages) > s.maxTurns {
		sh.Messages = sh.Messages[len(sh.Messages)-s.maxTurns:]
	}
	sh.mu.Unlock()

	// Persist asynchronously — never block the conversation on the database.
	if persister != nil && (msg.Role == "user" || msg.Role == "assistant") {
		go persister.SaveMessage(sessionID, userID, msg)
	}
}

// Get retrieves recent messages for a session, trimmed to fit token budget.
// Sessions missing from memory are restored from the persister if available.
func (s *Store) Get(sessionID string) []model.Message {
	s.mu.RLock()
	sh, ok := s.sessions[sessionID]
	persister := s.persister
	s.mu.RUnlock()

	if !ok {
		if persister == nil {
			return nil
		}
		loaded := persister.LoadRecent(sessionID, s.maxTurns)
		if len(loaded) == 0 {
			return nil
		}
		sh = &SessionHistory{Messages: loaded}
		s.mu.Lock()
		s.sessions[sessionID] = sh
		s.mu.Unlock()
	}

	sh.mu.RLock()
	defer sh.mu.RUnlock()

	// Return all if within budget
	tokens := countTokens(sh.Messages)
	if tokens <= s.maxTokens {
		// Return a copy
		result := make([]model.Message, len(sh.Messages))
		copy(result, sh.Messages)
		return result
	}

	// Slide window from the end, keeping most recent messages
	// Always keep at least the last 4 messages (2 turns)
	minKeep := 4
	if len(sh.Messages) < minKeep {
		minKeep = len(sh.Messages)
	}

	result := make([]model.Message, 0, minKeep)
	tokens = 0
	for i := len(sh.Messages) - 1; i >= 0; i-- {
		msgTokens := tokenEstimate(sh.Messages[i].Content)
		if tokens+msgTokens > s.maxTokens && len(result) >= minKeep {
			break
		}
		result = append([]model.Message{sh.Messages[i]}, result...)
		tokens += msgTokens
	}
	return result
}

// Clear removes all messages for a session.
func (s *Store) Clear(sessionID string) {
	s.mu.Lock()
	delete(s.sessions, sessionID)
	s.mu.Unlock()
}

// Builder assembles the full context for an LLM call.
type Builder struct {
	store *Store
}

// NewBuilder creates a context builder.
func NewBuilder(store *Store) *Builder {
	return &Builder{store: store}
}

// Build creates the full message array for an LLM request.
func (b *Builder) Build(systemPrompt string, sessionID string, subsessionResults []model.SubSessionResult, iotContext string) []model.LLMMessage {
	messages := make([]model.LLMMessage, 0)

	// 1. System prompt (persona + rules)
	if systemPrompt != "" {
		messages = append(messages, model.LLMMessage{Role: "system", Content: systemPrompt})
	}

	// 2. Context: merge all subsession/memory/iot into ONE system message
	//    with clear delimiters to prevent LLM from confusing context with user input
	var ctxParts []string

	for _, r := range subsessionResults {
		if r.Error != "" || r.Content == "" || r.Content == "no relevant memories" {
			continue
		}
		if r.Provider == "general" && strings.HasPrefix(r.Content, "Intent:") {
			continue // Internal data, not for LLM
		}
		if r.Provider == "memory" {
			ctxParts = append(ctxParts, fmt.Sprintf("【相关记忆】\n%s", r.Content))
		} else if r.Provider == "iot" {
			ctxParts = append(ctxParts, fmt.Sprintf("【设备状态】\n%s", r.Content))
		}
	}

	if iotContext != "" {
		ctxParts = append(ctxParts, fmt.Sprintf("【设备列表】\n%s", iotContext))
	}

	if len(ctxParts) > 0 {
		contextBlock := "以下是系统上下文，不是用户说的话：\n" + strings.Join(ctxParts, "\n\n")
		messages = append(messages, model.LLMMessage{Role: "system", Content: contextBlock})
	}

	// 3. Conversation history
	history := b.store.Get(sessionID)
	for _, m := range history {
		messages = append(messages, model.LLMMessage{Role: m.Role, Content: m.Content})
	}

	return messages
}

// --- Helpers ---

// tokenEstimate estimates the number of tokens in a string (rough: 1 token ≈ 2 Chinese chars ≈ 0.75 English words).
func tokenEstimate(text string) int {
	words := strings.Fields(text)
	chars := utf8.RuneCountInString(text)
	// Approximate: Chinese-heavy text ~ chars/2, English-heavy ~ words*1.3
	return (chars/2 + int(float64(len(words))*1.3)) / 2
}

// countTokens estimates total tokens in a message slice.
func countTokens(msgs []model.Message) int {
	total := 0
	for _, m := range msgs {
		total += tokenEstimate(m.Content)
	}
	return total
}