121 lines
3.6 KiB
TypeScript
121 lines
3.6 KiB
TypeScript
import { generateText, type ModelMessage } from "ai"
|
|
import { Session } from "."
|
|
import { Identifier } from "../id/id"
|
|
import { Instance } from "../project/instance"
|
|
import { Provider } from "../provider/provider"
|
|
import { defer } from "../util/defer"
|
|
import { MessageV2 } from "./message-v2"
|
|
import { SystemPrompt } from "./system"
|
|
import { Bus } from "../bus"
|
|
import z from "zod"
|
|
import type { ModelsDev } from "../provider/models"
|
|
import { SessionPrompt } from "./prompt"
|
|
|
|
export namespace SessionCompaction {
|
|
export const Event = {
|
|
Compacted: Bus.event(
|
|
"session.compacted",
|
|
z.object({
|
|
sessionID: z.string(),
|
|
}),
|
|
),
|
|
}
|
|
|
|
export function isOverflow(input: { tokens: MessageV2.Assistant["tokens"]; model: ModelsDev.Model }) {
|
|
const count = input.tokens.input + input.tokens.cache.read + input.tokens.output
|
|
const output = Math.min(input.model.limit.output, SessionPrompt.OUTPUT_TOKEN_MAX) || SessionPrompt.OUTPUT_TOKEN_MAX
|
|
const usable = input.model.limit.context - output
|
|
return count > usable
|
|
}
|
|
|
|
export async function run(input: { sessionID: string; providerID: string; modelID: string }) {
|
|
await Session.update(input.sessionID, (draft) => {
|
|
draft.time.compacting = Date.now()
|
|
})
|
|
await using _ = defer(async () => {
|
|
await Session.update(input.sessionID, (draft) => {
|
|
draft.time.compacting = undefined
|
|
})
|
|
})
|
|
const toSummarize = await Session.messages(input.sessionID).then(MessageV2.filterSummarized)
|
|
const model = await Provider.getModel(input.providerID, input.modelID)
|
|
const system = [
|
|
...SystemPrompt.summarize(model.providerID),
|
|
...(await SystemPrompt.environment()),
|
|
...(await SystemPrompt.custom()),
|
|
]
|
|
|
|
const msg = (await Session.updateMessage({
|
|
id: Identifier.ascending("message"),
|
|
role: "assistant",
|
|
sessionID: input.sessionID,
|
|
system,
|
|
mode: "build",
|
|
path: {
|
|
cwd: Instance.directory,
|
|
root: Instance.worktree,
|
|
},
|
|
cost: 0,
|
|
tokens: {
|
|
output: 0,
|
|
input: 0,
|
|
reasoning: 0,
|
|
cache: { read: 0, write: 0 },
|
|
},
|
|
modelID: input.modelID,
|
|
providerID: model.providerID,
|
|
time: {
|
|
created: Date.now(),
|
|
},
|
|
})) as MessageV2.Assistant
|
|
const generated = await generateText({
|
|
maxRetries: 10,
|
|
model: model.language,
|
|
messages: [
|
|
...system.map(
|
|
(x): ModelMessage => ({
|
|
role: "system",
|
|
content: x,
|
|
}),
|
|
),
|
|
...MessageV2.toModelMessage(toSummarize),
|
|
{
|
|
role: "user",
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text: "Provide a detailed but concise summary of our conversation above. Focus on information that would be helpful for continuing the conversation, including what we did, what we're doing, which files we're working on, and what we're going to do next.",
|
|
},
|
|
],
|
|
},
|
|
],
|
|
})
|
|
const usage = Session.getUsage(model.info, generated.usage, generated.providerMetadata)
|
|
msg.cost += usage.cost
|
|
msg.tokens = usage.tokens
|
|
msg.summary = true
|
|
msg.time.completed = Date.now()
|
|
await Session.updateMessage(msg)
|
|
const part = await Session.updatePart({
|
|
type: "text",
|
|
sessionID: input.sessionID,
|
|
messageID: msg.id,
|
|
id: Identifier.ascending("part"),
|
|
text: generated.text,
|
|
time: {
|
|
start: Date.now(),
|
|
end: Date.now(),
|
|
},
|
|
})
|
|
|
|
Bus.publish(Event.Compacted, {
|
|
sessionID: input.sessionID,
|
|
})
|
|
|
|
return {
|
|
info: msg,
|
|
parts: [part],
|
|
}
|
|
}
|
|
}
|