feat: AI SDK v6 support (#18433)
This commit is contained in:
@@ -215,7 +215,7 @@ When constructing the summary, try to stick to this template:
|
||||
tools: {},
|
||||
system: [],
|
||||
messages: [
|
||||
...MessageV2.toModelMessages(msgs, model, { stripMedia: true }),
|
||||
...(await MessageV2.toModelMessages(msgs, model, { stripMedia: true })),
|
||||
{
|
||||
role: "user",
|
||||
content: [
|
||||
|
||||
@@ -1,16 +1,6 @@
|
||||
import { Installation } from "@/installation"
|
||||
import { Provider } from "@/provider/provider"
|
||||
import { Log } from "@/util/log"
|
||||
import {
|
||||
streamText,
|
||||
wrapLanguageModel,
|
||||
type ModelMessage,
|
||||
type StreamTextResult,
|
||||
type Tool,
|
||||
type ToolSet,
|
||||
tool,
|
||||
jsonSchema,
|
||||
} from "ai"
|
||||
import { streamText, wrapLanguageModel, type ModelMessage, type Tool, tool, jsonSchema } from "ai"
|
||||
import { mergeDeep, pipe } from "remeda"
|
||||
import { GitLabWorkflowLanguageModel } from "gitlab-ai-provider"
|
||||
import { ProviderTransform } from "@/provider/transform"
|
||||
@@ -23,6 +13,7 @@ import { SystemPrompt } from "./system"
|
||||
import { Flag } from "@/flag/flag"
|
||||
import { Permission } from "@/permission"
|
||||
import { Auth } from "@/auth"
|
||||
import { Installation } from "@/installation"
|
||||
|
||||
export namespace LLM {
|
||||
const log = Log.create({ service: "llm" })
|
||||
@@ -43,8 +34,6 @@ export namespace LLM {
|
||||
toolChoice?: "auto" | "required" | "none"
|
||||
}
|
||||
|
||||
export type StreamOutput = StreamTextResult<ToolSet, unknown>
|
||||
|
||||
export async function stream(input: StreamInput) {
|
||||
const l = log
|
||||
.clone()
|
||||
@@ -273,8 +262,10 @@ export namespace LLM {
|
||||
model: language,
|
||||
middleware: [
|
||||
{
|
||||
specificationVersion: "v3" as const,
|
||||
async transformParams(args) {
|
||||
if (args.type === "stream") {
|
||||
// TODO: verify that LanguageModelV3Prompt is still compat here!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
|
||||
// @ts-expect-error
|
||||
args.params.prompt = ProviderTransform.message(args.params.prompt, input.model, options)
|
||||
}
|
||||
|
||||
@@ -573,11 +573,11 @@ export namespace MessageV2 {
|
||||
}))
|
||||
}
|
||||
|
||||
export function toModelMessages(
|
||||
export async function toModelMessages(
|
||||
input: WithParts[],
|
||||
model: Provider.Model,
|
||||
options?: { stripMedia?: boolean },
|
||||
): ModelMessage[] {
|
||||
): Promise<ModelMessage[]> {
|
||||
const result: UIMessage[] = []
|
||||
const toolNames = new Set<string>()
|
||||
// Track media from tool results that need to be injected as user messages
|
||||
@@ -601,7 +601,8 @@ export namespace MessageV2 {
|
||||
return false
|
||||
})()
|
||||
|
||||
const toModelOutput = (output: unknown) => {
|
||||
const toModelOutput = (options: { toolCallId: string; input: unknown; output: unknown }) => {
|
||||
const output = options.output
|
||||
if (typeof output === "string") {
|
||||
return { type: "text", value: output }
|
||||
}
|
||||
@@ -799,7 +800,7 @@ export namespace MessageV2 {
|
||||
|
||||
const tools = Object.fromEntries(Array.from(toolNames).map((toolName) => [toolName, { toModelOutput }]))
|
||||
|
||||
return convertToModelMessages(
|
||||
return await convertToModelMessages(
|
||||
result.filter((msg) => msg.parts.some((part) => part.type !== "step-start")),
|
||||
{
|
||||
//@ts-expect-error (convertToModelMessages expects a ToolSet but only actually needs tools[name]?.toModelOutput)
|
||||
@@ -871,7 +872,13 @@ export namespace MessageV2 {
|
||||
db.select().from(PartTable).where(eq(PartTable.message_id, message_id)).orderBy(PartTable.id).all(),
|
||||
)
|
||||
return rows.map(
|
||||
(row) => ({ ...row.data, id: row.id, sessionID: row.session_id, messageID: row.message_id }) as MessageV2.Part,
|
||||
(row) =>
|
||||
({
|
||||
...row.data,
|
||||
id: row.id,
|
||||
sessionID: row.session_id,
|
||||
messageID: row.message_id,
|
||||
}) as MessageV2.Part,
|
||||
)
|
||||
})
|
||||
|
||||
|
||||
@@ -11,7 +11,7 @@ import { Session } from "."
|
||||
import { Agent } from "../agent/agent"
|
||||
import { Provider } from "../provider/provider"
|
||||
import { ModelID, ProviderID } from "../provider/schema"
|
||||
import { type Tool as AITool, tool, jsonSchema, type ToolCallOptions, asSchema } from "ai"
|
||||
import { type Tool as AITool, tool, jsonSchema, type ToolExecutionOptions, asSchema } from "ai"
|
||||
import { SessionCompaction } from "./compaction"
|
||||
import { Instance } from "../project/instance"
|
||||
import { Bus } from "../bus"
|
||||
@@ -321,7 +321,13 @@ export namespace SessionPrompt {
|
||||
if (!lastUser) throw new Error("No user message found in stream. This should never happen.")
|
||||
if (
|
||||
lastAssistant?.finish &&
|
||||
!["tool-calls", "unknown"].includes(lastAssistant.finish) &&
|
||||
![
|
||||
"tool-calls",
|
||||
// in v6 unknown became other but other existed in v5 too and was distinctly different
|
||||
// I think there are certain providers that used to have bad stop reasons, not rlly sure which
|
||||
// ones if any still have this?
|
||||
// "unknown",
|
||||
].includes(lastAssistant.finish) &&
|
||||
lastUser.id < lastAssistant.id
|
||||
) {
|
||||
log.info("exiting loop", { sessionID })
|
||||
@@ -692,7 +698,7 @@ export namespace SessionPrompt {
|
||||
sessionID,
|
||||
system,
|
||||
messages: [
|
||||
...MessageV2.toModelMessages(msgs, model),
|
||||
...(await MessageV2.toModelMessages(msgs, model)),
|
||||
...(isLastStep
|
||||
? [
|
||||
{
|
||||
@@ -775,7 +781,7 @@ export namespace SessionPrompt {
|
||||
using _ = log.time("resolveTools")
|
||||
const tools: Record<string, AITool> = {}
|
||||
|
||||
const context = (args: any, options: ToolCallOptions): Tool.Context => ({
|
||||
const context = (args: any, options: ToolExecutionOptions): Tool.Context => ({
|
||||
sessionID: input.session.id,
|
||||
abort: options.abortSignal!,
|
||||
messageID: input.processor.message.id,
|
||||
@@ -861,7 +867,8 @@ export namespace SessionPrompt {
|
||||
const execute = item.execute
|
||||
if (!execute) continue
|
||||
|
||||
const transformed = ProviderTransform.schema(input.model, asSchema(item.inputSchema).jsonSchema)
|
||||
const schema = await asSchema(item.inputSchema).jsonSchema
|
||||
const transformed = ProviderTransform.schema(input.model, schema)
|
||||
item.inputSchema = jsonSchema(transformed)
|
||||
// Wrap execute to add plugin hooks and format output
|
||||
item.execute = async (args, opts) => {
|
||||
@@ -974,10 +981,10 @@ export namespace SessionPrompt {
|
||||
metadata: { valid: true },
|
||||
}
|
||||
},
|
||||
toModelOutput(result) {
|
||||
toModelOutput({ output }) {
|
||||
return {
|
||||
type: "text",
|
||||
value: result.output,
|
||||
value: output.output,
|
||||
}
|
||||
},
|
||||
})
|
||||
@@ -2010,28 +2017,28 @@ NOTE: At any point in time through this workflow you should feel free to ask the
|
||||
(await Provider.getSmallModel(input.providerID)) ?? (await Provider.getModel(input.providerID, input.modelID))
|
||||
)
|
||||
})
|
||||
const result = await LLM.stream({
|
||||
agent,
|
||||
user: firstRealUser.info as MessageV2.User,
|
||||
system: [],
|
||||
small: true,
|
||||
tools: {},
|
||||
model,
|
||||
abort: new AbortController().signal,
|
||||
sessionID: input.session.id,
|
||||
retries: 2,
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: "Generate a title for this conversation:\n",
|
||||
},
|
||||
...(hasOnlySubtaskParts
|
||||
? [{ role: "user" as const, content: subtaskParts.map((p) => p.prompt).join("\n") }]
|
||||
: MessageV2.toModelMessages(contextMessages, model)),
|
||||
],
|
||||
})
|
||||
const text = await result.text.catch((err) => log.error("failed to generate title", { error: err }))
|
||||
if (text) {
|
||||
try {
|
||||
const result = await LLM.stream({
|
||||
agent,
|
||||
user: firstRealUser.info as MessageV2.User,
|
||||
system: [],
|
||||
small: true,
|
||||
tools: {},
|
||||
model,
|
||||
abort: new AbortController().signal,
|
||||
sessionID: input.session.id,
|
||||
retries: 2,
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: "Generate a title for this conversation:\n",
|
||||
},
|
||||
...(hasOnlySubtaskParts
|
||||
? [{ role: "user" as const, content: subtaskParts.map((p) => p.prompt).join("\n") }]
|
||||
: await MessageV2.toModelMessages(contextMessages, model)),
|
||||
],
|
||||
})
|
||||
const text = await result.text
|
||||
const cleaned = text
|
||||
.replace(/<think>[\s\S]*?<\/think>\s*/g, "")
|
||||
.split("\n")
|
||||
@@ -2044,6 +2051,8 @@ NOTE: At any point in time through this workflow you should feel free to ask the
|
||||
if (NotFoundError.isInstance(err)) return
|
||||
throw err
|
||||
})
|
||||
} catch (error) {
|
||||
log.error("failed to generate title", { error })
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user