first commit
This commit is contained in:
@@ -0,0 +1,367 @@
|
||||
// This file (US-003) covers summarization: the structured prompts, the
|
||||
// conversation serializer, file-operation extraction, and GenerateSummary,
|
||||
// which drives a provider stream to turn compacted history into a structured
|
||||
// checkpoint summary. It mirrors pi's harness/compaction/compaction.ts prompts
|
||||
// and utils.ts helpers, adapted to pigo's flat message list and Provider stream.
|
||||
package compaction
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
"github.com/smallnest/pigo/internal/provider"
|
||||
)
|
||||
|
||||
// SUMMARIZATION_SYSTEM_PROMPT instructs the model to only emit the structured
|
||||
// summary and never continue the conversation. Verbatim from pi.
|
||||
const SUMMARIZATION_SYSTEM_PROMPT = `You are a context summarization assistant. Your task is to read a conversation between a user and an AI assistant, then produce a structured summary following the exact format specified.
|
||||
|
||||
Do NOT continue the conversation. Do NOT respond to any questions in the conversation. ONLY output the structured summary.`
|
||||
|
||||
// summarizationPrompt is the first-time summary template (pi's SUMMARIZATION_PROMPT).
|
||||
const summarizationPrompt = `The messages above are a conversation to summarize. Create a structured context checkpoint summary that another LLM will use to continue the work.
|
||||
|
||||
Use this EXACT format:
|
||||
|
||||
## Goal
|
||||
[What is the user trying to accomplish? Can be multiple items if the session covers different tasks.]
|
||||
|
||||
## Constraints & Preferences
|
||||
- [Any constraints, preferences, or requirements mentioned by user]
|
||||
- [Or "(none)" if none were mentioned]
|
||||
|
||||
## Progress
|
||||
### Done
|
||||
- [x] [Completed tasks/changes]
|
||||
|
||||
### In Progress
|
||||
- [ ] [Current work]
|
||||
|
||||
### Blocked
|
||||
- [Issues preventing progress, if any]
|
||||
|
||||
## Key Decisions
|
||||
- **[Decision]**: [Brief rationale]
|
||||
|
||||
## Next Steps
|
||||
1. [Ordered list of what should happen next]
|
||||
|
||||
## Critical Context
|
||||
- [Any data, examples, or references needed to continue]
|
||||
- [Or "(none)" if not applicable]
|
||||
|
||||
Keep each section concise. Preserve exact file paths, function names, and error messages.`
|
||||
|
||||
// updateSummarizationPrompt incorporates new messages into an existing summary
|
||||
// (pi's UPDATE_SUMMARIZATION_PROMPT).
|
||||
const updateSummarizationPrompt = `The messages above are NEW conversation messages to incorporate into the existing summary provided in <previous-summary> tags.
|
||||
|
||||
Update the existing structured summary with new information. RULES:
|
||||
- PRESERVE all existing information from the previous summary
|
||||
- ADD new progress, decisions, and context from the new messages
|
||||
- UPDATE the Progress section: move items from "In Progress" to "Done" when completed
|
||||
- UPDATE "Next Steps" based on what was accomplished
|
||||
- PRESERVE exact file paths, function names, and error messages
|
||||
- If something is no longer relevant, you may remove it
|
||||
|
||||
Use this EXACT format:
|
||||
|
||||
## Goal
|
||||
[Preserve existing goals, add new ones if the task expanded]
|
||||
|
||||
## Constraints & Preferences
|
||||
- [Preserve existing, add new ones discovered]
|
||||
|
||||
## Progress
|
||||
### Done
|
||||
- [x] [Include previously done items AND newly completed items]
|
||||
|
||||
### In Progress
|
||||
- [ ] [Current work - update based on progress]
|
||||
|
||||
### Blocked
|
||||
- [Current blockers - remove if resolved]
|
||||
|
||||
## Key Decisions
|
||||
- **[Decision]**: [Brief rationale] (preserve all previous, add new)
|
||||
|
||||
## Next Steps
|
||||
1. [Update based on current state]
|
||||
|
||||
## Critical Context
|
||||
- [Preserve important context, add new if needed]
|
||||
|
||||
Keep each section concise. Preserve exact file paths, function names, and error messages.`
|
||||
|
||||
// FileOperations accumulates the file paths a compaction range touched, split
|
||||
// by access kind. Mirrors pi's FileOperations.
|
||||
type FileOperations struct {
|
||||
// Read holds files read but not necessarily modified.
|
||||
Read map[string]struct{}
|
||||
// Written holds files written by full-file write operations.
|
||||
Written map[string]struct{}
|
||||
// Edited holds files modified by edit operations.
|
||||
Edited map[string]struct{}
|
||||
}
|
||||
|
||||
// NewFileOps returns an empty file-operation accumulator.
|
||||
func NewFileOps() FileOperations {
|
||||
return FileOperations{
|
||||
Read: map[string]struct{}{},
|
||||
Written: map[string]struct{}{},
|
||||
Edited: map[string]struct{}{},
|
||||
}
|
||||
}
|
||||
|
||||
// extractFileOpsFromMessage adds file operations from an assistant message's
|
||||
// tool calls to the accumulator, keyed on the tool name (read/write/edit) and
|
||||
// the string "path" argument. Non-assistant messages are ignored, matching pi.
|
||||
func extractFileOpsFromMessage(msg agentcore.Message, ops FileOperations) {
|
||||
a, ok := msg.(agentcore.AssistantMessage)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
for _, call := range a.ToolCalls() {
|
||||
path := toolCallPath(call.Arguments)
|
||||
if path == "" {
|
||||
continue
|
||||
}
|
||||
switch call.Name {
|
||||
case "read":
|
||||
ops.Read[path] = struct{}{}
|
||||
case "write":
|
||||
ops.Written[path] = struct{}{}
|
||||
case "edit":
|
||||
ops.Edited[path] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// toolCallPath extracts a string "path" argument from a tool call's raw
|
||||
// arguments, returning "" when absent or not a string.
|
||||
func toolCallPath(args json.RawMessage) string {
|
||||
if len(args) == 0 {
|
||||
return ""
|
||||
}
|
||||
var decoded struct {
|
||||
Path string `json:"path"`
|
||||
}
|
||||
if err := json.Unmarshal(args, &decoded); err != nil {
|
||||
return ""
|
||||
}
|
||||
return decoded.Path
|
||||
}
|
||||
|
||||
// computeFileLists derives sorted read-only and modified (edited∪written) file
|
||||
// lists, excluding modified files from the read-only list. Mirrors pi.
|
||||
func computeFileLists(ops FileOperations) (readFiles, modifiedFiles []string) {
|
||||
modified := map[string]struct{}{}
|
||||
for f := range ops.Edited {
|
||||
modified[f] = struct{}{}
|
||||
}
|
||||
for f := range ops.Written {
|
||||
modified[f] = struct{}{}
|
||||
}
|
||||
for f := range ops.Read {
|
||||
if _, isMod := modified[f]; !isMod {
|
||||
readFiles = append(readFiles, f)
|
||||
}
|
||||
}
|
||||
for f := range modified {
|
||||
modifiedFiles = append(modifiedFiles, f)
|
||||
}
|
||||
sort.Strings(readFiles)
|
||||
sort.Strings(modifiedFiles)
|
||||
return readFiles, modifiedFiles
|
||||
}
|
||||
|
||||
// formatFileOperations renders the file lists as <read-files>/<modified-files>
|
||||
// metadata blocks appended to a summary, or "" when both lists are empty.
|
||||
func formatFileOperations(readFiles, modifiedFiles []string) string {
|
||||
var sections []string
|
||||
if len(readFiles) > 0 {
|
||||
sections = append(sections, "<read-files>\n"+strings.Join(readFiles, "\n")+"\n</read-files>")
|
||||
}
|
||||
if len(modifiedFiles) > 0 {
|
||||
sections = append(sections, "<modified-files>\n"+strings.Join(modifiedFiles, "\n")+"\n</modified-files>")
|
||||
}
|
||||
if len(sections) == 0 {
|
||||
return ""
|
||||
}
|
||||
return "\n\n" + strings.Join(sections, "\n\n")
|
||||
}
|
||||
|
||||
// toolResultMaxChars caps a tool result's serialized text in the summarization
|
||||
// prompt, matching pi's TOOL_RESULT_MAX_CHARS.
|
||||
const toolResultMaxChars = 2000
|
||||
|
||||
// truncateForSummary caps text at maxChars, appending a truncation marker.
|
||||
func truncateForSummary(text string, maxChars int) string {
|
||||
if len(text) <= maxChars {
|
||||
return text
|
||||
}
|
||||
return fmt.Sprintf("%s\n\n[... %d more characters truncated]", text[:maxChars], len(text)-maxChars)
|
||||
}
|
||||
|
||||
// textOf concatenates the text blocks of a content list.
|
||||
func textOf(content agentcore.ContentList) string {
|
||||
var b strings.Builder
|
||||
for _, c := range content {
|
||||
if t, ok := c.(agentcore.TextContent); ok {
|
||||
b.WriteString(t.Text)
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// serializeConversation renders messages as a plain-text transcript for the
|
||||
// summarization prompt, mirroring pi's serializeConversation: user text,
|
||||
// assistant thinking/text/tool-call lines, and truncated tool results.
|
||||
func serializeConversation(msgs []agentcore.Message) string {
|
||||
var parts []string
|
||||
for _, msg := range msgs {
|
||||
switch m := msg.(type) {
|
||||
case agentcore.UserMessage:
|
||||
if s := textOf(m.Content); s != "" {
|
||||
parts = append(parts, "[User]: "+s)
|
||||
}
|
||||
case agentcore.AssistantMessage:
|
||||
var textParts, thinkingParts, toolCalls []string
|
||||
for _, block := range m.Content {
|
||||
switch c := block.(type) {
|
||||
case agentcore.TextContent:
|
||||
textParts = append(textParts, c.Text)
|
||||
case agentcore.ThinkingContent:
|
||||
thinkingParts = append(thinkingParts, c.Thinking)
|
||||
case agentcore.ToolCallContent:
|
||||
toolCalls = append(toolCalls, formatToolCall(c))
|
||||
}
|
||||
}
|
||||
if len(thinkingParts) > 0 {
|
||||
parts = append(parts, "[Assistant thinking]: "+strings.Join(thinkingParts, "\n"))
|
||||
}
|
||||
if len(textParts) > 0 {
|
||||
parts = append(parts, "[Assistant]: "+strings.Join(textParts, "\n"))
|
||||
}
|
||||
if len(toolCalls) > 0 {
|
||||
parts = append(parts, "[Assistant tool calls]: "+strings.Join(toolCalls, "; "))
|
||||
}
|
||||
case agentcore.ToolResultMessage:
|
||||
if s := textOf(m.Content); s != "" {
|
||||
parts = append(parts, "[Tool result]: "+truncateForSummary(s, toolResultMaxChars))
|
||||
}
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, "\n\n")
|
||||
}
|
||||
|
||||
// formatToolCall renders a tool call as name(key=value, ...) with each value
|
||||
// JSON-encoded, mirroring pi's serialization of tool-call arguments.
|
||||
func formatToolCall(c agentcore.ToolCallContent) string {
|
||||
if len(c.Arguments) == 0 {
|
||||
return c.Name + "()"
|
||||
}
|
||||
var m map[string]json.RawMessage
|
||||
if err := json.Unmarshal(c.Arguments, &m); err != nil {
|
||||
// Not an object: fall back to the raw argument text.
|
||||
return fmt.Sprintf("%s(%s)", c.Name, string(c.Arguments))
|
||||
}
|
||||
keys := make([]string, 0, len(m))
|
||||
for k := range m {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
sort.Strings(keys) // deterministic order (Go map iteration is random)
|
||||
pairs := make([]string, 0, len(keys))
|
||||
for _, k := range keys {
|
||||
pairs = append(pairs, k+"="+string(m[k]))
|
||||
}
|
||||
return fmt.Sprintf("%s(%s)", c.Name, strings.Join(pairs, ", "))
|
||||
}
|
||||
|
||||
// GenerateSummary drives a provider stream to summarize msgs into a structured
|
||||
// checkpoint. When previousSummary is non-empty it uses the update template and
|
||||
// embeds the prior summary in <previous-summary> tags; otherwise it uses the
|
||||
// first-time template. maxTokens is bounded to min(0.8*reserveTokens,
|
||||
// model.MaxOutputTokens) as in pi. The returned text is the concatenation of
|
||||
// the assistant response's text blocks. A terminal error/aborted response is
|
||||
// surfaced as an error.
|
||||
func GenerateSummary(
|
||||
ctx context.Context,
|
||||
stream provider.StreamFn,
|
||||
model provider.Model,
|
||||
msgs []agentcore.Message,
|
||||
reserveTokens int,
|
||||
previousSummary string,
|
||||
cfg provider.StreamConfig,
|
||||
) (string, error) {
|
||||
base := summarizationPrompt
|
||||
if previousSummary != "" {
|
||||
base = updateSummarizationPrompt
|
||||
}
|
||||
|
||||
conversation := serializeConversation(msgs)
|
||||
var b strings.Builder
|
||||
b.WriteString("<conversation>\n")
|
||||
b.WriteString(conversation)
|
||||
b.WriteString("\n</conversation>\n\n")
|
||||
if previousSummary != "" {
|
||||
b.WriteString("<previous-summary>\n")
|
||||
b.WriteString(previousSummary)
|
||||
b.WriteString("\n</previous-summary>\n\n")
|
||||
}
|
||||
b.WriteString(base)
|
||||
|
||||
promptMsg := agentcore.UserMessage{
|
||||
RoleField: agentcore.RoleUser,
|
||||
Content: agentcore.ContentList{agentcore.NewTextContent(b.String())},
|
||||
}
|
||||
|
||||
// Bound the summary output to min(0.8*reserveTokens, model max output).
|
||||
maxTokens := (reserveTokens * 8) / 10
|
||||
if model.MaxOutputTokens > 0 && model.MaxOutputTokens < maxTokens {
|
||||
maxTokens = model.MaxOutputTokens
|
||||
}
|
||||
extra := map[string]any{}
|
||||
for k, v := range cfg.Extra {
|
||||
extra[k] = v
|
||||
}
|
||||
if maxTokens > 0 {
|
||||
extra["max_tokens"] = maxTokens
|
||||
}
|
||||
cfg.Extra = extra
|
||||
|
||||
llm := provider.LlmContext{
|
||||
SystemPrompt: SUMMARIZATION_SYSTEM_PROMPT,
|
||||
Messages: agentcore.MessageList{promptMsg},
|
||||
}
|
||||
s, err := stream(ctx, model.ID, llm, cfg)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("compaction: build summary stream: %w", err)
|
||||
}
|
||||
|
||||
// Drain events so the stream's result is populated, then read the final
|
||||
// message. Failures ride the stream as a terminal message per the provider
|
||||
// contract, so inspect StopReason rather than only the returned error.
|
||||
for range s.Events() {
|
||||
}
|
||||
final, resErr := s.Result(ctx)
|
||||
if resErr != nil {
|
||||
return "", fmt.Errorf("compaction: summary stream: %w", resErr)
|
||||
}
|
||||
switch final.StopReason {
|
||||
case agentcore.StopReasonAborted:
|
||||
return "", fmt.Errorf("compaction: summarization aborted: %s", final.ErrorMessage)
|
||||
case agentcore.StopReasonError:
|
||||
return "", fmt.Errorf("compaction: summarization failed: %s", final.ErrorMessage)
|
||||
}
|
||||
|
||||
summary := textOf(final.Content)
|
||||
if strings.TrimSpace(summary) == "" {
|
||||
return "", fmt.Errorf("compaction: summarization produced empty output")
|
||||
}
|
||||
return summary, nil
|
||||
}
|
||||
Reference in New Issue
Block a user