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.5 kB · 218 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219import type { AgentMessage } from "@earendil-works/pi-agent-core";import type { AssistantMessage, CacheRetention } from "@earendil-works/pi-ai";
export interface SessionCacheState { provider: string; modelId: string; cachedTokens: number; timestamp: number; ttlSeconds?: number; deterministic: boolean;}
export interface PromptCacheSnapshot { cachedTokens: number; ttlSeconds?: number; expiresAt?: number; deterministic: boolean; status: "warm" | "cold" | "mismatch"; canRestore?: boolean;}
export type CacheTtlStatus = | { status: "none" } | { status: "cold"; cachedTokens?: number } | { status: "mismatch"; cachedModelId: string; canRestore: boolean } | { status: "warm"; cachedTokens: number; ttlSeconds: number; expiresAt: number; remainingSeconds: number; deterministic: boolean; label: string; };
export interface ModelIdentity { provider: string; id: string;}
export function formatCacheTtl(seconds: number): string { if (!Number.isFinite(seconds) || seconds <= 0) return "0s"; const total = Math.floor(seconds); if (total <= 0) return "0s"; if (total < 60) return `${total}s`; const minutes = Math.floor(total / 60); const remSeconds = total % 60; if (total < 3600) { return remSeconds > 0 ? `${minutes}m ${remSeconds}s` : `${minutes}m`; } const hours = Math.floor(total / 3600); const remMinutes = Math.floor((total % 3600) / 60); return remMinutes > 0 ? `${hours}h ${remMinutes}m` : `${hours}h`;}
export function resolvePromptCacheTtl( provider: string, modelId: string, cacheRetention?: CacheRetention, supportsLongCacheRetention = true,): { ttlSeconds: number; deterministic: boolean } | undefined { const prov = provider.toLowerCase(); const id = modelId.toLowerCase(); const isAnthropic = prov === "anthropic" || id.startsWith("claude-"); const isBedrockClaude = prov === "amazon-bedrock" && id.includes("claude"); if (isAnthropic || isBedrockClaude) { const isLong = supportsLongCacheRetention && (cacheRetention === "long" || (typeof process !== "undefined" && process.env.PI_CACHE_RETENTION === "long")); return { ttlSeconds: isLong ? 3600 : 300, deterministic: true }; }
if (prov === "openai-codex") { return { ttlSeconds: 1800, deterministic: true }; }
const isOpenAI = prov === "openai" || id.startsWith("gpt-") || prov.includes("azure-openai"); if (isOpenAI) { const isLong = supportsLongCacheRetention && (cacheRetention === "long" || (typeof process !== "undefined" && process.env.PI_CACHE_RETENTION === "long")); if (isLong) return { ttlSeconds: 86400, deterministic: true }; return { ttlSeconds: 1800, deterministic: true }; }
if (prov === "deepseek" || id.startsWith("deepseek-")) { return { ttlSeconds: 43200, deterministic: false }; }
if (prov.startsWith("google") || id.startsWith("gemini-")) { return { ttlSeconds: 300, deterministic: false }; }
if (prov === "groq") { return { ttlSeconds: 7200, deterministic: false }; }
if (prov === "zai" || prov === "zhipu" || id.startsWith("glm-")) { return { ttlSeconds: 1800, deterministic: false }; }
if (prov === "mistral" || id.startsWith("mistral-") || id.startsWith("codestral-")) { return { ttlSeconds: 1800, deterministic: false }; }
if (prov === "together" || prov === "together-ai") { return { ttlSeconds: 1800, deterministic: false }; }
if (prov === "fireworks") { return { ttlSeconds: 900, deterministic: false }; }
return undefined;}
export function findLatestCompletedAssistant( messages: readonly AgentMessage[], latestCompactionTimestamp?: number,): AssistantMessage | undefined { for (let i = messages.length - 1; i >= 0; i--) { const msg = messages[i]; if (msg.role !== "assistant" || !("usage" in msg)) continue; const assistant = msg as AssistantMessage; if (assistant.stopReason === "aborted" || assistant.stopReason === "error" || !assistant.usage) continue; if (latestCompactionTimestamp !== undefined && assistant.timestamp <= latestCompactionTimestamp) { return undefined; } return assistant; } return undefined;}
export function extractSessionCacheState( messages: readonly AgentMessage[], latestCompactionTimestamp?: number, cacheRetention?: CacheRetention, supportsLongCacheRetention?: boolean,): SessionCacheState | undefined { const assistant = findLatestCompletedAssistant(messages, latestCompactionTimestamp); if (!assistant || !assistant.usage) return undefined; const cachedTokens = (assistant.usage.cacheRead ?? 0) + (assistant.usage.cacheWrite ?? 0); if (cachedTokens <= 0) return undefined;
const provider = assistant.provider || "unknown"; const modelId = assistant.model || "unknown"; const ttlInfo = resolvePromptCacheTtl(provider, modelId, cacheRetention, supportsLongCacheRetention);
return { provider, modelId, cachedTokens, timestamp: assistant.timestamp, ttlSeconds: ttlInfo?.ttlSeconds, deterministic: ttlInfo?.deterministic ?? false, };}
export function evaluateActiveCacheTtl( currentModel: ModelIdentity | undefined, cache: SessionCacheState | undefined, now: number = Date.now(),): CacheTtlStatus { if (!currentModel || !cache || cache.cachedTokens <= 0) { return { status: "none" }; }
if ( currentModel.provider.toLowerCase() !== cache.provider.toLowerCase() || currentModel.id.toLowerCase() !== cache.modelId.toLowerCase() ) { const canRestore = cache.ttlSeconds !== undefined && cache.timestamp + cache.ttlSeconds * 1000 > now; return { status: "mismatch", cachedModelId: cache.modelId, canRestore, }; }
if (cache.ttlSeconds === undefined) { return { status: "warm", cachedTokens: cache.cachedTokens, ttlSeconds: 0, expiresAt: 0, remainingSeconds: 0, deterministic: false, label: "", }; }
const expiresAt = cache.timestamp + cache.ttlSeconds * 1000; const remainingSeconds = Math.max(0, Math.floor((expiresAt - now) / 1000)); if (remainingSeconds <= 0) { return { status: "cold", cachedTokens: cache.cachedTokens }; }
return { status: "warm", cachedTokens: cache.cachedTokens, ttlSeconds: cache.ttlSeconds, expiresAt, remainingSeconds, deterministic: cache.deterministic, label: formatCacheTtl(remainingSeconds), };}
export { assertPromptPayloadWithinCeiling, calculatePromptPayloadCeiling, compilePromptPayload, DEFAULT_PROMPT_PAYLOAD_CEILING_RATIO, estimateMessageTokens, estimatePayloadTokens, formatModelContextCeilingStatement, PromptPayloadCeilingExceededError,} from "./prompts/payload-ceiling.js";