feat: 支持传图

This commit is contained in:
ZyGuo committed 2026-06-03 10:51:42 +08:00
1 parent 34a9f15cb4
commit 5d39aa55bd
3 files changed
+331 -86

No files matched your search

+13 -86
View File
@@ -30,21 +30,10 @@ import { useThreadStream } from "@/core/threads/hooks";
import { textOfMessage } from "@/core/threads/utils";
import { env } from "@/env";
import { cn } from "@/lib/utils";
type BrainyChatSubmitMessage = {
source?: string;
type?: string;
payload?: {
text?: string;
prompt?: string;
pageContext?: {
title?: string;
content?: string;
url?: string;
};
};
requestId?: string;
};
import {
type BrainyChatSubmitMessage,
parseBrainyChatSubmitMessage,
} from "./iframe-chat-submit";
type BrainyChatStatus = "started" | "finished";
@@ -62,47 +51,6 @@ function postBrainyChatStatus(status: BrainyChatStatus, requestId: string) {
);
}
function buildIframeMessageText({
text,
prompt,
pageContext,
}: {
text: string;
prompt: string;
pageContext?:
| {
title?: string;
content?: string;
url?: string;
}
| undefined;
}) {
const primaryText = text || prompt;
const sections = [primaryText];
if (prompt && prompt !== primaryText) {
sections.push(`<extension_prompt>\n${prompt}\n</extension_prompt>`);
}
if (pageContext) {
const pageContextParts = [
pageContext.title ? `<title>\n${pageContext.title}\n</title>` : null,
pageContext.url ? `<url>\n${pageContext.url}\n</url>` : null,
pageContext.content
? `<content>\n${pageContext.content}\n</content>`
: null,
].filter(Boolean);
if (pageContextParts.length > 0) {
sections.push(
`<page_context>\n${pageContextParts.join("\n\n")}\n</page_context>`,
);
}
}
return sections.join("\n\n");
}
export default function IframeChatPage() {
const routeBase = useRouteBase();
const router = useRouter();
@@ -265,44 +213,23 @@ export default function IframeChatPage() {
if (data?.source !== "brainy-extension") return;
if (data?.type !== "brainy:chat-submit") return;
const text = data.payload?.text?.trim() ?? "";
const prompt = data.payload?.prompt?.trim() ?? "";
const rawPageContext = data.payload?.pageContext;
const pageContext = rawPageContext
? {
...(rawPageContext.title?.trim()
? { title: rawPageContext.title.trim() }
: {}),
...(rawPageContext.content?.trim()
? { content: rawPageContext.content.trim() }
: {}),
...(rawPageContext.url?.trim()
? { url: rawPageContext.url.trim() }
: {}),
}
: undefined;
const submitText = buildIframeMessageText({ text, prompt, pageContext });
const { submitText, files, iframeContext } =
parseBrainyChatSubmitMessage(data);
console.log("收到插件消息", {
text,
prompt,
pageContext,
text: data.payload?.text,
prompt: data.payload?.prompt,
pageContext: data.payload?.pageContext,
screenshot: data.payload?.screenshot,
requestId: data.requestId,
filesCount: files.length,
});
if (!submitText) return;
if (!submitText && files.length === 0) return;
pendingRequestIdRef.current = data.requestId ?? null;
const iframeContext: Record<string, unknown> = {};
if (data.requestId) {
iframeContext.brainy_request_id = data.requestId;
}
handleSubmit(
{ text: submitText, files: [] },
Object.keys(iframeContext).length > 0 ? iframeContext : undefined,
);
handleSubmit({ text: submitText, files }, iframeContext);
};
window.addEventListener("message", handleMessage);
@@ -0,0 +1,71 @@
import assert from "node:assert/strict";
import test from "node:test";
import { parseBrainyChatSubmitMessage } from "./iframe-chat-submit";
test("parses iframe text, page context, and screenshot into a prompt message", () => {
const result = parseBrainyChatSubmitMessage({
source: "brainy-extension",
type: "brainy:chat-submit",
requestId: "iframe-123",
payload: {
text: "提取并翻译",
prompt: "提取图片中的文字并翻译成简体中文",
pageContext: {
title: "UN News",
content: "Banning children from social media is not the answer.",
url: "https://news.un.org/example",
},
screenshot: {
dataUrl: "data:image/png;base64,aGVsbG8=",
fileName: "capture.png",
images: [
{
filename: "capture.png",
url: "C:\\Users\\Rocky\\AppData\\Local\\Temp\\capture.png",
size: 0,
},
],
},
},
});
assert.equal(result.submitText, [
"提取并翻译",
"<extension_prompt>\n提取图片中的文字并翻译成简体中文\n</extension_prompt>",
"<page_context>\n<title>\nUN News\n</title>\n\n<url>\nhttps://news.un.org/example\n</url>\n\n<content>\nBanning children from social media is not the answer.\n</content>\n</page_context>",
].join("\n\n"));
assert.deepEqual(result.iframeContext, { brainy_request_id: "iframe-123" });
assert.equal(result.files.length, 1);
assert.equal(result.files[0]?.type, "file");
assert.equal(result.files[0]?.filename, "capture.png");
assert.equal(result.files[0]?.mediaType, "image/png");
assert.ok(result.files[0]?.file instanceof File);
assert.equal(result.files[0]?.file?.name, "capture.png");
assert.equal(result.files[0]?.file?.type, "image/png");
assert.equal(result.files[0]?.url, "data:image/png;base64,aGVsbG8=");
});
test("keeps screenshot-only iframe submissions in the conversation flow", () => {
const result = parseBrainyChatSubmitMessage({
source: "brainy-extension",
type: "brainy:chat-submit",
payload: {
screenshot: {
dataUrl: "data:image/jpeg;base64,aGVsbG8=",
images: [
{
filename: "shot.jpg",
url: "C:\\Users\\Rocky\\AppData\\Local\\Temp\\shot.jpg",
size: 0,
},
],
},
},
});
assert.equal(result.submitText, "");
assert.equal(result.files.length, 1);
assert.equal(result.files[0]?.filename, "shot.jpg");
assert.equal(result.files[0]?.mediaType, "image/jpeg");
});
@@ -0,0 +1,247 @@
type BrainyPageContext = {
title?: string;
content?: string;
url?: string;
};
type BrainyScreenshotPath = {
image_url?: string;
};
type BrainyScreenshotImage = {
filename?: string;
url?: string;
size?: number;
};
type BrainyScreenshot = {
dataUrl?: string;
fileName?: string;
filePath?: BrainyScreenshotPath[];
images?: BrainyScreenshotImage[];
};
export type BrainyChatSubmitMessage = {
source?: string;
type?: string;
payload?: {
text?: string;
prompt?: string;
pageContext?: BrainyPageContext;
screenshot?: BrainyScreenshot;
};
requestId?: string;
};
export type IframePromptInputFilePart = {
type: "file";
url: string;
mediaType: string;
filename: string;
file?: File;
};
export type ParsedBrainyChatSubmitMessage = {
submitText: string;
files: IframePromptInputFilePart[];
iframeContext?: Record<string, unknown>;
};
function trimToUndefined(value: string | undefined): string | undefined {
const trimmed = value?.trim();
return trimmed ? trimmed : undefined;
}
function fileNameFromPath(path: string | undefined): string | undefined {
if (!path) return undefined;
const parts = path.split(/[\\/]/).filter(Boolean);
return parts.at(-1);
}
function isLikelyLocalPath(path: string): boolean {
return /^[a-zA-Z]:[\\/]/.test(path) || path.startsWith("\\\\");
}
function mimeTypeFromFilename(filename: string | undefined): string | undefined {
const lower = filename?.toLowerCase() ?? "";
if (lower.endsWith(".png")) return "image/png";
if (lower.endsWith(".jpg") || lower.endsWith(".jpeg")) return "image/jpeg";
if (lower.endsWith(".webp")) return "image/webp";
if (lower.endsWith(".gif")) return "image/gif";
return undefined;
}
function parseDataUrlMeta(dataUrl: string): { mimeType: string; payload: string } | null {
const match = /^data:([^;,]+)?(?:;base64)?,(.*)$/i.exec(dataUrl);
if (!match) return null;
return {
mimeType: match[1] || "application/octet-stream",
payload: match[2] || "",
};
}
function decodeBase64(base64: string): Uint8Array {
if (typeof Buffer !== "undefined") {
return Uint8Array.from(Buffer.from(base64, "base64"));
}
const binary = atob(base64);
const bytes = new Uint8Array(binary.length);
for (let i = 0; i < binary.length; i += 1) {
bytes[i] = binary.charCodeAt(i);
}
return bytes;
}
function fileFromDataUrl(
dataUrl: string,
filename: string,
mediaType: string,
): File | undefined {
const meta = parseDataUrlMeta(dataUrl);
if (!meta) return undefined;
const bytes = decodeBase64(meta.payload);
const arrayBuffer = bytes.buffer.slice(
bytes.byteOffset,
bytes.byteOffset + bytes.byteLength,
);
return new File([arrayBuffer], filename, {
type: mediaType || meta.mimeType,
});
}
function normalizePageContext(
pageContext: BrainyPageContext | undefined,
): BrainyPageContext | undefined {
if (!pageContext) return undefined;
const normalized = {
title: trimToUndefined(pageContext.title),
content: trimToUndefined(pageContext.content),
url: trimToUndefined(pageContext.url),
};
return normalized.title || normalized.content || normalized.url
? normalized
: undefined;
}
export function buildIframeMessageText({
text,
prompt,
pageContext,
}: {
text?: string;
prompt?: string;
pageContext?: BrainyPageContext;
}) {
const primaryText = trimToUndefined(text) || trimToUndefined(prompt) || "";
const sections: string[] = [];
if (primaryText) {
sections.push(primaryText);
}
const normalizedPrompt = trimToUndefined(prompt);
if (normalizedPrompt && normalizedPrompt !== primaryText) {
sections.push(
`<extension_prompt>\n${normalizedPrompt}\n</extension_prompt>`,
);
}
const normalizedPageContext = normalizePageContext(pageContext);
if (normalizedPageContext) {
const pageContextParts = [
normalizedPageContext.title
? `<title>\n${normalizedPageContext.title}\n</title>`
: null,
normalizedPageContext.url
? `<url>\n${normalizedPageContext.url}\n</url>`
: null,
normalizedPageContext.content
? `<content>\n${normalizedPageContext.content}\n</content>`
: null,
].filter(Boolean);
if (pageContextParts.length > 0) {
sections.push(
`<page_context>\n${pageContextParts.join("\n\n")}\n</page_context>`,
);
}
}
return sections.join("\n\n");
}
function buildScreenshotFileParts(
screenshot: BrainyScreenshot | undefined,
): IframePromptInputFilePart[] {
if (!screenshot) return [];
const fallbackPath =
screenshot.filePath?.map((item) => item.image_url).find(Boolean) ??
screenshot.images?.map((item) => item.url).find(Boolean);
const fallbackFilename =
trimToUndefined(screenshot.fileName) ??
screenshot.images?.map((item) => trimToUndefined(item.filename)).find(Boolean) ??
fileNameFromPath(fallbackPath) ??
"screenshot.png";
const dataUrl = trimToUndefined(screenshot.dataUrl);
if (dataUrl) {
const parsed = parseDataUrlMeta(dataUrl);
const mediaType =
parsed?.mimeType || mimeTypeFromFilename(fallbackFilename) || "image/png";
const file = fileFromDataUrl(dataUrl, fallbackFilename, mediaType);
return [
{
type: "file",
url: dataUrl,
mediaType,
filename: fallbackFilename,
...(file ? { file } : {}),
},
];
}
const remoteImage = screenshot.images?.find((item) => {
const url = trimToUndefined(item.url);
return !!url && !isLikelyLocalPath(url);
});
if (remoteImage?.url && remoteImage.filename) {
return [
{
type: "file",
url: remoteImage.url,
mediaType:
mimeTypeFromFilename(remoteImage.filename) || "application/octet-stream",
filename: remoteImage.filename,
},
];
}
return [];
}
export function parseBrainyChatSubmitMessage(
data: BrainyChatSubmitMessage,
): ParsedBrainyChatSubmitMessage {
const text = trimToUndefined(data.payload?.text);
const prompt = trimToUndefined(data.payload?.prompt);
const pageContext = normalizePageContext(data.payload?.pageContext);
const submitText = buildIframeMessageText({ text, prompt, pageContext });
const files = buildScreenshotFileParts(data.payload?.screenshot);
const requestId = trimToUndefined(data.requestId);
return {
submitText,
files,
...(requestId
? { iframeContext: { brainy_request_id: requestId } }
: {}),
};
}