DemoCoder/context_compaction.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"

	"gocoder/agents"
	"gocoder/db"
)

const (
	defaultContextCompactionTriggerRatio = 0.80
	defaultContextCompactionTargetRatio  = 0.70
	contextCompactionSummaryFloor        = 256
	contextCompactionSummaryCeiling      = 1024
	contextCompactionSummaryDivisor      = 20
)

func configuredCompactionTriggerRatio() float64 {
	cfg := currentConfig()
	if cfg.LLM.ContextCompactionTriggerRatio > 0 && cfg.LLM.ContextCompactionTriggerRatio < 1 {
		return cfg.LLM.ContextCompactionTriggerRatio
	}
	return defaultContextCompactionTriggerRatio
}

func configuredCompactionTargetRatio() float64 {
	trigger := configuredCompactionTriggerRatio()
	ratio := defaultContextCompactionTargetRatio
	cfg := currentConfig()
	if cfg.LLM.ContextCompactionTargetRatio > 0 && cfg.LLM.ContextCompactionTargetRatio < 1 {
		ratio = cfg.LLM.ContextCompactionTargetRatio
	}
	if ratio >= trigger {
		ratio = trigger * 0.875
		if ratio <= 0 {
			ratio = defaultContextCompactionTargetRatio
		}
	}
	return ratio
}

func effectiveContextTokenLimit() int {
	cfg := currentConfig()
	if cfg.LLM.ContextTokens > 0 {
		return cfg.LLM.ContextTokens
	}
	if cfg.LLM.MaxTokens > 0 {
		return cfg.LLM.MaxTokens
	}
	return 0
}

func estimateTextTokens(text string) int {
	trimmed := strings.TrimSpace(text)
	if trimmed == "" {
		return 0
	}

	tokens := len(trimmed)/4 + 1
	if tokens < 1 {
		tokens = 1
	}
	return tokens
}

func estimateMessageTokens(msg db.Message) int {
	tokens := estimateTextTokens(msg.Content)
	if tokens == 0 {
		return 0
	}
	return tokens + 4
}

func estimatePromptTokens(messages []db.Message) int {
	total := 0
	for _, msg := range messages {
		total += estimateMessageTokens(msg)
	}
	return total
}

func formatConversationSummary(summary string, summaryCutoffOrder int) string {
	trimmed := strings.TrimSpace(summary)
	if trimmed == "" {
		return ""
	}
	if summaryCutoffOrder > 0 {
		return fmt.Sprintf("Conversation memory summary up to message %d:\n%s", summaryCutoffOrder, trimmed)
	}
	return "Conversation memory summary:\n" + trimmed
}

func buildPromptHistory(history []db.Message, summary string, summaryCutoffOrder int) []db.Message {
	promptHistory := make([]db.Message, 0, len(history)+1)
	if summaryText := formatConversationSummary(summary, summaryCutoffOrder); summaryText != "" {
		promptHistory = append(promptHistory, db.Message{
			Role:    "system",
			Kind:    "message",
			Content: summaryText,
		})
	}
	for _, msg := range history {
		if summaryCutoffOrder > 0 && msg.Order <= summaryCutoffOrder {
			continue
		}
		promptHistory = append(promptHistory, msg)
	}
	return promptHistory
}

func filterMessagesAfterOrder(history []db.Message, order int) []db.Message {
	if order <= 0 {
		return append([]db.Message(nil), history...)
	}

	filtered := make([]db.Message, 0, len(history))
	for _, msg := range history {
		if msg.Order > order {
			filtered = append(filtered, msg)
		}
	}
	return filtered
}

func selectMessagesThroughOrder(history []db.Message, startExclusive, endInclusive int) []db.Message {
	if endInclusive <= startExclusive {
		return nil
	}

	selected := make([]db.Message, 0)
	for _, msg := range history {
		if msg.Order <= startExclusive {
			continue
		}
		if msg.Order > endInclusive {
			break
		}
		selected = append(selected, msg)
	}
	return selected
}

func estimateSummaryReserveTokens(limit int) int {
	reserve := limit / contextCompactionSummaryDivisor
	if reserve < contextCompactionSummaryFloor {
		reserve = contextCompactionSummaryFloor
	}
	if reserve > contextCompactionSummaryCeiling {
		reserve = contextCompactionSummaryCeiling
	}
	return reserve
}

func selectCompactionCutoff(systemMsgs []db.Message, history []db.Message, summary string, summaryCutoffOrder int, userMessage string, contextLimit int) (int, bool) {
	if contextLimit <= 0 || len(history) == 0 {
		return 0, false
	}

	trigger := int(float64(contextLimit) * configuredCompactionTriggerRatio())
	target := int(float64(contextLimit) * configuredCompactionTargetRatio())
	if target < 1 {
		target = 1
	}

	currentEstimate := estimatePromptTokens(buildPromptMessages(systemMsgs, buildPromptHistory(history, summary, summaryCutoffOrder), userMessage))
	if currentEstimate <= trigger {
		return 0, false
	}

	summaryReserve := estimateSummaryReserveTokens(contextLimit)
	for _, msg := range history {
		if msg.Order <= summaryCutoffOrder {
			continue
		}

		tailHistory := filterMessagesAfterOrder(history, msg.Order)
		estimate := estimatePromptTokens(buildPromptMessages(systemMsgs, buildPromptHistory(tailHistory, "", 0), userMessage)) + summaryReserve
		if estimate <= target {
			return msg.Order, true
		}
	}

	lastOrder := 0
	for i := len(history) - 1; i >= 0; i-- {
		if history[i].Order > summaryCutoffOrder {
			lastOrder = history[i].Order
			break
		}
	}
	if lastOrder > summaryCutoffOrder {
		return lastOrder, true
	}

	return 0, false
}

func renderCompactionTranscript(messages []db.Message) string {
	if len(messages) == 0 {
		return ""
	}

	var b strings.Builder
	for i, msg := range messages {
		if i > 0 {
			b.WriteString("\n")
		}
		fmt.Fprintf(&b, "Order %d | role=%s | kind=%s", msg.Order, strings.TrimSpace(msg.Role), strings.TrimSpace(msg.Kind))
		if name := strings.TrimSpace(msg.Name); name != "" {
			fmt.Fprintf(&b, " | name=%s", name)
		}
		if agent := strings.TrimSpace(msg.Agent); agent != "" {
			fmt.Fprintf(&b, " | agent=%s", agent)
		}
		b.WriteString("\n")
		content := strings.TrimSpace(msg.Content)
		if content == "" {
			content = "[empty]"
		}
		b.WriteString(content)
		b.WriteString("\n")
	}
	return strings.TrimSpace(b.String())
}

func callLMStudioCompletion(ctx context.Context, messages []db.Message, maxTokens int, temperature float64) (string, error) {
	requestCtx := ctx
	if requestCtx == nil {
		requestCtx = context.Background()
	}

	cfg := currentConfig()
	reqBody := buildLMStudioRequestWithLimits(messages, false, cfg.LLM, maxTokens, temperature)
	jsonData, err := json.Marshal(reqBody)
	if err != nil {
		return "", err
	}

	req, err := http.NewRequestWithContext(requestCtx, http.MethodPost, cfg.LLM.BaseURL+"/chat/completions", bytes.NewBuffer(jsonData))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		_, _ = io.ReadAll(resp.Body)
		return "", fmt.Errorf("LMStudio API error: status=%d", resp.StatusCode)
	}

	var lmResp LMStudioResponse
	if err := json.NewDecoder(resp.Body).Decode(&lmResp); err != nil {
		return "", err
	}
	if len(lmResp.Choices) == 0 {
		return "", fmt.Errorf("empty response from LLM")
	}

	content := strings.TrimSpace(lmResp.Choices[0].Message.Content)
	if content == "" {
		return "", fmt.Errorf("empty response from LLM")
	}
	return content, nil
}

func summarizeConversationChunk(ctx context.Context, existingSummary string, messages []db.Message, summaryReserveTokens int) (string, error) {
	messagesToSend := []db.Message{
		{
			Role:    "system",
			Content: agents.ContextCompactorSystemMessage(),
		},
	}

	if trimmed := strings.TrimSpace(existingSummary); trimmed != "" {
		messagesToSend = append(messagesToSend, db.Message{
			Role:    "system",
			Content: "Existing memory to preserve and compress:\n" + trimmed,
		})
	}

	transcript := renderCompactionTranscript(messages)
	if transcript != "" {
		messagesToSend = append(messagesToSend, db.Message{
			Role:    "user",
			Content: "Compress the following conversation chunk into durable memory:\n\n" + transcript,
		})
	}

	return callLMStudioCompletion(ctx, messagesToSend, summaryReserveTokens, 0.1)
}

func maybeCompactConversationContext(ctx context.Context, conv *db.Conversation, history []db.Message, systemMsgs []db.Message, userMessage string) bool {
	if conv == nil || store == nil {
		return false
	}

	contextLimit := effectiveContextTokenLimit()
	cutoff, ok := selectCompactionCutoff(systemMsgs, history, conv.ContextSummary, conv.ContextSummaryOrder, userMessage, contextLimit)
	if !ok || cutoff <= conv.ContextSummaryOrder {
		return false
	}

	chunk := selectMessagesThroughOrder(history, conv.ContextSummaryOrder, cutoff)
	if len(chunk) == 0 {
		return false
	}

	summaryReserveTokens := estimateSummaryReserveTokens(contextLimit)
	nextSummary, err := summarizeConversationChunk(ctx, conv.ContextSummary, chunk, summaryReserveTokens)
	if err != nil {
		log.Printf("failed to compact conversation %s through order %d: %v", conv.ID, cutoff, err)
		return false
	}

	if err := store.UpdateConversationContextSummary(conv.ID, nextSummary, cutoff); err != nil {
		log.Printf("failed to store compacted conversation memory for %s: %v", conv.ID, err)
		return false
	}

	conv.ContextSummary = nextSummary
	conv.ContextSummaryOrder = cutoff
	log.Printf("compacted conversation %s context through message %d", conv.ID, cutoff)
	return true
}