mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-16 05:03:26 +00:00
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
parent
fd13b53b1c
commit
5cbaf2be88
2 changed files with 190 additions and 31 deletions
|
|
@ -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" };
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue