568 lines
20 KiB
TypeScript
568 lines
20 KiB
TypeScript
"use client";
|
||
|
||
import { useChat } from "@ai-sdk/react";
|
||
import { DefaultChatTransport } from "ai";
|
||
import { useRouter, useSearchParams } from "next/navigation";
|
||
import { useEffect, useRef, useState } from "react";
|
||
import useSWR, { useSWRConfig } from "swr";
|
||
import { unstable_serialize } from "swr/infinite";
|
||
import { ChatHeader } from "@/components/chat-header";
|
||
import {
|
||
AlertDialog,
|
||
AlertDialogAction,
|
||
AlertDialogCancel,
|
||
AlertDialogContent,
|
||
AlertDialogDescription,
|
||
AlertDialogFooter,
|
||
AlertDialogHeader,
|
||
AlertDialogTitle,
|
||
} from "@/components/ui/alert-dialog";
|
||
import { useArtifactSelector } from "@/hooks/use-artifact";
|
||
import { useAutoResume } from "@/hooks/use-auto-resume";
|
||
import { useChatVisibility } from "@/hooks/use-chat-visibility";
|
||
import { usePendingGeneration, type PendingGeneration } from "@/hooks/use-pending-generation";
|
||
import type { Vote } from "@/lib/db/schema";
|
||
import type { UserType } from "@/app/(auth)/auth";
|
||
import { ChatSDKError } from "@/lib/errors";
|
||
import type { Attachment, ChatMessage } from "@/lib/types";
|
||
import type { AppUsage } from "@/lib/usage";
|
||
import { fetcher, fetchWithErrorHandlers, generateUUID } from "@/lib/utils";
|
||
import { Artifact } from "@/components/artifact";
|
||
import { useDataStream } from "@/components/data-stream-provider";
|
||
import { Messages } from "@/components/messages";
|
||
import { MultimodalInput } from "@/components/multimodal-input";
|
||
import { getChatHistoryPaginationKey } from "@/components/sidebar-history";
|
||
import { toast } from "@/components/toast";
|
||
import type { VisibilityType } from "@/components/visibility-selector";
|
||
import type { VideoGenerationParams } from "@/components/video-params-bar";
|
||
import type { ImageGenerationParams } from "@/components/image-params-bar";
|
||
|
||
export function Chat({
|
||
id,
|
||
initialMessages,
|
||
initialChatModel,
|
||
initialVisibilityType,
|
||
isReadonly,
|
||
autoResume,
|
||
initialLastContext,
|
||
userType,
|
||
}: {
|
||
id: string;
|
||
initialMessages: ChatMessage[];
|
||
initialChatModel: string;
|
||
initialVisibilityType: VisibilityType;
|
||
isReadonly: boolean;
|
||
autoResume: boolean;
|
||
initialLastContext?: AppUsage;
|
||
userType?: UserType;
|
||
}) {
|
||
const router = useRouter();
|
||
|
||
const { visibilityType } = useChatVisibility({
|
||
chatId: id,
|
||
initialVisibilityType,
|
||
});
|
||
|
||
// 🆕 持久化生成任务状态（解决用户离开再返回时加载状态丢失）
|
||
const { pendingGeneration, startGeneration, completeGeneration } = usePendingGeneration(id);
|
||
|
||
// 🆕 标记是否已经尝试刷新过（避免无限刷新）
|
||
const hasTriedRefreshRef = useRef(false);
|
||
|
||
// 🆕 当检测到恢复的 pendingGeneration 时，刷新页面获取最新消息
|
||
useEffect(() => {
|
||
if (pendingGeneration && !hasTriedRefreshRef.current) {
|
||
// 检查最后一条消息是否已经是 AI 回复
|
||
const lastMessage = initialMessages.length > 0 ? initialMessages[initialMessages.length - 1] : null;
|
||
|
||
if (lastMessage?.role === "assistant") {
|
||
// 最后一条已经是 AI 回复，说明生成已完成，清除 pending 状态
|
||
console.log("[Chat] 🎉 检测到生成已完成（最后消息是 AI 回复），清除 pending 状态");
|
||
completeGeneration();
|
||
} else if (lastMessage?.role === "user") {
|
||
// 最后一条是用户消息，可能还在生成中，或者已经完成但页面数据过期
|
||
// 标记已尝试刷新，然后刷新页面获取最新数据
|
||
hasTriedRefreshRef.current = true;
|
||
console.log("[Chat] 🔄 检测到恢复的 pending 状态，刷新页面获取最新消息...");
|
||
|
||
// 延迟一小段时间再刷新，避免立即刷新
|
||
setTimeout(() => {
|
||
router.refresh();
|
||
}, 500);
|
||
}
|
||
}
|
||
}, [pendingGeneration, initialMessages, completeGeneration, router]);
|
||
|
||
const { mutate } = useSWRConfig();
|
||
|
||
// Handle browser back/forward navigation
|
||
useEffect(() => {
|
||
const handlePopState = () => {
|
||
// When user navigates back/forward, refresh to sync with URL
|
||
router.refresh();
|
||
};
|
||
|
||
window.addEventListener("popstate", handlePopState);
|
||
return () => window.removeEventListener("popstate", handlePopState);
|
||
}, [router]);
|
||
const { setDataStream } = useDataStream();
|
||
|
||
const [input, setInput] = useState<string>("");
|
||
const [usage, setUsage] = useState<AppUsage | undefined>(initialLastContext);
|
||
const [showCreditCardAlert, setShowCreditCardAlert] = useState(false);
|
||
const [currentModelId, setCurrentModelId] = useState(initialChatModel);
|
||
const currentModelIdRef = useRef(currentModelId);
|
||
|
||
// 🆕 视频参数状态
|
||
const [videoParams, setVideoParams] = useState<VideoGenerationParams | null>(null);
|
||
|
||
// 🆕 图片参数状态（默认值：auto比例，2K分辨率）
|
||
const [imageParams, setImageParams] = useState<ImageGenerationParams>({
|
||
aspectRatio: "auto",
|
||
imageSize: "2K"
|
||
});
|
||
|
||
// 🆕 从 localStorage 读取保存的图片参数
|
||
useEffect(() => {
|
||
try {
|
||
const saved = localStorage.getItem("gemini-image-params");
|
||
if (saved) {
|
||
setImageParams(JSON.parse(saved));
|
||
}
|
||
} catch (e) {
|
||
console.error("Failed to load image params:", e);
|
||
}
|
||
}, []);
|
||
|
||
// 🆕 保存图片参数到 localStorage
|
||
const handleImageParamsChange = (params: ImageGenerationParams) => {
|
||
setImageParams(params);
|
||
try {
|
||
localStorage.setItem("gemini-image-params", JSON.stringify(params));
|
||
} catch (e) {
|
||
console.error("Failed to save image params:", e);
|
||
}
|
||
console.log("[Chat] Image params updated:", params);
|
||
};
|
||
|
||
// 🆕 获取模型列表以检查是否是视频模型/图片模型
|
||
const { data: modelsData } = useSWR<{
|
||
models: Array<{
|
||
id: string;
|
||
modelId: string;
|
||
displayName: string;
|
||
providerType: string;
|
||
supportsVision?: boolean;
|
||
supportsImageGen?: boolean;
|
||
}>;
|
||
}>("/api/models", fetcher);
|
||
|
||
const currentModel = modelsData?.models?.find(m => m.id === currentModelId);
|
||
const isVideoModel = currentModel?.providerType === "volcano-video";
|
||
const isImageModel = currentModel?.supportsImageGen === true;
|
||
|
||
// 🔍 调试日志
|
||
console.log('[Chat] Render - currentModelId:', currentModelId);
|
||
console.log('[Chat] Render - modelsData loaded:', !!modelsData);
|
||
console.log('[Chat] Render - currentModel:', currentModel?.displayName, {
|
||
id: currentModel?.id,
|
||
supportsImageGen: currentModel?.supportsImageGen,
|
||
providerType: currentModel?.providerType,
|
||
isImageModel,
|
||
imageParams
|
||
});
|
||
|
||
// 使用 ref 保存视频/图片相关状态，以便在 prepareSendMessagesRequest 中访问最新值
|
||
const isVideoModelRef = useRef(isVideoModel);
|
||
const videoParamsRef = useRef(videoParams);
|
||
const isImageModelRef = useRef(isImageModel);
|
||
const imageParamsRef = useRef(imageParams);
|
||
|
||
useEffect(() => {
|
||
isVideoModelRef.current = isVideoModel;
|
||
videoParamsRef.current = videoParams;
|
||
isImageModelRef.current = isImageModel;
|
||
imageParamsRef.current = imageParams;
|
||
console.log("[Chat] Refs updated - imageParams:", imageParams);
|
||
}, [isVideoModel, videoParams, isImageModel, imageParams]);
|
||
|
||
// 调试日志
|
||
console.log("[Chat Debug] currentModelId:", currentModelId);
|
||
console.log("[Chat Debug] currentModel:", currentModel);
|
||
console.log("[Chat Debug] isVideoModel:", isVideoModel);
|
||
|
||
// 🆕 当切换到视频模型时，自动初始化 videoParams
|
||
useEffect(() => {
|
||
if (isVideoModel && !videoParams) {
|
||
const defaultParams: VideoGenerationParams = {
|
||
mode: "text-to-video",
|
||
resolution: "720p",
|
||
ratio: "16:9",
|
||
duration: 5,
|
||
watermark: false,
|
||
cameraFixed: false,
|
||
cameraMove: "none",
|
||
fps: 24,
|
||
};
|
||
setVideoParams(defaultParams);
|
||
console.log("[Chat] 视频模型已选择，初始化默认参数:", defaultParams);
|
||
} else if (!isVideoModel && videoParams) {
|
||
// 切换到非视频模型时清空 videoParams
|
||
setVideoParams(null);
|
||
console.log("[Chat] 已切换到非视频模型，清空视频参数");
|
||
}
|
||
}, [isVideoModel, videoParams]);
|
||
|
||
// 🆕 当切换到图片模型时，强制更新 imageParams（确保组件重新渲染）
|
||
useEffect(() => {
|
||
if (isImageModel) {
|
||
// 强制更新一次 imageParams 以触发重新渲染
|
||
setImageParams(prev => ({...prev}));
|
||
console.log("[Chat] 图片模型已激活，触发参数更新");
|
||
}
|
||
}, [isImageModel]);
|
||
|
||
useEffect(() => {
|
||
currentModelIdRef.current = currentModelId;
|
||
// 调试：在控制台显示当前模型切换
|
||
console.log(`[Chat] Model changed to: ${currentModelId}`);
|
||
}, [currentModelId]);
|
||
|
||
const {
|
||
messages,
|
||
setMessages,
|
||
sendMessage,
|
||
status,
|
||
stop,
|
||
regenerate,
|
||
resumeStream,
|
||
} = useChat<ChatMessage>({
|
||
id,
|
||
messages: initialMessages,
|
||
experimental_throttle: 100,
|
||
generateId: generateUUID,
|
||
transport: new DefaultChatTransport({
|
||
api: "/api/chat",
|
||
fetch: fetchWithErrorHandlers,
|
||
prepareSendMessagesRequest(request) {
|
||
console.log("[Chat] Preparing request, videoParams:", videoParamsRef.current);
|
||
console.log("[Chat] Preparing request, imageParams:", imageParamsRef.current);
|
||
console.log("[Chat] Is video model:", isVideoModelRef.current);
|
||
console.log("[Chat] Is image model:", isImageModelRef.current);
|
||
return {
|
||
body: {
|
||
id: request.id,
|
||
message: request.messages.at(-1),
|
||
selectedChatModel: currentModelIdRef.current,
|
||
selectedVisibilityType: visibilityType,
|
||
videoParams: videoParamsRef.current || undefined, // 🆕 传递视频参数（使用ref获取最新值）
|
||
imageParams: imageParamsRef.current || undefined, // 🆕 传递图片参数（使用ref获取最新值）
|
||
...request.body,
|
||
},
|
||
};
|
||
},
|
||
}),
|
||
onData: (dataPart: any) => {
|
||
setDataStream((ds) => (ds ? [...ds, dataPart] : []));
|
||
if (dataPart.type === "data-usage") {
|
||
setUsage(dataPart.data);
|
||
}
|
||
|
||
// 处理后端返回的自定义消息（用于图片/视频生成的非流式响应）
|
||
if (dataPart.type === "data-appendMessage") {
|
||
try {
|
||
// dataPart.data 是 JSON 字符串
|
||
const messageData = typeof dataPart.data === 'string'
|
||
? JSON.parse(dataPart.data)
|
||
: dataPart.data;
|
||
|
||
console.log("[Chat] Received appendMessage:", messageData);
|
||
|
||
// 🆕 收到消息，清除持久化状态
|
||
completeGeneration();
|
||
|
||
setMessages((prevMessages) => {
|
||
if (prevMessages.some(m => m.id === messageData.id)) {
|
||
return prevMessages;
|
||
}
|
||
|
||
// 构造 ChatMessage 对象
|
||
// 注意：这里需要确保符合 ChatMessage 类型定义
|
||
const newMessage: any = {
|
||
id: messageData.id,
|
||
role: messageData.role,
|
||
content: messageData.parts?.[0]?.text || "",
|
||
createdAt: new Date(),
|
||
parts: messageData.parts || [],
|
||
attachments: [],
|
||
};
|
||
|
||
return [...prevMessages, newMessage];
|
||
});
|
||
} catch (e) {
|
||
console.error("[Chat] Failed to parse data-appendMessage:", e);
|
||
}
|
||
}
|
||
},
|
||
onFinish: () => {
|
||
mutate(unstable_serialize(getChatHistoryPaginationKey));
|
||
// 🆕 生成完成，清除持久化状态
|
||
completeGeneration();
|
||
},
|
||
onError: (error) => {
|
||
if (error instanceof ChatSDKError) {
|
||
// Check if it's a credit card error
|
||
if (
|
||
error.message?.includes("AI Gateway requires a valid credit card")
|
||
) {
|
||
setShowCreditCardAlert(true);
|
||
} else {
|
||
toast({
|
||
type: "error",
|
||
description: error.message,
|
||
});
|
||
}
|
||
}
|
||
},
|
||
});
|
||
|
||
// 调试：监控聊天状态变化 & 持久化生成状态
|
||
// 使用 ref 来访问最新的 modelsData，避免依赖数组问题
|
||
const modelsDataRef = useRef(modelsData);
|
||
useEffect(() => {
|
||
modelsDataRef.current = modelsData;
|
||
}, [modelsData]);
|
||
|
||
useEffect(() => {
|
||
console.log(`[Chat] Status changed: ${status}`);
|
||
|
||
// 🆕 当状态变为 submitted 时，记录生成任务
|
||
if (status === "submitted") {
|
||
// 判断模型类型 - 使用 ref 获取最新的 modelsData
|
||
const models = modelsDataRef.current?.models;
|
||
const model = models?.find((m: any) => m.id === currentModelId);
|
||
|
||
let modelType: PendingGeneration["modelType"] = "other";
|
||
if (model?.providerType === "gemini") modelType = "gemini";
|
||
else if (model?.providerType === "volcano-vision") modelType = "volcano-vision";
|
||
else if (model?.providerType === "volcano-video") modelType = "volcano-video";
|
||
else if (model?.providerType === "qianwen-video") modelType = "qianwen-video";
|
||
|
||
console.log(`[Chat] 🚀 记录生成任务: modelType=${modelType}, modelId=${currentModelId}`);
|
||
startGeneration(currentModelId, modelType);
|
||
}
|
||
}, [status, currentModelId, startGeneration]);
|
||
|
||
const searchParams = useSearchParams();
|
||
const query = searchParams.get("query");
|
||
const editImageUrl = searchParams.get("editImage");
|
||
|
||
const [hasAppendedQuery, setHasAppendedQuery] = useState(false);
|
||
const [hasLoadedEditImage, setHasLoadedEditImage] = useState(false);
|
||
|
||
useEffect(() => {
|
||
if (query && !hasAppendedQuery) {
|
||
sendMessage({
|
||
role: "user" as const,
|
||
parts: [{ type: "text", text: query }],
|
||
});
|
||
|
||
setHasAppendedQuery(true);
|
||
window.history.replaceState({}, "", `/chat/${id}`);
|
||
}
|
||
}, [query, sendMessage, hasAppendedQuery, id]);
|
||
|
||
// 🆕 处理从Gallery跳转过来的图片编辑请求
|
||
useEffect(() => {
|
||
if (editImageUrl && !hasLoadedEditImage) {
|
||
const fileName = editImageUrl.split('/').pop()?.split('?')[0] || `image-${Date.now()}.jpg`;
|
||
|
||
// 推断图片类型
|
||
let contentType = "image/jpeg";
|
||
if (fileName.toLowerCase().endsWith(".png")) contentType = "image/png";
|
||
if (fileName.toLowerCase().endsWith(".webp")) contentType = "image/webp";
|
||
if (fileName.toLowerCase().endsWith(".gif")) contentType = "image/gif";
|
||
|
||
setAttachments([{
|
||
url: decodeURIComponent(editImageUrl),
|
||
name: fileName,
|
||
contentType: contentType,
|
||
}]);
|
||
|
||
// 自动填充提示词
|
||
setInput("请帮我编辑这张图片：");
|
||
|
||
setHasLoadedEditImage(true);
|
||
window.history.replaceState({}, "", `/chat/${id}`);
|
||
}
|
||
}, [editImageUrl, hasLoadedEditImage, id, setInput]);
|
||
|
||
const { data: votes } = useSWR<Vote[]>(
|
||
messages.length >= 2 ? `/api/vote?chatId=${id}` : null,
|
||
fetcher
|
||
);
|
||
|
||
const [attachments, setAttachments] = useState<Attachment[]>([]);
|
||
const isArtifactVisible = useArtifactSelector((state) => state.isVisible);
|
||
|
||
// 处理图片编辑请求
|
||
const handleImageEdit = (imageUrl: string) => {
|
||
// 优化：直接使用 URL 构造附件对象，无需重新下载上传
|
||
// 后端 API 会在发送时处理 URL（如果需要 Base64 转换）
|
||
const fileName = imageUrl.split('/').pop()?.split('?')[0] || `image-${Date.now()}.jpg`;
|
||
|
||
// 简单的类型推断
|
||
let contentType = "image/jpeg";
|
||
if (fileName.toLowerCase().endsWith(".png")) contentType = "image/png";
|
||
if (fileName.toLowerCase().endsWith(".webp")) contentType = "image/webp";
|
||
if (fileName.toLowerCase().endsWith(".gif")) contentType = "image/gif";
|
||
|
||
setAttachments((prev) => [
|
||
...prev,
|
||
{
|
||
url: imageUrl,
|
||
name: fileName,
|
||
contentType: contentType,
|
||
},
|
||
]);
|
||
|
||
// 如果输入框为空，自动填入提示词
|
||
setInput((prev) => prev ? prev : "请帮我编辑这张图片：");
|
||
};
|
||
|
||
// 处理画图标注完成
|
||
const handleDrawingComplete = (annotatedImageUrl: string) => {
|
||
// annotatedImageUrl 是 base64 data URL 格式
|
||
const fileName = `annotated-${Date.now()}.png`;
|
||
|
||
setAttachments((prev) => [
|
||
...prev,
|
||
{
|
||
url: annotatedImageUrl,
|
||
name: fileName,
|
||
contentType: "image/png",
|
||
},
|
||
]);
|
||
|
||
// 如果输入框为空，自动填入提示词
|
||
setInput((prev) => prev ? prev : "请根据我的标注修改这张图片：");
|
||
};
|
||
|
||
useAutoResume({
|
||
autoResume,
|
||
initialMessages,
|
||
resumeStream,
|
||
setMessages,
|
||
});
|
||
|
||
return (
|
||
<>
|
||
<div className="overscroll-behavior-contain flex h-dvh min-w-0 touch-pan-y flex-col bg-background">
|
||
<ChatHeader
|
||
chatId={id}
|
||
isReadonly={isReadonly}
|
||
selectedVisibilityType={initialVisibilityType}
|
||
/>
|
||
|
||
<Messages
|
||
chatId={id}
|
||
isArtifactVisible={isArtifactVisible}
|
||
isReadonly={isReadonly}
|
||
messages={messages}
|
||
regenerate={regenerate}
|
||
sendMessage={sendMessage}
|
||
selectedModelId={currentModelId}
|
||
setMessages={setMessages}
|
||
status={status}
|
||
votes={votes}
|
||
onImageEdit={handleImageEdit}
|
||
onDrawingComplete={handleDrawingComplete}
|
||
pendingGeneration={pendingGeneration}
|
||
imageCount={imageParams?.imageCount || 1}
|
||
/>
|
||
|
||
<div className="sticky bottom-0 z-1 mx-auto w-full max-w-4xl bg-background px-2 pb-3 md:px-4 md:pb-4">
|
||
{!isReadonly && (
|
||
<MultimodalInput
|
||
attachments={attachments}
|
||
chatId={id}
|
||
input={input}
|
||
messages={messages}
|
||
onModelChange={setCurrentModelId}
|
||
selectedModelId={currentModelId}
|
||
selectedVisibilityType={visibilityType}
|
||
sendMessage={sendMessage}
|
||
setAttachments={setAttachments}
|
||
setInput={setInput}
|
||
setMessages={setMessages}
|
||
status={status}
|
||
stop={stop}
|
||
usage={usage}
|
||
isVideoModel={isVideoModel}
|
||
videoModelName={currentModel?.displayName}
|
||
videoParams={videoParams}
|
||
onVideoParamsChange={(params) => {
|
||
setVideoParams(params);
|
||
console.log("[Chat] Video params updated:", params);
|
||
}}
|
||
isImageModel={isImageModel}
|
||
imageModelName={currentModel?.displayName}
|
||
imageParams={imageParams}
|
||
onImageParamsChange={handleImageParamsChange}
|
||
imageProviderType={currentModel?.providerType}
|
||
/>
|
||
)}
|
||
</div>
|
||
</div>
|
||
|
||
<Artifact
|
||
attachments={attachments}
|
||
chatId={id}
|
||
input={input}
|
||
isReadonly={isReadonly}
|
||
messages={messages}
|
||
regenerate={regenerate}
|
||
selectedModelId={currentModelId}
|
||
selectedVisibilityType={visibilityType}
|
||
sendMessage={sendMessage}
|
||
setAttachments={setAttachments}
|
||
setInput={setInput}
|
||
setMessages={setMessages}
|
||
status={status}
|
||
stop={stop}
|
||
votes={votes}
|
||
/>
|
||
|
||
<AlertDialog
|
||
onOpenChange={setShowCreditCardAlert}
|
||
open={showCreditCardAlert}
|
||
>
|
||
<AlertDialogContent>
|
||
<AlertDialogHeader>
|
||
<AlertDialogTitle>启用 AI Gateway</AlertDialogTitle>
|
||
<AlertDialogDescription>
|
||
此应用需要
|
||
{process.env.NODE_ENV === "production" ? "项目拥有者" : "你"}
|
||
激活 Vercel AI Gateway。
|
||
</AlertDialogDescription>
|
||
</AlertDialogHeader>
|
||
<AlertDialogFooter>
|
||
<AlertDialogCancel>取消</AlertDialogCancel>
|
||
<AlertDialogAction
|
||
onClick={() => {
|
||
window.open(
|
||
"https://vercel.com/d?to=%2F%5Bteam%5D%2F%7E%2Fai%3Fmodal%3Dadd-credit-card",
|
||
"_blank"
|
||
);
|
||
window.location.href = "/";
|
||
}}
|
||
>
|
||
立即启用
|
||
</AlertDialogAction>
|
||
</AlertDialogFooter>
|
||
</AlertDialogContent>
|
||
</AlertDialog>
|
||
</>
|
||
);
|
||
}
|