compaction.go

package tok

import (
	"encoding/json"
	"fmt"
	"strings"
)

// CompactionSchema is a 5-field structured schema for LLM-based context
// compaction. It preserves critical information while maximizing compression.
type CompactionSchema struct {
	TaskOverview         string   `json:"task_overview"`
	CurrentState         string   `json:"current_state"`
	ImportantDiscoveries []string `json:"important_discoveries"`
	NextSteps            []string `json:"next_steps"`
	ContextToPreserve    []string `json:"context_to_preserve"`
}

// ToPrompt renders the schema as a structured prompt for reinsertion into context.
func (s *CompactionSchema) ToPrompt() string {
	var parts []string

	if s.TaskOverview != "" {
		parts = append(parts, "## Task Overview\n"+s.TaskOverview)
	}
	if s.CurrentState != "" {
		parts = append(parts, "## Current State\n"+s.CurrentState)
	}
	if len(s.ImportantDiscoveries) > 0 {
		items := make([]string, len(s.ImportantDiscoveries))
		for i, d := range s.ImportantDiscoveries {
			items[i] = "- " + d
		}
		parts = append(parts, "## Important Discoveries\n"+strings.Join(items, "\n"))
	}
	if len(s.NextSteps) > 0 {
		items := make([]string, len(s.NextSteps))
		for i, s := range s.NextSteps {
			items[i] = "- " + s
		}
		parts = append(parts, "## Next Steps\n"+strings.Join(items, "\n"))
	}
	if len(s.ContextToPreserve) > 0 {
		items := make([]string, len(s.ContextToPreserve))
		for i, c := range s.ContextToPreserve {
			items[i] = "- " + c
		}
		parts = append(parts, "## Context to Preserve\n"+strings.Join(items, "\n"))
	}

	return strings.Join(parts, "\n\n")
}

// CompactionSystemPrompt is the system prompt for LLM-based structured compaction.
const CompactionSystemPrompt = `You are a context compressor. Given a conversation history, extract and structure the essential information into exactly these 5 fields:

1. task_overview: One sentence describing what's being accomplished
2. current_state: Where we are now (1-2 sentences)
3. important_discoveries: Bullet list of key findings, decisions, constraints (max 10 items)
4. next_steps: What needs to happen next (max 5 items)
5. context_to_preserve: Specific details that must not be lost — file paths, variable values, error messages, URLs, version numbers (max 10 items)

Respond ONLY with valid JSON matching this schema. Be maximally concise — every word must earn its place.`

// BuildCompactionPrompt builds a prompt for LLM-based structured compaction.
func BuildCompactionPrompt(context string, maxChars int) string {
	if maxChars <= 0 {
		maxChars = 8000
	}
	if len(context) > maxChars {
		context = context[:maxChars]
	}
	return fmt.Sprintf("%s\n\n---\nContext to compress:\n%s", CompactionSystemPrompt, context)
}

// ParseCompactionResponse parses an LLM's JSON response into a CompactionSchema.
func ParseCompactionResponse(response string) (*CompactionSchema, error) {
	text := strings.TrimSpace(response)

	// Strip markdown code fences if present
	if strings.HasPrefix(text, "```") {
		lines := strings.Split(text, "\n")
		if len(lines) > 2 {
			end := len(lines) - 1
			if strings.TrimSpace(lines[end]) == "```" {
				lines = lines[1:end]
			} else {
				lines = lines[1:]
			}
			text = strings.Join(lines, "\n")
		}
	}

	var schema CompactionSchema
	if err := json.Unmarshal([]byte(text), &schema); err != nil {
		// Fallback: treat entire response as task_overview
		fallback := text
		if len(fallback) > 500 {
			fallback = fallback[:500]
		}
		return &CompactionSchema{TaskOverview: fallback}, nil
	}
	return &schema, nil
}