Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
6.4 kB · 194 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195/** * Truncated Tool Example - Demonstrates proper output truncation for custom tools * * Custom tools MUST truncate their output to avoid overwhelming the LLM context. * The built-in limit is 50KB (~10k tokens) and 2000 lines, whichever is hit first. * * This example shows how to: * 1. Use the built-in truncation utilities * 2. Write full output to a temp file when truncated * 3. Inform the LLM where to find the complete output * 4. Custom rendering of tool calls and results * * The `rg` tool here wraps ripgrep with proper truncation and custom rendering. */
import { mkdtemp, writeFile } from "node:fs/promises";import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, formatSize, type TruncationResult, truncateHead, withFileMutationQueue,} from "@earendil-works/pi-coding-agent";import { Text } from "@earendil-works/pi-tui";import { execSync } from "child_process";import { tmpdir } from "os";import { join } from "path";import { Type } from "typebox";
const RgParams = Type.Object({ pattern: Type.String({ description: "Search pattern (regex)" }), path: Type.Optional(Type.String({ description: "Directory to search (default: current directory)" })), glob: Type.Optional(Type.String({ description: "File glob pattern, e.g. '*.ts'" })),});
interface RgDetails { pattern: string; path?: string; glob?: string; matchCount: number; truncation?: TruncationResult; fullOutputPath?: string;}
export default function (pi: ExtensionAPI) { pi.registerTool({ name: "rg", label: "ripgrep", // Document the truncation limits in the tool description so the LLM knows description: `Search file contents using ripgrep. Output is truncated to ${DEFAULT_MAX_LINES} lines or ${formatSize(DEFAULT_MAX_BYTES)} (whichever is hit first). If truncated, full output is saved to a temp file.`, parameters: RgParams,
async execute(_toolCallId, params, _signal, _onUpdate, ctx) { const { pattern, path: searchPath, glob } = params;
// Build the ripgrep command const args = ["rg", "--line-number", "--color=never"]; if (glob) args.push("--glob", glob); args.push(pattern); args.push(searchPath || ".");
let output: string; try { output = execSync(args.join(" "), { cwd: ctx.cwd, encoding: "utf-8", maxBuffer: 100 * 1024 * 1024, // 100MB buffer to capture full output }); } catch (err: any) { // ripgrep exits with 1 when no matches found if (err.status === 1) { return { content: [{ type: "text", text: "No matches found" }], details: { pattern, path: searchPath, glob, matchCount: 0 } as RgDetails, }; } throw new Error(`ripgrep failed: ${err.message}`); }
if (!output.trim()) { return { content: [{ type: "text", text: "No matches found" }], details: { pattern, path: searchPath, glob, matchCount: 0 } as RgDetails, }; }
// Apply truncation using built-in utilities // truncateHead keeps the first N lines/bytes (good for search results) // truncateTail keeps the last N lines/bytes (good for logs/command output) const truncation = truncateHead(output, { maxLines: DEFAULT_MAX_LINES, maxBytes: DEFAULT_MAX_BYTES, });
// Count matches (each non-empty line with a match) const matchCount = output.split("\n").filter((line) => line.trim()).length;
const details: RgDetails = { pattern, path: searchPath, glob, matchCount, };
let resultText = truncation.content;
if (truncation.truncated) { // Save full output to a temp file so LLM can access it if needed const tempDir = await mkdtemp(join(tmpdir(), "pi-rg-")); const tempFile = join(tempDir, "output.txt"); await withFileMutationQueue(tempFile, async () => { await writeFile(tempFile, output, "utf8"); });
details.truncation = truncation; details.fullOutputPath = tempFile;
// Add truncation notice - this helps the LLM understand the output is incomplete const truncatedLines = truncation.totalLines - truncation.outputLines; const truncatedBytes = truncation.totalBytes - truncation.outputBytes;
resultText += `\n\n[Output truncated: showing ${truncation.outputLines} of ${truncation.totalLines} lines`; resultText += ` (${formatSize(truncation.outputBytes)} of ${formatSize(truncation.totalBytes)}).`; resultText += ` ${truncatedLines} lines (${formatSize(truncatedBytes)}) omitted.`; resultText += ` Full output saved to: ${tempFile}]`; }
return { content: [{ type: "text", text: resultText }], details, }; },
// Custom rendering of the tool call (shown before/during execution) renderCall(args, theme, _context) { let text = theme.fg("toolTitle", theme.bold("rg ")); text += theme.fg("accent", `"${args.pattern}"`); if (args.path) { text += theme.fg("muted", ` in ${args.path}`); } if (args.glob) { text += theme.fg("dim", ` --glob ${args.glob}`); } return new Text(text, 0, 0); },
// Custom rendering of the tool result renderResult(result, { expanded, isPartial }, theme, _context) { const details = result.details as RgDetails | undefined;
// Handle streaming/partial results if (isPartial) { return new Text(theme.fg("warning", "Searching..."), 0, 0); }
// No matches if (!details || details.matchCount === 0) { return new Text(theme.fg("dim", "No matches found"), 0, 0); }
// Build result display let text = theme.fg("success", `${details.matchCount} matches`);
// Show truncation warning if applicable if (details.truncation?.truncated) { text += theme.fg("warning", " (truncated)"); }
// In expanded view, show the actual matches if (expanded) { const content = result.content[0]; if (content?.type === "text") { // Show first 20 lines in expanded view, or all if fewer const lines = content.text.split("\n").slice(0, 20); for (const line of lines) { text += `\n${theme.fg("dim", line)}`; } if (content.text.split("\n").length > 20) { text += `\n${theme.fg("muted", "... (use ipython to see full output)")}`; } }
// Show temp file path if truncated if (details.fullOutputPath) { text += `\n${theme.fg("dim", `Full output: ${details.fullOutputPath}`)}`; } }
return new Text(text, 0, 0); }, });}