| 1 | import { type UIMessage } from "ai"; |
| 2 | import { |
| 3 | AIMessage, |
| 4 | ToolMessage, |
| 5 | type BaseMessage, |
| 6 | } from "@langchain/core/messages"; |
| 7 | |
| 8 | function getTextContent(message: BaseMessage) { |
| 9 | if (typeof message.content === "string") { |
| 10 | return message.content; |
| 11 | } |
| 12 | |
| 13 | return message.contentBlocks |
| 14 | .filter((part) => part.type === "text") |
| 15 | .map((part) => ("text" in part ? part.text ?? "" : "")) |
| 16 | .join(""); |
| 17 | } |
| 18 | |
| 19 | function createTextMessage(message: BaseMessage, role: "user" | "assistant") { |
| 20 | const text = getTextContent(message); |
| 21 | |
| 22 | return { |
| 23 | id: message.id ?? crypto.randomUUID(), |
| 24 | role, |
| 25 | parts: text |
| 26 | ? [ |
| 27 | { |
| 28 | type: "text" as const, |
| 29 | text, |
| 30 | }, |
| 31 | ] |
| 32 | : [], |
| 33 | } satisfies UIMessage; |
| 34 | } |
| 35 | |
| 36 | export function toHydratedUiMessages(messages: BaseMessage[]): UIMessage[] { |
| 37 | const hydratedMessages: UIMessage[] = []; |
| 38 | const toolCallLocations = new Map< |
| 39 | string, |
| 40 | { |
| 41 | messageIndex: number; |
| 42 | partIndex: number; |
| 43 | } |
| 44 | >(); |
| 45 | |
| 46 | for (const message of messages) { |
| 47 | if (message.type === "system" || message.type === "developer") { |
| 48 | continue; |
| 49 | } |
| 50 | |
| 51 | if (message.type === "human") { |
| 52 | hydratedMessages.push(createTextMessage(message, "user")); |
| 53 | continue; |
| 54 | } |
| 55 | |
| 56 | if (AIMessage.isInstance(message)) { |
| 57 | const parts: UIMessage["parts"] = []; |
| 58 | const text = getTextContent(message); |
| 59 | |
| 60 | if (text.trim().length > 0) { |
| 61 | parts.push({ |
| 62 | type: "text", |
| 63 | text, |
| 64 | }); |
| 65 | } |
| 66 | |
| 67 | for (const toolCall of message.tool_calls ?? []) { |
| 68 | const toolCallId = toolCall.id ?? crypto.randomUUID(); |
| 69 | const partIndex = parts.length; |
| 70 | |
| 71 | parts.push({ |
| 72 | type: "dynamic-tool", |
| 73 | toolName: toolCall.name, |
| 74 | toolCallId, |
| 75 | state: "input-available", |
| 76 | input: toolCall.args ?? {}, |
| 77 | }); |
| 78 | |
| 79 | toolCallLocations.set(toolCallId, { |
| 80 | messageIndex: hydratedMessages.length, |
| 81 | partIndex, |
| 82 | }); |
| 83 | } |
| 84 | |
| 85 | if (parts.length > 0) { |
| 86 | hydratedMessages.push({ |
| 87 | id: message.id ?? crypto.randomUUID(), |
| 88 | role: "assistant", |
| 89 | parts, |
| 90 | }); |
| 91 | } |
| 92 | |
| 93 | continue; |
| 94 | } |
| 95 | |
| 96 | if (ToolMessage.isInstance(message)) { |
| 97 | const existingLocation = toolCallLocations.get(message.tool_call_id); |
| 98 | const outputText = getTextContent(message); |
| 99 | |
| 100 | if (existingLocation) { |
| 101 | const existingMessage = hydratedMessages[existingLocation.messageIndex]; |
| 102 | const existingPart = existingMessage?.parts[existingLocation.partIndex]; |
| 103 | |
| 104 | if (existingMessage && existingPart?.type === "dynamic-tool") { |
| 105 | existingMessage.parts[existingLocation.partIndex] = { |
| 106 | type: "dynamic-tool", |
| 107 | toolName: existingPart.toolName, |
| 108 | toolCallId: existingPart.toolCallId, |
| 109 | state: |
| 110 | message.status === "error" ? "output-error" : "output-available", |
| 111 | input: "input" in existingPart ? existingPart.input : {}, |
| 112 | ...(message.status === "error" |
| 113 | ? { |
| 114 | errorText: outputText || "Tool execution failed", |
| 115 | } |
| 116 | : { |
| 117 | output: outputText, |
| 118 | }), |
| 119 | } as UIMessage["parts"][number]; |
| 120 | continue; |
| 121 | } |
| 122 | } |
| 123 | |
| 124 | hydratedMessages.push({ |
| 125 | id: message.id ?? crypto.randomUUID(), |
| 126 | role: "assistant", |
| 127 | parts: [ |
| 128 | { |
| 129 | type: "dynamic-tool", |
| 130 | toolName: message.name ?? "tool", |
| 131 | toolCallId: message.tool_call_id, |
| 132 | state: |
| 133 | message.status === "error" ? "output-error" : "output-available", |
| 134 | input: {}, |
| 135 | ...(message.status === "error" |
| 136 | ? { |
| 137 | errorText: outputText || "Tool execution failed", |
| 138 | } |
| 139 | : { |
| 140 | output: outputText, |
| 141 | }), |
| 142 | } as UIMessage["parts"][number], |
| 143 | ], |
| 144 | }); |
| 145 | continue; |
| 146 | } |
| 147 | |
| 148 | hydratedMessages.push(createTextMessage(message, "assistant")); |
| 149 | } |
| 150 | |
| 151 | return hydratedMessages; |
| 152 | } |
| 153 |