初始提交：边缘AI算力机统一AI通讯层

2026-08-03 07:44:05 +08:00
commit 93a469061d
51 changed files with 11565 additions and 0 deletions
@@ -0,0 +1,208 @@
+package context
+
+import (
+	"github.com/edgeai/gateway/internal/config"
+	"github.com/edgeai/gateway/pkg/api"
+)
+
+// Assembler assembles context messages for a chat request.
+type Assembler struct {
+	estimator *TokenEstimator
+	cfg       *config.ContextConfig
+}
+
+// NewAssembler creates a new context assembler.
+func NewAssembler(cfg *config.ContextConfig) *Assembler {
+	return &Assembler{
+		estimator: NewTokenEstimator(),
+		cfg:       cfg,
+	}
+}
+
+// AssembleResult contains the assembled messages and metadata.
+type AssembleResult struct {
+	Messages     []api.Message
+	InputTokens  int
+	Trimmed      bool
+	TrimmedCount int
+}
+
+// Assemble combines session history with new messages, applying context window limits.
+func (a *Assembler) Assemble(history []api.Message, newMessages []api.Message, contextWindow int, maxOutputTokens int, policy string) *AssembleResult {
+	// Calculate available context for history
+	availableForHistory := contextWindow - maxOutputTokens
+	if availableForHistory < 0 {
+		availableForHistory = contextWindow / 2
+	}
+
+	// Apply safety margin
+	availableForHistory = int(float64(availableForHistory) * (1.0 - a.cfg.SafetyMarginRatio))
+
+	// Combine all messages
+	allMessages := make([]api.Message, 0, len(history)+len(newMessages))
+	allMessages = append(allMessages, history...)
+	allMessages = append(allMessages, newMessages...)
+
+	// Estimate total tokens
+	totalTokens := a.estimateAllTokens(allMessages)
+
+	if totalTokens <= availableForHistory {
+		return &AssembleResult{
+			Messages:    allMessages,
+			InputTokens: totalTokens,
+			Trimmed:     false,
+		}
+	}
+
+	// Need to trim — apply policy
+	trimmed := a.applyPolicy(allMessages, availableForHistory, policy)
+
+	return &AssembleResult{
+		Messages:     trimmed.messages,
+		InputTokens:  trimmed.tokens,
+		Trimmed:      true,
+		TrimmedCount: len(allMessages) - len(trimmed.messages),
+	}
+}
+
+type trimResult struct {
+	messages []api.Message
+	tokens   int
+}
+
+func (a *Assembler) applyPolicy(messages []api.Message, budget int, policy string) trimResult {
+	switch policy {
+	case "recent_only":
+		return a.trimRecentOnly(messages, budget)
+	case "summary_and_recent":
+		return a.trimSummaryAndRecent(messages, budget)
+	case "full":
+		return a.trimFull(messages, budget)
+	default:
+		return a.trimSummaryAndRecent(messages, budget)
+	}
+}
+
+// trimRecentOnly keeps only the most recent messages within budget.
+func (a *Assembler) trimRecentOnly(messages []api.Message, budget int) trimResult {
+	result := make([]api.Message, 0)
+	tokens := 0
+
+	// Iterate from the end (most recent first)
+	for i := len(messages) - 1; i >= 0; i-- {
+		msgTokens := a.estimateMsgTokens(messages[i])
+		if tokens+msgTokens > budget && len(result) > 0 {
+			break
+		}
+		// Prepend to maintain order
+		result = append([]api.Message{messages[i]}, result...)
+		tokens += msgTokens
+	}
+
+	return trimResult{messages: result, tokens: tokens}
+}
+
+// trimSummaryAndRecent keeps system message + a summary placeholder + recent messages.
+func (a *Assembler) trimSummaryAndRecent(messages []api.Message, budget int) trimResult {
+	if len(messages) == 0 {
+		return trimResult{}
+	}
+
+	// Always keep system messages at the front
+	systemMsgs := []api.Message{}
+	rest := []api.Message{}
+	for _, m := range messages {
+		if m.Role == "system" {
+			systemMsgs = append(systemMsgs, m)
+		} else {
+			rest = append(rest, m)
+		}
+	}
+
+	systemTokens := 0
+	for _, m := range systemMsgs {
+		systemTokens += a.estimateMsgTokens(m)
+	}
+
+	// Reserve space for a summary placeholder (~50 tokens)
+	summaryTokens := 50
+	availableForRecent := budget - systemTokens - summaryTokens
+	if availableForRecent < 0 {
+		availableForRecent = budget / 2
+	}
+
+	// Keep most recent messages
+	recentMsgs := []api.Message{}
+	recentTokens := 0
+	for i := len(rest) - 1; i >= 0; i-- {
+		msgTokens := a.estimateMsgTokens(rest[i])
+		if recentTokens+msgTokens > availableForRecent && len(recentMsgs) > 0 {
+			break
+		}
+		recentMsgs = append([]api.Message{rest[i]}, recentMsgs...)
+		recentTokens += msgTokens
+	}
+
+	// Add summary placeholder if we trimmed anything
+	result := make([]api.Message, 0, len(systemMsgs)+1+len(recentMsgs))
+	result = append(result, systemMsgs...)
+	if len(recentMsgs) < len(rest) {
+		result = append(result, api.Message{
+			Role:    "system",
+			Content: "[Earlier conversation history has been summarized and omitted.]",
+		})
+	}
+	result = append(result, recentMsgs...)
+
+	return trimResult{
+		messages: result,
+		tokens:   systemTokens + summaryTokens + recentTokens,
+	}
+}
+
+// trimFull keeps messages as-is but truncates the oldest if over budget.
+func (a *Assembler) trimFull(messages []api.Message, budget int) trimResult {
+	result := make([]api.Message, 0, len(messages))
+	tokens := 0
+
+	// Keep system messages, trim oldest non-system messages
+	systemMsgs := []api.Message{}
+	rest := []api.Message{}
+	for _, m := range messages {
+		if m.Role == "system" {
+			systemMsgs = append(systemMsgs, m)
+		} else {
+			rest = append(rest, m)
+		}
+	}
+
+	for _, m := range systemMsgs {
+		t := a.estimateMsgTokens(m)
+		tokens += t
+		result = append(result, m)
+	}
+
+	for _, m := range rest {
+		t := a.estimateMsgTokens(m)
+		if tokens+t > budget {
+			break
+		}
+		tokens += t
+		result = append(result, m)
+	}
+
+	return trimResult{messages: result, tokens: tokens}
+}
+
+func (a *Assembler) estimateAllTokens(messages []api.Message) int {
+	total := 0
+	for _, m := range messages {
+		total += a.estimateMsgTokens(m)
+	}
+	return total
+}
+
+func (a *Assembler) estimateMsgTokens(msg api.Message) int {
+	content, _ := msg.Content.(string)
+	return a.estimator.EstimateText(msg.Role) + a.estimator.EstimateText(content) + 4
+}
@@ -0,0 +1,130 @@
+package context
+
+import (
+	"testing"
+
+	"github.com/edgeai/gateway/internal/config"
+	"github.com/edgeai/gateway/pkg/api"
+)
+
+func TestTokenEstimator(t *testing.T) {
+	est := NewTokenEstimator()
+
+	// Empty string
+	if got := est.EstimateText(""); got != 0 {
+		t.Errorf("empty string: expected 0, got %d", got)
+	}
+
+	// English text
+	tokens := est.EstimateText("Hello world, this is a test.")
+	if tokens <= 0 {
+		t.Errorf("expected positive tokens for English, got %d", tokens)
+	}
+
+	// Chinese text (each char ~1 token)
+	cjkTokens := est.EstimateText("你好世界")
+	if cjkTokens != 4 {
+		t.Errorf("expected 4 tokens for 4 CJK chars, got %d", cjkTokens)
+	}
+}
+
+func TestEstimateKVCache(t *testing.T) {
+	// 1000 tokens, 32 layers, 4096 hidden dim, 2 bytes/element
+	result := EstimateKVCache(1000, 32, 4096, 2)
+	expected := int64(1000) * 32 * 2 * 4096 * 2
+	if result != expected {
+		t.Errorf("expected %d, got %d", expected, result)
+	}
+}
+
+func TestAssemblerNoTrim(t *testing.T) {
+	cfg := &config.ContextConfig{SafetyMarginRatio: 0.08}
+	a := NewAssembler(cfg)
+
+	history := []api.Message{
+		{Role: "user", Content: "Hi"},
+		{Role: "assistant", Content: "Hello!"},
+	}
+	newMsgs := []api.Message{
+		{Role: "user", Content: "How are you?"},
+	}
+
+	result := a.Assemble(history, newMsgs, 1000, 100, "summary_and_recent")
+	if result.Trimmed {
+		t.Error("expected no trimming for small context")
+	}
+	if len(result.Messages) != 3 {
+		t.Errorf("expected 3 messages, got %d", len(result.Messages))
+	}
+}
+
+func TestAssemblerTrimRecentOnly(t *testing.T) {
+	cfg := &config.ContextConfig{SafetyMarginRatio: 0.08}
+	a := NewAssembler(cfg)
+
+	// Create many messages that exceed budget
+	msgs := make([]api.Message, 20)
+	for i := range msgs {
+		msgs[i] = api.Message{Role: "user", Content: "This is message number " + string(rune('A'+i))}
+	}
+
+	result := a.Assemble(msgs, []api.Message{}, 50, 10, "recent_only")
+	if !result.Trimmed {
+		t.Error("expected trimming for large context")
+	}
+	if len(result.Messages) >= 20 {
+		t.Error("expected fewer messages after trimming")
+	}
+}
+
+func TestAssemblerSummaryAndRecent(t *testing.T) {
+	cfg := &config.ContextConfig{SafetyMarginRatio: 0.08}
+	a := NewAssembler(cfg)
+
+	msgs := make([]api.Message, 0, 22)
+	msgs = append(msgs, api.Message{Role: "system", Content: "You are a helpful assistant."})
+	for i := 0; i < 20; i++ {
+		msgs = append(msgs, api.Message{Role: "user", Content: "Message " + string(rune('A'+i%26))})
+		msgs = append(msgs, api.Message{Role: "assistant", Content: "Response " + string(rune('A'+i%26))})
+	}
+
+	result := a.Assemble(msgs, []api.Message{}, 80, 20, "summary_and_recent")
+	if !result.Trimmed {
+		t.Error("expected trimming")
+	}
+
+	// System message should be preserved
+	hasSystem := false
+	hasSummary := false
+	for _, m := range result.Messages {
+		if m.Role == "system" {
+			if content, ok := m.Content.(string); ok {
+				if content == "You are a helpful assistant." {
+					hasSystem = true
+				}
+				if contains(content, "summarized") {
+					hasSummary = true
+				}
+			}
+		}
+	}
+	if !hasSystem {
+		t.Error("system message should be preserved")
+	}
+	if !hasSummary {
+		t.Error("summary placeholder should be present when trimmed")
+	}
+}
+
+func contains(s, substr string) bool {
+	return len(s) >= len(substr) && (s == substr || (len(s) > len(substr) && (indexOf(s, substr) >= 0)))
+}
+
+func indexOf(s, substr string) int {
+	for i := 0; i <= len(s)-len(substr); i++ {
+		if s[i:i+len(substr)] == substr {
+			return i
+		}
+	}
+	return -1
+}
@@ -0,0 +1,82 @@
+package context
+
+import (
+	"strings"
+	"unicode"
+)
+
+// TokenEstimator estimates token counts for text using a simple heuristic.
+// For production use, replace with a proper tokenizer (tiktoken, etc.).
+type TokenEstimator struct {
+	charsPerToken float64
+}
+
+// NewTokenEstimator creates a new estimator with the default ratio.
+// English text averages ~4 chars/token, Chinese ~1.5 chars/token.
+func NewTokenEstimator() *TokenEstimator {
+	return &TokenEstimator{charsPerToken: 3.0}
+}
+
+// EstimateText estimates token count for a given text.
+func (e *TokenEstimator) EstimateText(text string) int {
+	if text == "" {
+		return 0
+	}
+
+	// Count CJK characters as individual tokens
+	cjkCount := 0
+	nonCJKChars := 0
+	for _, r := range text {
+		if unicode.Is(unicode.Han, r) || unicode.Is(unicode.Hiragana, r) || unicode.Is(unicode.Katakana, r) || unicode.Is(unicode.Hangul, r) {
+			cjkCount++
+		} else {
+			nonCJKChars++
+		}
+	}
+
+	// Non-CJK: estimate by chars/token ratio
+	nonCJKTokens := int(float64(nonCJKChars) / e.charsPerToken)
+	if nonCJKChars > 0 && nonCJKTokens == 0 {
+		nonCJKTokens = 1
+	}
+
+	return cjkCount + nonCJKTokens
+}
+
+// EstimateMessage estimates token count for a single message (including role overhead).
+func (e *TokenEstimator) EstimateMessage(msg interface{ GetRole() string; GetContent() string }) int {
+	role := msg.GetRole()
+	content := msg.GetContent()
+	// Role tokens: ~1-2 tokens for role name
+	roleTokens := len(strings.Fields(role)) + 1
+	return roleTokens + e.EstimateText(content)
+}
+
+// EstimateMessages estimates total token count for a list of messages.
+func (e *TokenEstimator) EstimateMessages(messages []Message) int {
+	total := 0
+	for _, m := range messages {
+		total += e.EstimateText(m.Role) + e.EstimateText(m.Content) + 4 // role + content + formatting overhead
+	}
+	return total
+}
+
+// Message is a simplified message structure for estimation.
+type Message struct {
+	Role    string
+	Content string
+}
+
+func (m Message) GetRole() string    { return m.Role }
+func (m Message) GetContent() string { return m.Content }
+
+// EstimateKVCache estimates the KV cache memory usage in bytes.
+// Formula: input_tokens × layers × 2 (K+V) × hidden_dim × bytes_per_element
+func EstimateKVCache(inputTokens, layers, hiddenDim, bytesPerElement int) int64 {
+	return int64(inputTokens) * int64(layers) * 2 * int64(hiddenDim) * int64(bytesPerElement)
+}
+
+// EstimateKVCachePerToken estimates KV cache per token in bytes.
+func EstimateKVCachePerToken(layers, hiddenDim, bytesPerElement int) int64 {
+	return int64(layers) * 2 * int64(hiddenDim) * int64(bytesPerElement)
+}