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"
)

// 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
}

// 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,
	}
}

// Add appends a message to the session history.
func (s *Store) Add(sessionID string, msg model.Message) {
	s.mu.Lock()
	sh, ok := s.sessions[sessionID]
	if !ok {
		sh = &SessionHistory{}
		s.sessions[sessionID] = sh
	}
	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()
}

// Get retrieves recent messages for a session, trimmed to fit token budget.
func (s *Store) Get(sessionID string) []model.Message {
	s.mu.RLock()
	sh, ok := s.sessions[sessionID]
	s.mu.RUnlock()
	if !ok {
		return nil
	}

	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.
// systemPrompt is injected as the first system message.
// session history, subsession results, and IoT context are merged.
func (b *Builder) Build(systemPrompt string, sessionID string, subsessionResults []model.SubSessionResult, iotContext string) []model.LLMMessage {
	messages := make([]model.LLMMessage, 0)

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

	// 2. Subsession results (injected as system context)
	for _, r := range subsessionResults {
		if r.Error != "" {
			continue
		}
		if r.Content != "" {
			msg := model.LLMMessage{
				Role:    "system",
				Content: fmt.Sprintf("[%s context]\n%s", r.Provider, r.Content),
			}
			messages = append(messages, msg)
		}
	}

	// 3. IoT device context
	if iotContext != "" {
		messages = append(messages, model.LLMMessage{
			Role:    "system",
			Content: fmt.Sprintf("[IoT devices]\n%s", iotContext),
		})
	}

	// 4. 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
}