LLM cleanup (#5462)

Co-authored-by: GitHub Action <action@github.com>
Co-authored-by: Aiden Cline <63023139+rekram1-node@users.noreply.github.com>
This commit is contained in:
Dax
2025-12-14 21:11:30 -05:00
committed by GitHub
parent fdf560c343
commit fed4776451
24 changed files with 548 additions and 609 deletions

View File

@@ -5,32 +5,22 @@ import z from "zod"
import { Identifier } from "../id/id"
import { MessageV2 } from "./message-v2"
import { Log } from "../util/log"
import { Flag } from "../flag/flag"
import { SessionRevert } from "./revert"
import { Session } from "."
import { Agent } from "../agent/agent"
import { Provider } from "../provider/provider"
import {
generateText,
type ModelMessage,
type Tool as AITool,
tool,
wrapLanguageModel,
stepCountIs,
jsonSchema,
} from "ai"
import { type Tool as AITool, tool, jsonSchema } from "ai"
import { SessionCompaction } from "./compaction"
import { Instance } from "../project/instance"
import { Bus } from "../bus"
import { ProviderTransform } from "../provider/transform"
import { SystemPrompt } from "./system"
import { Plugin } from "../plugin"
import PROMPT_PLAN from "../session/prompt/plan.txt"
import BUILD_SWITCH from "../session/prompt/build-switch.txt"
import MAX_STEPS from "../session/prompt/max-steps.txt"
import { defer } from "../util/defer"
import { clone, mergeDeep, pipe } from "remeda"
import { mergeDeep, pipe } from "remeda"
import { ToolRegistry } from "../tool/registry"
import { Wildcard } from "../util/wildcard"
import { MCP } from "../mcp"
@@ -44,12 +34,13 @@ import { Command } from "../command"
import { $, fileURLToPath } from "bun"
import { ConfigMarkdown } from "../config/markdown"
import { SessionSummary } from "./summary"
import { Config } from "../config/config"
import { NamedError } from "@opencode-ai/util/error"
import { fn } from "@/util/fn"
import { SessionProcessor } from "./processor"
import { TaskTool } from "@/tool/task"
import { SessionStatus } from "./status"
import { LLM } from "./llm"
import { iife } from "@/util/iife"
import { Shell } from "@/shell/shell"
// @ts-ignore
@@ -96,8 +87,8 @@ export namespace SessionPrompt {
.optional(),
agent: z.string().optional(),
noReply: z.boolean().optional(),
system: z.string().optional(),
tools: z.record(z.string(), z.boolean()).optional(),
system: z.string().optional(),
parts: z.array(
z.discriminatedUnion("type", [
MessageV2.TextPart.omit({
@@ -145,6 +136,20 @@ export namespace SessionPrompt {
})
export type PromptInput = z.infer<typeof PromptInput>
export const prompt = fn(PromptInput, async (input) => {
const session = await Session.get(input.sessionID)
await SessionRevert.cleanup(session)
const message = await createUserMessage(input)
await Session.touch(input.sessionID)
if (input.noReply === true) {
return message
}
return loop(input.sessionID)
})
export async function resolvePromptParts(template: string): Promise<PromptInput["parts"]> {
const parts: PromptInput["parts"] = [
{
@@ -196,20 +201,6 @@ export namespace SessionPrompt {
return parts
}
export const prompt = fn(PromptInput, async (input) => {
const session = await Session.get(input.sessionID)
await SessionRevert.cleanup(session)
const message = await createUserMessage(input)
await Session.touch(input.sessionID)
if (input.noReply === true) {
return message
}
return loop(input.sessionID)
})
function start(sessionID: string) {
const s = state()
if (s[sessionID]) return
@@ -291,7 +282,6 @@ export namespace SessionPrompt {
})
const model = await Provider.getModel(lastUser.model.providerID, lastUser.model.modelID)
const language = await Provider.getLanguage(model)
const task = tasks.pop()
// pending subtask
@@ -304,6 +294,7 @@ export namespace SessionPrompt {
parentID: lastUser.id,
sessionID,
mode: task.agent,
agent: task.agent,
path: {
cwd: Instance.directory,
root: Instance.worktree,
@@ -414,11 +405,6 @@ export namespace SessionPrompt {
messages: msgs,
parentID: lastUser.id,
abort,
agent: lastUser.agent,
model: {
providerID: model.providerID,
modelID: model.id,
},
sessionID,
auto: task.auto,
})
@@ -442,7 +428,6 @@ export namespace SessionPrompt {
}
// normal processing
const cfg = await Config.get()
const agent = await Agent.get(lastUser.agent)
const maxSteps = agent.maxSteps ?? Infinity
const isLastStep = step >= maxSteps
@@ -450,12 +435,14 @@ export namespace SessionPrompt {
messages: msgs,
agent,
})
const processor = SessionProcessor.create({
assistantMessage: (await Session.updateMessage({
id: Identifier.ascending("message"),
parentID: lastUser.id,
role: "assistant",
mode: agent.name,
agent: agent.name,
path: {
cwd: Instance.directory,
root: Instance.worktree,
@@ -478,12 +465,6 @@ export namespace SessionPrompt {
model,
abort,
})
const system = await resolveSystemPrompt({
model,
agent,
system: lastUser.system,
isLastStep,
})
const tools = await resolveTools({
agent,
sessionID,
@@ -491,30 +472,6 @@ export namespace SessionPrompt {
tools: lastUser.tools,
processor,
})
const provider = await Provider.getProvider(model.providerID)
const params = await Plugin.trigger(
"chat.params",
{
sessionID: sessionID,
agent: lastUser.agent,
model: model,
provider,
message: lastUser,
},
{
temperature: model.capabilities.temperature
? (agent.temperature ?? ProviderTransform.temperature(model))
: undefined,
topP: agent.topP ?? ProviderTransform.topP(model),
topK: ProviderTransform.topK(model),
options: pipe(
{},
mergeDeep(ProviderTransform.options(model, sessionID, provider?.options)),
mergeDeep(model.options),
mergeDeep(agent.options),
),
},
)
if (step === 1) {
SessionSummary.summarize({
@@ -523,135 +480,25 @@ export namespace SessionPrompt {
})
}
// Deep copy message history so that modifications made by plugins do not
// affect the original messages
const sessionMessages = clone(
msgs.filter((m) => {
if (m.info.role !== "assistant" || m.info.error === undefined) {
return true
}
if (
MessageV2.AbortedError.isInstance(m.info.error) &&
m.parts.some((part) => part.type !== "step-start" && part.type !== "reasoning")
) {
return true
}
return false
}),
)
await Plugin.trigger("experimental.chat.messages.transform", {}, { messages: sessionMessages })
const messages: ModelMessage[] = [
...system.map(
(x): ModelMessage => ({
role: "system",
content: x,
}),
),
...MessageV2.toModelMessage(sessionMessages),
...(isLastStep
? [
{
role: "assistant" as const,
content: MAX_STEPS,
},
]
: []),
]
const result = await processor.process({
onError(error) {
log.error("stream error", {
error,
})
},
async experimental_repairToolCall(input) {
const lower = input.toolCall.toolName.toLowerCase()
if (lower !== input.toolCall.toolName && tools[lower]) {
log.info("repairing tool call", {
tool: input.toolCall.toolName,
repaired: lower,
})
return {
...input.toolCall,
toolName: lower,
}
}
return {
...input.toolCall,
input: JSON.stringify({
tool: input.toolCall.toolName,
error: input.error.message,
}),
toolName: "invalid",
}
},
headers: {
...(model.providerID.startsWith("opencode")
? {
"x-opencode-project": Instance.project.id,
"x-opencode-session": sessionID,
"x-opencode-request": lastUser.id,
"x-opencode-client": Flag.OPENCODE_CLIENT,
}
: undefined),
...model.headers,
},
// set to 0, we handle loop
maxRetries: 0,
activeTools: Object.keys(tools).filter((x) => x !== "invalid"),
maxOutputTokens: ProviderTransform.maxOutputTokens(
model.api.npm,
params.options,
model.limit.output,
OUTPUT_TOKEN_MAX,
),
abortSignal: abort,
providerOptions: ProviderTransform.providerOptions(model, params.options),
stopWhen: stepCountIs(1),
temperature: params.temperature,
topP: params.topP,
topK: params.topK,
toolChoice: isLastStep ? "none" : undefined,
messages,
tools: model.capabilities.toolcall === false ? undefined : tools,
model: wrapLanguageModel({
model: language,
middleware: [
{
async transformParams(args) {
if (args.type === "stream") {
// @ts-expect-error - prompt types are compatible at runtime
args.params.prompt = ProviderTransform.message(args.params.prompt, model)
}
// Transform tool schemas for provider compatibility
if (args.params.tools && Array.isArray(args.params.tools)) {
args.params.tools = args.params.tools.map((tool: any) => {
// Tools at middleware level have inputSchema, not parameters
if (tool.inputSchema && typeof tool.inputSchema === "object") {
// Transform the inputSchema for provider compatibility
return {
...tool,
inputSchema: ProviderTransform.schema(model, tool.inputSchema),
}
}
// If no inputSchema, return tool unchanged
return tool
})
}
return args.params
},
},
],
}),
experimental_telemetry: {
isEnabled: cfg.experimental?.openTelemetry,
metadata: {
userId: cfg.username ?? "unknown",
sessionId: sessionID,
},
},
user: lastUser,
agent,
abort,
sessionID,
system: [...(await SystemPrompt.environment()), ...(await SystemPrompt.custom())],
messages: [
...MessageV2.toModelMessage(msgs),
...(isLastStep
? [
{
role: "assistant" as const,
content: MAX_STEPS,
},
]
: []),
],
tools,
model,
})
if (result === "stop") break
continue
@@ -675,33 +522,6 @@ export namespace SessionPrompt {
return Provider.defaultModel()
}
async function resolveSystemPrompt(input: {
system?: string
agent: Agent.Info
model: Provider.Model
isLastStep?: boolean
}) {
let system = SystemPrompt.header(input.model.providerID)
system.push(
...(() => {
if (input.system) return [input.system]
if (input.agent.prompt) return [input.agent.prompt]
return SystemPrompt.provider(input.model)
})(),
)
system.push(...(await SystemPrompt.environment()))
system.push(...(await SystemPrompt.custom()))
if (input.isLastStep) {
system.push(MAX_STEPS)
}
// max 2 system prompt messages for caching purposes
const [first, ...rest] = system
system = [first, rest.join("\n")]
return system
}
async function resolveTools(input: {
agent: Agent.Info
model: Provider.Model
@@ -709,6 +529,7 @@ export namespace SessionPrompt {
tools?: Record<string, boolean>
processor: SessionProcessor.Info
}) {
using _ = log.time("resolveTools")
const tools: Record<string, AITool> = {}
const enabledTools = pipe(
input.agent.tools,
@@ -778,7 +599,6 @@ export namespace SessionPrompt {
},
})
}
for (const [key, item] of Object.entries(await MCP.tools())) {
if (Wildcard.all(key, enabledTools) === false) continue
const execute = item.execute
@@ -857,7 +677,6 @@ export namespace SessionPrompt {
created: Date.now(),
},
tools: input.tools,
system: input.system,
agent: agent.name,
model: input.model ?? agent.model ?? (await lastModel(input.sessionID)),
}
@@ -1148,7 +967,7 @@ export namespace SessionPrompt {
synthetic: true,
})
}
const wasPlan = input.messages.some((msg) => msg.info.role === "assistant" && msg.info.mode === "plan")
const wasPlan = input.messages.some((msg) => msg.info.role === "assistant" && msg.info.agent === "plan")
if (wasPlan && input.agent.name === "build") {
userMessage.parts.push({
id: Identifier.ascending("part"),
@@ -1216,6 +1035,7 @@ export namespace SessionPrompt {
sessionID: input.sessionID,
parentID: userMsg.id,
mode: input.agent,
agent: input.agent,
cost: 0,
path: {
cwd: Instance.directory,
@@ -1510,28 +1330,24 @@ export namespace SessionPrompt {
input.history.filter((m) => m.info.role === "user" && !m.parts.every((p) => "synthetic" in p && p.synthetic))
.length === 1
if (!isFirst) return
const cfg = await Config.get()
const small =
(await Provider.getSmallModel(input.providerID)) ?? (await Provider.getModel(input.providerID, input.modelID))
const language = await Provider.getLanguage(small)
const provider = await Provider.getProvider(small.providerID)
const options = pipe(
{},
mergeDeep(ProviderTransform.options(small, input.session.id, provider?.options)),
mergeDeep(ProviderTransform.smallOptions(small)),
mergeDeep(small.options),
)
await generateText({
// use higher # for reasoning models since reasoning tokens eat up a lot of the budget
maxOutputTokens: small.capabilities.reasoning ? 3000 : 20,
providerOptions: ProviderTransform.providerOptions(small, options),
const agent = await Agent.get("title")
if (!agent) return
const result = await LLM.stream({
agent,
user: input.message.info as MessageV2.User,
system: [],
small: true,
tools: {},
model: await iife(async () => {
if (agent.model) return await Provider.getModel(agent.model.providerID, agent.model.modelID)
return (
(await Provider.getSmallModel(input.providerID)) ?? (await Provider.getModel(input.providerID, input.modelID))
)
}),
abort: new AbortController().signal,
sessionID: input.session.id,
retries: 2,
messages: [
...SystemPrompt.title(small.providerID).map(
(x): ModelMessage => ({
role: "system",
content: x,
}),
),
{
role: "user",
content: "Generate a title for this conversation:\n",
@@ -1555,32 +1371,19 @@ export namespace SessionPrompt {
},
]),
],
headers: small.headers,
model: language,
experimental_telemetry: {
isEnabled: cfg.experimental?.openTelemetry,
metadata: {
userId: cfg.username ?? "unknown",
sessionId: input.session.id,
},
},
})
.then((result) => {
if (result.text)
return Session.update(input.session.id, (draft) => {
const cleaned = result.text
.replace(/<think>[\s\S]*?<\/think>\s*/g, "")
.split("\n")
.map((line) => line.trim())
.find((line) => line.length > 0)
if (!cleaned) return
const text = await result.text.catch((err) => log.error("failed to generate title", { error: err }))
if (text)
return Session.update(input.session.id, (draft) => {
const cleaned = text
.replace(/<think>[\s\S]*?<\/think>\s*/g, "")
.split("\n")
.map((line) => line.trim())
.find((line) => line.length > 0)
if (!cleaned) return
const title = cleaned.length > 100 ? cleaned.substring(0, 97) + "..." : cleaned
draft.title = title
})
})
.catch((error) => {
log.error("failed to generate title", { error, model: small.id })
const title = cleaned.length > 100 ? cleaned.substring(0, 97) + "..." : cleaned
draft.title = title
})
}
}