Improve branch summarization with preparation and file ops extraction

- Add prepareBranchEntries() to extract messages and file operations
- Extract read/write/edit operations from tool calls
- Append static file operations section to summary
- Improve prompt for branch summarization
- Skip toolResult messages (context already in assistant message)
- Export new types: BranchPreparation, FileOperations
This commit is contained in:
Mario Zechner 2025-12-29 21:05:23 +01:00
parent fd13b53b1c
commit 5cbaf2be88
2 changed files with 190 additions and 31 deletions

View file

@ -9,8 +9,9 @@ import type { Model } from "@mariozechner/pi-ai";
import { complete } from "@mariozechner/pi-ai";
import type { SessionEntry } from "../session-manager.js";
const DEFAULT_INSTRUCTIONS =
"Summarize this conversation branch concisely, capturing key decisions, actions taken, and outcomes.";
// ============================================================================
// Types
// ============================================================================
export interface BranchSummaryResult {
summary?: string;
@ -18,6 +19,25 @@ export interface BranchSummaryResult {
error?: string;
}
export interface FileOperations {
read: Set<string>;
written: Set<string>;
edited: Set<string>;
}
export interface BranchPreparation {
/** Messages extracted for summarization */
messages: Array<{ role: string; content: string }>;
/** File operations extracted from tool calls */
fileOps: FileOperations;
/** Previous summaries found in entries */
previousSummaries: string[];
}
// ============================================================================
// Entry Parsing
// ============================================================================
/**
* Extract text content from any message type.
*/
@ -33,6 +53,152 @@ function extractMessageText(message: any): string {
return "";
}
/**
* Extract file operations from tool calls in an assistant message.
*/
function extractFileOpsFromToolCalls(message: any, fileOps: FileOperations): void {
if (!message.content || !Array.isArray(message.content)) return;
for (const block of message.content) {
if (typeof block !== "object" || block === null) continue;
if (block.type !== "toolCall") continue;
const args = block.arguments as Record<string, unknown> | undefined;
if (!args) continue;
const path = typeof args.path === "string" ? args.path : undefined;
if (!path) continue;
switch (block.name) {
case "read":
fileOps.read.add(path);
break;
case "write":
fileOps.written.add(path);
break;
case "edit":
fileOps.edited.add(path);
break;
}
}
}
/**
* Prepare entries for summarization.
*
* Extracts:
* - Messages (user, assistant text, custom_message)
* - File operations from tool calls
* - Previous branch summaries
*
* Skips:
* - toolResult messages (context already in assistant message)
* - thinking_level_change, model_change, custom, label entries
* - compaction entries (these are boundaries, shouldn't be in the input)
*/
export function prepareBranchEntries(entries: SessionEntry[]): BranchPreparation {
const messages: Array<{ role: string; content: string }> = [];
const fileOps: FileOperations = {
read: new Set(),
written: new Set(),
edited: new Set(),
};
const previousSummaries: string[] = [];
for (const entry of entries) {
switch (entry.type) {
case "message": {
const role = entry.message.role;
// Skip tool results - the context is in the assistant's tool call
if (role === "toolResult") continue;
// Extract file ops from assistant tool calls
if (role === "assistant") {
extractFileOpsFromToolCalls(entry.message, fileOps);
}
// Extract text content
const text = extractMessageText(entry.message);
if (text) {
messages.push({ role, content: text });
}
break;
}
case "custom_message": {
const text =
typeof entry.content === "string"
? entry.content
: entry.content
.filter((c): c is { type: "text"; text: string } => c.type === "text")
.map((c) => c.text)
.join("");
if (text) {
messages.push({ role: "user", content: text });
}
break;
}
case "branch_summary": {
previousSummaries.push(entry.summary);
break;
}
// Skip these entry types - they don't contribute to conversation content
case "compaction":
case "thinking_level_change":
case "model_change":
case "custom":
case "label":
break;
}
}
return { messages, fileOps, previousSummaries };
}
// ============================================================================
// Summary Generation
// ============================================================================
const BRANCH_SUMMARY_PROMPT = `Summarize this conversation branch concisely for context when returning later:
- Key decisions made and actions taken
- Important context, constraints, or preferences discovered
- Current state and any pending work
- Critical information needed to continue from a different point
Be brief and focused on what matters for future reference.`;
/**
* Format file operations as a static section to append to summary.
*/
function formatFileOperations(fileOps: FileOperations): string {
const sections: string[] = [];
if (fileOps.read.size > 0) {
const files = [...fileOps.read].sort();
sections.push(`**Read:** ${files.join(", ")}`);
}
if (fileOps.edited.size > 0) {
const files = [...fileOps.edited].sort();
sections.push(`**Edited:** ${files.join(", ")}`);
}
if (fileOps.written.size > 0) {
// Exclude files that were also edited (edit implies write)
const writtenOnly = [...fileOps.written].filter((f) => !fileOps.edited.has(f)).sort();
if (writtenOnly.length > 0) {
sections.push(`**Created:** ${writtenOnly.join(", ")}`);
}
}
if (sections.length === 0) return "";
return `\n\n---\n**Files:**\n${sections.join("\n")}`;
}
/**
* Generate a summary of abandoned branch entries.
*
@ -49,39 +215,26 @@ export async function generateBranchSummary(
signal: AbortSignal,
customInstructions?: string,
): Promise<BranchSummaryResult> {
// Convert entries to messages for summarization
const messages: Array<{ role: string; content: string }> = [];
for (const entry of entries) {
if (entry.type === "message") {
const text = extractMessageText(entry.message);
if (text) {
messages.push({ role: entry.message.role, content: text });
}
} else if (entry.type === "custom_message") {
const text =
typeof entry.content === "string"
? entry.content
: entry.content
.filter((c): c is { type: "text"; text: string } => c.type === "text")
.map((c) => c.text)
.join("");
if (text) {
messages.push({ role: "user", content: text });
}
} else if (entry.type === "branch_summary") {
messages.push({ role: "system", content: `[Previous branch summary: ${entry.summary}]` });
}
}
const { messages, fileOps, previousSummaries } = prepareBranchEntries(entries);
if (messages.length === 0) {
return { summary: "No content to summarize" };
}
// Build prompt for summarization
const conversationText = messages.map((m) => `${m.role}: ${m.content}`).join("\n\n");
const instructions = customInstructions ? `${customInstructions}\n\n` : `${DEFAULT_INSTRUCTIONS}\n\n`;
const prompt = `${instructions}Conversation:\n${conversationText}`;
// Build conversation text
const parts: string[] = [];
// Include previous summaries as context
if (previousSummaries.length > 0) {
parts.push(`[Previous context: ${previousSummaries.join(" | ")}]`);
}
// Add conversation
parts.push(messages.map((m) => `${m.role}: ${m.content}`).join("\n\n"));
const conversationText = parts.join("\n\n");
const instructions = customInstructions || BRANCH_SUMMARY_PROMPT;
const prompt = `${instructions}\n\nConversation:\n${conversationText}`;
// Call LLM for summarization
const response = await complete(
@ -106,10 +259,13 @@ export async function generateBranchSummary(
return { error: response.errorMessage || "Summarization failed" };
}
const summary = response.content
let summary = response.content
.filter((c): c is { type: "text"; text: string } => c.type === "text")
.map((c) => c.text)
.join("\n");
// Append static file operations section
summary += formatFileOperations(fileOps);
return { summary: summary || "No summary generated" };
}

View file

@ -12,6 +12,7 @@ export {
export { type ApiKeyCredential, type AuthCredential, AuthStorage, type OAuthCredential } from "./core/auth-storage.js";
// Compaction
export {
type BranchPreparation,
type BranchSummaryResult,
type CompactionResult,
type CutPointResult,
@ -19,11 +20,13 @@ export {
compact,
DEFAULT_COMPACTION_SETTINGS,
estimateTokens,
type FileOperations,
findCutPoint,
findTurnStartIndex,
generateBranchSummary,
generateSummary,
getLastAssistantUsage,
prepareBranchEntries,
shouldCompact,
} from "./core/compaction/index.js";
// Custom tools