gochat/backend/internal/service/token_estimator.go

package service

import (
	"strings"
)

// TokenEstimator provides approximate token counting for LLM context management.
// Uses the industry-standard heuristic of ~4 characters per token for English text.
// This is not exact but sufficient for context window truncation decisions.
//
// Reference: AI_FEATURE_ROADMAP.md §3.4 — token counting + sliding window
type TokenEstimator struct{}

// NewTokenEstimator creates a new TokenEstimator.
func NewTokenEstimator() *TokenEstimator {
	return &TokenEstimator{}
}

// EstimateText returns an approximate token count for the given text.
// Heuristic: ~4 characters per token (works well for English; slightly
// underestimates for CJK text which is ~2 chars/token).
func (e *TokenEstimator) EstimateText(text string) int {
	if text == "" {
		return 0
	}
	// Count characters (not bytes) for better accuracy with Unicode
	charCount := len([]rune(text))
	return (charCount + 3) / 4 // ceiling division
}

// EstimateMessages returns the total approximate token count for a list of LLM messages.
// Each message has overhead (~4 tokens for role + delimiters).
func (e *TokenEstimator) EstimateMessages(messages []llmMessageWithContent) int {
	total := 0
	for _, msg := range messages {
		total += 4 // role + delimiter overhead
		total += e.EstimateText(msg.Content)
	}
	total += 2 // priming tokens
	return total
}

// llmMessageWithContent is a minimal interface for messages with content.
type llmMessageWithContent struct {
	Role    string
	Content string
}

// TruncateMessages applies a sliding window to fit messages within a token budget.
// It keeps the most recent messages (by index) and drops older ones until
// the total estimated tokens fit within maxTokens.
//
// If the most recent single message exceeds maxTokens, it is truncated to fit.
//
// Returns the truncated message list and the number of messages dropped.
func (e *TokenEstimator) TruncateMessages(messages []llmMessageWithContent, maxTokens int) ([]llmMessageWithContent, int) {
	if maxTokens <= 0 {
		return messages, 0
	}

	// Calculate total tokens
	totalTokens := e.EstimateMessages(messages)
	if totalTokens <= maxTokens {
		return messages, 0
	}

	// Drop oldest messages until we fit
	dropped := 0
	for len(messages) > 1 {
		messages = messages[1:]
		dropped++
		totalTokens = e.EstimateMessages(messages)
		if totalTokens <= maxTokens {
			return messages, dropped
		}
	}

	// If single message still too long, truncate its content
	if len(messages) == 1 {
		overhead := 6 // role + delimiters + priming
		availableTokens := maxTokens - overhead
		if availableTokens > 0 {
			maxChars := availableTokens * 4
			content := messages[0].Content
			if len([]rune(content)) > maxChars {
				runes := []rune(content)
				messages[0].Content = string(runes[:maxChars]) + "..."
			}
		}
	}

	return messages, dropped
}

// BuildContextWindow retrieves messages and applies token-budgeted truncation.
// This is the main entry point for CopilotContextService and other services
// that need to build LLM context from conversation history.
//
// Parameters:
//   - rawMessages: all messages from the conversation (oldest first)
//   - maxContextTokens: token budget for the message context (default 4096 if <= 0)
//
// Returns the truncated messages (most recent kept) formatted as role/content pairs.
func BuildContextWindow(rawMessages []ConversationMessage, maxContextTokens int) ([]llmMessageWithContent, int) {
	if maxContextTokens <= 0 {
		maxContextTokens = 4096
	}

	// Convert to estimator format
	msgs := make([]llmMessageWithContent, 0, len(rawMessages))
	for _, m := range rawMessages {
		content := strings.TrimSpace(m.Content)
		if content == "" {
			continue
		}
		// Skip activity/template messages
		if m.MessageType == "activity" {
			continue
		}
		role := "user"
		if m.MessageType == "outgoing" {
			role = "assistant"
		}
		msgs = append(msgs, llmMessageWithContent{Role: role, Content: content})
	}

	estimator := NewTokenEstimator()
	return estimator.TruncateMessages(msgs, maxContextTokens)
}

// ConversationMessage is a minimal message representation for context building.
// Used to avoid importing model.Message in the estimator (keeps it testable).
type ConversationMessage struct {
	Content     string
	MessageType string // "incoming", "outgoing", "activity"
}