|
import { readFileSync } from "node:fs"; |
|
import { homedir } from "node:os"; |
|
import { join } from "node:path"; |
|
import type { AgentMessage } from "@earendil-works/pi-agent-core"; |
|
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; |
|
|
|
type JsonObject = Record<string, unknown>; |
|
|
|
type ImageLimitRule = { |
|
provider: string; |
|
models: string[]; |
|
maxImages: number; |
|
}; |
|
|
|
export type ImageContextConfig = { |
|
defaultMaxImages: number | null; |
|
rules: ImageLimitRule[]; |
|
placeholder: string; |
|
}; |
|
|
|
type ImageBlock = { |
|
type: "image"; |
|
}; |
|
|
|
type ContentBlock = { |
|
type?: string; |
|
text?: string; |
|
}; |
|
|
|
type MessageWithContent = AgentMessage & { |
|
content?: string | ContentBlock[]; |
|
timestamp?: number; |
|
}; |
|
|
|
type PruneResult<T> = { |
|
messages: T[]; |
|
totalImages: number; |
|
keptImages: number; |
|
prunedImages: number; |
|
}; |
|
|
|
const CONFIG_PATH = join(homedir(), ".pi", "agent", "image-context.json"); |
|
const STATE_TYPE = "image-context"; |
|
const LEGACY_STATE_TYPE = "qwen-image-context"; |
|
const STATUS_KEY = "image-context"; |
|
|
|
function isObject(value: unknown): value is JsonObject { |
|
return typeof value === "object" && value !== null && !Array.isArray(value); |
|
} |
|
|
|
function parseImageLimit(value: unknown, name: string): number | null { |
|
if (value === null) return null; |
|
if (typeof value !== "number" || !Number.isInteger(value) || value < 0) { |
|
throw new Error(image-context: ${name} must be null or a non-negative integer); |
|
} |
|
return value; |
|
} |
|
|
|
function parseRule(value: unknown, index: number): ImageLimitRule { |
|
if (!isObject(value) || typeof value.provider !== "string") { |
|
throw new Error(image-context: rules[${index}].provider must be a string); |
|
} |
|
const models = Array.isArray(value.models) |
|
? value.models.filter((model): model is string => typeof model === "string") |
|
: []; |
|
if (models.length === 0) { |
|
throw new Error(image-context: rules[${index}].models must contain a model id or *); |
|
} |
|
const maxImages = parseImageLimit(value.maxImages, rules[${index}].maxImages); |
|
if (maxImages === null) { |
|
throw new Error(image-context: rules[${index}].maxImages cannot be null); |
|
} |
|
return { provider: value.provider, models, maxImages }; |
|
} |
|
|
|
function loadConfig(): ImageContextConfig { |
|
const raw: unknown = JSON.parse(readFileSync(CONFIG_PATH, "utf8")); |
|
if (!isObject(raw)) throw new Error("image-context: config must be an object"); |
|
if (typeof raw.placeholder !== "string" || raw.placeholder.length === 0) { |
|
throw new Error("image-context: placeholder must be a non-empty string"); |
|
} |
|
if (!Array.isArray(raw.rules)) { |
|
throw new Error("image-context: rules must be an array"); |
|
} |
|
return { |
|
defaultMaxImages: parseImageLimit(raw.defaultMaxImages, "defaultMaxImages"), |
|
rules: raw.rules.map(parseRule), |
|
placeholder: raw.placeholder, |
|
}; |
|
} |
|
|
|
export function resolveImageLimit( |
|
config: ImageContextConfig, |
|
provider: string | undefined, |
|
model: string | undefined, |
|
): number | null { |
|
if (!provider || !model) return config.defaultMaxImages; |
|
const rule = config.rules.find( |
|
(candidate) => |
|
(candidate.provider === "" || candidate.provider === provider) && |
|
(candidate.models.includes("") || candidate.models.includes(model)), |
|
); |
|
return rule?.maxImages ?? config.defaultMaxImages; |
|
} |
|
|
|
function isImageBlock(value: unknown): value is ImageBlock { |
|
return isObject(value) && value.type === "image"; |
|
} |
|
|
|
function messageTimestamp(message: MessageWithContent): number | undefined { |
|
return typeof message.timestamp === "number" ? message.timestamp : undefined; |
|
} |
|
|
|
export function pruneImageContext<T extends MessageWithContent>( |
|
messages: T[], |
|
maxImages: number, |
|
clearBeforeTimestamp: number, |
|
placeholder: string, |
|
): PruneResult<T> { |
|
let totalImages = 0; |
|
let keptImages = 0; |
|
let prunedImages = 0; |
|
let remaining = maxImages; |
|
|
|
for (let messageIndex = messages.length - 1; messageIndex >= 0; messageIndex--) { |
|
const message = messages[messageIndex]; |
|
if (!message || !Array.isArray(message.content)) continue; |
|
const timestamp = messageTimestamp(message); |
|
const explicitlyCleared = |
|
clearBeforeTimestamp > 0 && |
|
(timestamp === undefined || timestamp <= clearBeforeTimestamp); |
|
|
|
for (let blockIndex = 0; blockIndex < message.content.length; blockIndex++) { |
|
const block = message.content[blockIndex]; |
|
if (!isImageBlock(block)) continue; |
|
totalImages++; |
|
if (!explicitlyCleared && remaining > 0) { |
|
remaining--; |
|
keptImages++; |
|
continue; |
|
} |
|
message.content[blockIndex] = { type: "text", text: placeholder }; |
|
prunedImages++; |
|
} |
|
} |
|
|
|
return { messages, totalImages, keptImages, prunedImages }; |
|
} |
|
|
|
export default function imageContext(pi: ExtensionAPI): void { |
|
let config = loadConfig(); |
|
let clearBeforeTimestamp = 0; |
|
|
|
pi.on("session_start", async (_event, ctx) => { |
|
clearBeforeTimestamp = 0; |
|
for (const entry of ctx.sessionManager.getEntries()) { |
|
if ( |
|
entry.type === "custom" && |
|
(entry.customType === STATE_TYPE || entry.customType === LEGACY_STATE_TYPE) && |
|
isObject(entry.data) && |
|
typeof entry.data.clearBeforeTimestamp === "number" |
|
) { |
|
clearBeforeTimestamp = entry.data.clearBeforeTimestamp; |
|
} |
|
} |
|
ctx.ui.setStatus(STATUS_KEY, undefined); |
|
}); |
|
|
|
pi.on("context", async (event, ctx) => { |
|
try { |
|
config = loadConfig(); |
|
} catch (error) { |
|
console.error(error); |
|
} |
|
|
|
const limit = resolveImageLimit(config, ctx.model?.provider, ctx.model?.id); |
|
if (limit === null && clearBeforeTimestamp === 0) return; |
|
|
|
const result = pruneImageContext( |
|
event.messages as MessageWithContent[], |
|
limit ?? Number.MAX_SAFE_INTEGER, |
|
clearBeforeTimestamp, |
|
config.placeholder, |
|
); |
|
if (result.totalImages === 0) { |
|
ctx.ui.setStatus(STATUS_KEY, undefined); |
|
return { messages: result.messages }; |
|
} |
|
|
|
const status = |
|
result.prunedImages > 0 |
|
? images ${result.keptImages}/${result.totalImages} |
|
: images ${result.keptImages}; |
|
ctx.ui.setStatus(STATUS_KEY, ctx.ui.theme.fg("muted", status)); |
|
return { messages: result.messages }; |
|
}); |
|
|
|
pi.registerCommand("clear-images", { |
|
description: "Exclude every image already in this session from future LLM context", |
|
handler: async (_args, ctx) => { |
|
clearBeforeTimestamp = Date.now(); |
|
pi.appendEntry(STATE_TYPE, { clearBeforeTimestamp }); |
|
ctx.ui.notify( |
|
"Existing session images will no longer be sent to the model. New images follow the active model limit.", |
|
"info", |
|
); |
|
}, |
|
}); |
|
|
|
pi.registerCommand("image-context", { |
|
description: "Show the active model image-history limit and clear cutoff", |
|
handler: async (_args, ctx) => { |
|
config = loadConfig(); |
|
const provider = ctx.model?.provider; |
|
const model = ctx.model?.id; |
|
const limit = resolveImageLimit(config, provider, model); |
|
ctx.ui.notify( |
|
${provider ?? "?"}/${model ?? "?"}: + |
|
maxImages=${limit ?? "unlimited"}; + |
|
clearBefore=${clearBeforeTimestamp || "none"}, |
|
"info", |
|
); |
|
}, |
|
}); |
|
} |
Agents Shouldn't Blink