fix: ensure images are properly returned as tool results
This commit is contained in:
@@ -435,6 +435,40 @@ export namespace MessageV2 {
|
||||
|
||||
export function toModelMessages(input: WithParts[], model: Provider.Model): ModelMessage[] {
|
||||
const result: UIMessage[] = []
|
||||
const toolNames = new Set<string>()
|
||||
|
||||
const toModelOutput = (output: unknown) => {
|
||||
if (typeof output === "string") {
|
||||
return { type: "text", value: output }
|
||||
}
|
||||
|
||||
if (typeof output === "object") {
|
||||
const outputObject = output as {
|
||||
text: string
|
||||
attachments?: Array<{ mime: string; url: string }>
|
||||
}
|
||||
const attachments = (outputObject.attachments ?? []).filter((attachment) => {
|
||||
return attachment.url.startsWith("data:") && attachment.url.includes(",")
|
||||
})
|
||||
|
||||
return {
|
||||
type: "content",
|
||||
value: [
|
||||
{ type: "text", text: outputObject.text },
|
||||
...attachments.map((attachment) => ({
|
||||
type: "media",
|
||||
mediaType: attachment.mime,
|
||||
data: iife(() => {
|
||||
const commaIndex = attachment.url.indexOf(",")
|
||||
return commaIndex === -1 ? attachment.url : attachment.url.slice(commaIndex + 1)
|
||||
}),
|
||||
})),
|
||||
],
|
||||
}
|
||||
}
|
||||
|
||||
return { type: "json", value: output as never }
|
||||
}
|
||||
|
||||
for (const msg of input) {
|
||||
if (msg.parts.length === 0) continue
|
||||
@@ -505,31 +539,24 @@ export namespace MessageV2 {
|
||||
type: "step-start",
|
||||
})
|
||||
if (part.type === "tool") {
|
||||
toolNames.add(part.tool)
|
||||
if (part.state.status === "completed") {
|
||||
if (part.state.attachments?.length) {
|
||||
result.push({
|
||||
id: Identifier.ascending("message"),
|
||||
role: "user",
|
||||
parts: [
|
||||
{
|
||||
type: "text",
|
||||
text: `The tool ${part.tool} returned the following attachments:`,
|
||||
},
|
||||
...part.state.attachments.map((attachment) => ({
|
||||
type: "file" as const,
|
||||
url: attachment.url,
|
||||
mediaType: attachment.mime,
|
||||
filename: attachment.filename,
|
||||
})),
|
||||
],
|
||||
})
|
||||
}
|
||||
const outputText = part.state.time.compacted ? "[Old tool result content cleared]" : part.state.output
|
||||
const attachments = part.state.time.compacted ? [] : (part.state.attachments ?? [])
|
||||
const output =
|
||||
attachments.length > 0
|
||||
? {
|
||||
text: outputText,
|
||||
attachments,
|
||||
}
|
||||
: outputText
|
||||
|
||||
assistantMessage.parts.push({
|
||||
type: ("tool-" + part.tool) as `tool-${string}`,
|
||||
state: "output-available",
|
||||
toolCallId: part.callID,
|
||||
input: part.state.input,
|
||||
output: part.state.time.compacted ? "[Old tool result content cleared]" : part.state.output,
|
||||
output,
|
||||
...(differentModel ? {} : { callProviderMetadata: part.metadata }),
|
||||
})
|
||||
}
|
||||
@@ -568,7 +595,15 @@ export namespace MessageV2 {
|
||||
}
|
||||
}
|
||||
|
||||
return convertToModelMessages(result.filter((msg) => msg.parts.some((part) => part.type !== "step-start")))
|
||||
const tools = Object.fromEntries(Array.from(toolNames).map((toolName) => [toolName, { toModelOutput }]))
|
||||
|
||||
return 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)
|
||||
tools,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
export const stream = fn(Identifier.schema("session"), async function* (sessionID) {
|
||||
|
||||
@@ -722,12 +722,6 @@ export namespace SessionPrompt {
|
||||
)
|
||||
return result
|
||||
},
|
||||
toModelOutput(result) {
|
||||
return {
|
||||
type: "text",
|
||||
value: result.output,
|
||||
}
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
@@ -819,12 +813,6 @@ export namespace SessionPrompt {
|
||||
content: result.content, // directly return content to preserve ordering when outputting to model
|
||||
}
|
||||
}
|
||||
item.toModelOutput = (result) => {
|
||||
return {
|
||||
type: "text",
|
||||
value: result.output,
|
||||
}
|
||||
}
|
||||
tools[key] = item
|
||||
}
|
||||
|
||||
|
||||
@@ -262,7 +262,7 @@ describe("session.message-v2.toModelMessage", () => {
|
||||
])
|
||||
})
|
||||
|
||||
test("converts assistant tool completion into tool-call + tool-result messages and emits attachment message", () => {
|
||||
test("converts assistant tool completion into tool-call + tool-result messages with attachments", () => {
|
||||
const userID = "m-user"
|
||||
const assistantID = "m-assistant"
|
||||
|
||||
@@ -304,7 +304,7 @@ describe("session.message-v2.toModelMessage", () => {
|
||||
type: "file",
|
||||
mime: "image/png",
|
||||
filename: "attachment.png",
|
||||
url: "https://example.com/attachment.png",
|
||||
url: "data:image/png;base64,Zm9v",
|
||||
},
|
||||
],
|
||||
},
|
||||
@@ -319,18 +319,6 @@ describe("session.message-v2.toModelMessage", () => {
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "run tool" }],
|
||||
},
|
||||
{
|
||||
role: "user",
|
||||
content: [
|
||||
{ type: "text", text: "The tool bash returned the following attachments:" },
|
||||
{
|
||||
type: "file",
|
||||
mediaType: "image/png",
|
||||
filename: "attachment.png",
|
||||
data: "https://example.com/attachment.png",
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
role: "assistant",
|
||||
content: [
|
||||
@@ -352,7 +340,13 @@ describe("session.message-v2.toModelMessage", () => {
|
||||
type: "tool-result",
|
||||
toolCallId: "call-1",
|
||||
toolName: "bash",
|
||||
output: { type: "text", value: "ok" },
|
||||
output: {
|
||||
type: "content",
|
||||
value: [
|
||||
{ type: "text", text: "ok" },
|
||||
{ type: "media", mediaType: "image/png", data: "Zm9v" },
|
||||
],
|
||||
},
|
||||
providerOptions: { openai: { tool: "meta" } },
|
||||
},
|
||||
],
|
||||
|
||||
Reference in New Issue
Block a user