Exched/DeepSeek_Coder
0
1import { useQuery, useQueryClient, useMutation } from "@tanstack/react-query";2import { useMemo, useRef, useState } from "react";3import { toast } from "sonner";4import { useLocalStorage } from "react-use";5 6import { MODELS } from "@/lib/providers";7import { useEditor } from "./useEditor";8import { Page, EnhancedSettings } from "@/types";9import { api } from "@/lib/api";10import { useRouter } from "next/navigation";11import { useUser } from "./useUser";12import { isTheSameHtml } from "@/lib/compare-html-diff";13 14export const useAi = (onScrollToBottom?: () => void) => {15 const client = useQueryClient();16 const audio = useRef<HTMLAudioElement | null>(null);17 const { setPages, setCurrentPage, setPrompts, prompts, pages, project, setProject, commits, setCommits, setLastSavedPages, isSameHtml } = useEditor();18 const [controller, setController] = useState<AbortController | null>(null);19 const [storageProvider, setStorageProvider] = useLocalStorage("provider", "auto");20 const [storageModel, setStorageModel] = useLocalStorage("model", MODELS[0].value);21 const router = useRouter();22 const { projects, setProjects } = useUser();23 24 const { data: isAiWorking = false } = useQuery({25 queryKey: ["ai.isAiWorking"],26 queryFn: async () => false,27 refetchOnWindowFocus: false,28 refetchOnReconnect: false,29 refetchOnMount: false,30 });31 const setIsAiWorking = (newIsAiWorking: boolean) => {32 client.setQueryData(["ai.isAiWorking"], newIsAiWorking);33 };34 35 const { data: isThinking = false } = useQuery({36 queryKey: ["ai.isThinking"],37 queryFn: async () => false,38 refetchOnWindowFocus: false,39 refetchOnReconnect: false,40 refetchOnMount: false,41 });42 const setIsThinking = (newIsThinking: boolean) => {43 client.setQueryData(["ai.isThinking"], newIsThinking);44 };45 46 const { data: selectedElement } = useQuery<HTMLElement | null>({47 queryKey: ["ai.selectedElement"],48 queryFn: async () => null,49 refetchOnWindowFocus: false,50 refetchOnReconnect: false,51 refetchOnMount: false,52 initialData: null53 });54 const setSelectedElement = (newSelectedElement: HTMLElement | null) => {55 client.setQueryData(["ai.selectedElement"], newSelectedElement);56 };57 58 const { data: isEditableModeEnabled = false } = useQuery({59 queryKey: ["ai.isEditableModeEnabled"],60 queryFn: async () => false,61 refetchOnWindowFocus: false,62 refetchOnReconnect: false,63 refetchOnMount: false,64 });65 const setIsEditableModeEnabled = (newIsEditableModeEnabled: boolean) => {66 client.setQueryData(["ai.isEditableModeEnabled"], newIsEditableModeEnabled);67 };68 69 const { data: selectedFiles } = useQuery<string[]>({70 queryKey: ["ai.selectedFiles"],71 queryFn: async () => [],72 refetchOnWindowFocus: false,73 refetchOnReconnect: false,74 refetchOnMount: false,75 initialData: []76 });77 const setSelectedFiles = (newFiles: string[]) => {78 client.setQueryData(["ai.selectedFiles"], newFiles)79 };80 81 const { data: provider } = useQuery({82 queryKey: ["ai.provider"],83 queryFn: async () => storageProvider ?? "auto",84 refetchOnWindowFocus: false,85 refetchOnReconnect: false,86 refetchOnMount: false,87 initialData: storageProvider ?? "auto"88 });89 const setProvider = (newProvider: string) => {90 setStorageProvider(newProvider);91 client.setQueryData(["ai.provider"], newProvider);92 };93 94 const { data: model } = useQuery({95 queryKey: ["ai.model"],96 queryFn: async () => {97 // check if the model exist in the MODELS array98 const selectedModel = MODELS.find(m => m.value === storageModel || m.label === storageModel);99 if (selectedModel) {100 return selectedModel.value;101 }102 return MODELS[0].value;103 },104 refetchOnWindowFocus: false,105 refetchOnReconnect: false,106 refetchOnMount: false,107 initialData: undefined,108 });109 const setModel = (newModel: string) => {110 setStorageModel(newModel);111 client.setQueryData(["ai.model"], newModel);112 };113 114 const createNewProject = async (prompt: string, htmlPages: Page[], projectName: string | undefined, isLoggedIn?: boolean) => {115 if (isLoggedIn) {116 const response = await api.post("/me/projects", {117 title: projectName,118 pages: htmlPages,119 prompt,120 });121 if (response.data.ok) {122 setIsAiWorking(false);123 router.replace(`/${response.data.space.project.space_id}`);124 setProject(response.data.space);125 setProjects([...projects, response.data.space]);126 toast.success("AI responded successfully");127 if (audio.current) audio.current.play();128 }129 } else {130 setIsAiWorking(false);131 toast.success("AI responded successfully");132 if (audio.current) audio.current.play();133 }134 }135 136 const callAiNewProject = async (prompt: string, enhancedSettings?: EnhancedSettings, redesignMarkdown?: string, isLoggedIn?: boolean) => {137 if (isAiWorking) return;138 if (!redesignMarkdown && !prompt.trim()) return;139 140 setIsAiWorking(true);141 142 const abortController = new AbortController();143 setController(abortController);144 145 try {146 const request = await fetch("/api/ask", {147 method: "POST",148 body: JSON.stringify({149 prompt,150 provider,151 model,152 redesignMarkdown,153 enhancedSettings,154 }),155 headers: {156 "Content-Type": "application/json",157 "x-forwarded-for": window.location.hostname,158 },159 signal: abortController.signal,160 });161 162 if (request && request.body) {163 const reader = request.body.getReader();164 const decoder = new TextDecoder("utf-8");165 let contentResponse = "";166 167 const read = async (): Promise<any> => {168 const { done, value } = await reader.read();169 170 if (done) {171 const trimmedResponse = contentResponse.trim();172 if (trimmedResponse.startsWith("{") && trimmedResponse.endsWith("}")) {173 try {174 const jsonResponse = JSON.parse(trimmedResponse);175 if (jsonResponse && !jsonResponse.ok) {176 setIsAiWorking(false);177 if (jsonResponse.openLogin) {178 return { error: "login_required" };179 } else if (jsonResponse.openSelectProvider) {180 return { error: "provider_required", message: jsonResponse.message };181 } else if (jsonResponse.openProModal) {182 return { error: "pro_required" };183 } else {184 toast.error(jsonResponse.message);185 return { error: "api_error", message: jsonResponse.message };186 }187 }188 } catch (e) {189 // Not valid JSON, treat as normal content190 }191 }192 193 const newPages = formatPages(contentResponse);194 let projectName = contentResponse.match(/<<<<<<< PROJECT_NAME_START ([\s\S]*?) >>>>>>> PROJECT_NAME_END/)?.[1]?.trim();195 if (!projectName) {196 projectName = prompt.substring(0, 40).replace(/[^a-zA-Z0-9]/g, "-").slice(0, 40);197 }198 setPages(newPages);199 setLastSavedPages([...newPages]);200 if (newPages.length > 0 && !isTheSameHtml(newPages[0].html)) {201 createNewProject(prompt, newPages, projectName, isLoggedIn);202 }203 setPrompts([...prompts, prompt]);204 205 return { success: true, pages: newPages };206 }207 208 const chunk = decoder.decode(value, { stream: true });209 contentResponse += chunk;210 211 const trimmedResponse = contentResponse.trim();212 if (trimmedResponse.startsWith("{") && trimmedResponse.endsWith("}")) {213 try {214 const jsonResponse = JSON.parse(trimmedResponse);215 if (jsonResponse && !jsonResponse.ok) {216 setIsAiWorking(false);217 if (jsonResponse.openLogin) {218 return { error: "login_required" };219 } else if (jsonResponse.openSelectProvider) {220 return { error: "provider_required", message: jsonResponse.message };221 } else if (jsonResponse.openProModal) {222 return { error: "pro_required" };223 } else {224 toast.error(jsonResponse.message);225 return { error: "api_error", message: jsonResponse.message };226 }227 }228 } catch (e) {229 // Not a complete JSON yet, continue reading230 }231 }232 233 formatPages(contentResponse);234 235 // Continue reading236 return read();237 };238 239 return await read();240 }241 // eslint-disable-next-line @typescript-eslint/no-explicit-any242 } catch (error: any) {243 setIsAiWorking(false);244 setIsThinking(false);245 setController(null);246 247 if (!abortController.signal.aborted) {248 toast.error(error.message || "Network error occurred");249 }250 251 if (error.openLogin) {252 return { error: "login_required" };253 }254 return { error: "network_error", message: error.message };255 }256 };257 258 const callAiFollowUp = async (prompt: string, enhancedSettings?: EnhancedSettings, isNew?: boolean) => {259 if (isAiWorking) return;260 if (!prompt.trim()) return;261 262 263 setIsAiWorking(true);264 265 const abortController = new AbortController();266 setController(abortController);267 268 try {269 const request = await fetch("/api/ask", {270 method: "PUT",271 body: JSON.stringify({272 prompt,273 provider,274 previousPrompts: prompts,275 model,276 pages,277 selectedElementHtml: selectedElement?.outerHTML,278 files: selectedFiles,279 repoId: project?.space_id,280 isNew,281 enhancedSettings,282 }),283 headers: {284 "Content-Type": "application/json",285 "x-forwarded-for": window.location.hostname,286 },287 signal: abortController.signal,288 });289 290 if (request && request.body) {291 const res = await request.json();292 293 if (!request.ok) {294 if (res.openLogin) {295 setIsAiWorking(false);296 return { error: "login_required" };297 } else if (res.openSelectProvider) {298 setIsAiWorking(false);299 return { error: "provider_required", message: res.message };300 } else if (res.openProModal) {301 setIsAiWorking(false);302 return { error: "pro_required" };303 } else {304 toast.error(res.message);305 setIsAiWorking(false);306 return { error: "api_error", message: res.message };307 }308 }309 310 toast.success("AI responded successfully");311 const iframe = document.getElementById(312 "preview-iframe"313 ) as HTMLIFrameElement;314 315 if (isNew && res.repoId) {316 router.push(`/${res.repoId}`);317 setIsAiWorking(false);318 } else {319 setPages(res.pages);320 setLastSavedPages([...res.pages]); // Mark AI changes as saved321 setCommits([res.commit, ...commits]);322 setPrompts(323 [...prompts, prompt]324 )325 setSelectedElement(null);326 setSelectedFiles([]);327 setIsEditableModeEnabled(false);328 setIsAiWorking(false); // This was missing!329 }330 331 if (audio.current) audio.current.play();332 if (iframe) {333 setTimeout(() => {334 iframe.src = iframe.src;335 }, 500);336 }337 338 return { success: true, html: res.html, updatedLines: res.updatedLines };339 }340 // eslint-disable-next-line @typescript-eslint/no-explicit-any341 } catch (error: any) {342 setIsAiWorking(false);343 toast.error(error.message);344 if (error.openLogin) {345 return { error: "login_required" };346 }347 return { error: "network_error", message: error.message };348 }349 };350 351 const formatPages = (content: string) => {352 const pages: Page[] = [];353 if (!content.match(/<<<<<<< START_TITLE (.*?) >>>>>>> END_TITLE/)) {354 return pages;355 }356 357 const cleanedContent = content.replace(358 /[\s\S]*?<<<<<<< START_TITLE (.*?) >>>>>>> END_TITLE/,359 "<<<<<<< START_TITLE $1 >>>>>>> END_TITLE"360 );361 const htmlChunks = cleanedContent.split(362 /<<<<<<< START_TITLE (.*?) >>>>>>> END_TITLE/363 );364 const processedChunks = new Set<number>();365 366 htmlChunks.forEach((chunk, index) => {367 if (processedChunks.has(index) || !chunk?.trim()) {368 return;369 }370 const htmlContent = extractHtmlContent(htmlChunks[index + 1]);371 372 if (htmlContent) {373 const page: Page = {374 path: chunk.trim(),375 html: htmlContent,376 };377 pages.push(page);378 379 if (htmlContent.length > 200) {380 onScrollToBottom?.();381 }382 383 processedChunks.add(index);384 processedChunks.add(index + 1);385 }386 });387 if (pages.length > 0) {388 setPages(pages);389 const lastPagePath = pages[pages.length - 1]?.path;390 setCurrentPage(lastPagePath || "index.html");391 }392 393 return pages;394 };395 396 const formatPage = (content: string, currentPagePath: string) => {397 if (!content.match(/<<<<<<< START_TITLE (.*?) >>>>>>> END_TITLE/)) {398 return null;399 }400 401 const cleanedContent = content.replace(402 /[\s\S]*?<<<<<<< START_TITLE (.*?) >>>>>>> END_TITLE/,403 "<<<<<<< START_TITLE $1 >>>>>>> END_TITLE"404 );405 406 const htmlChunks = cleanedContent.split(407 /<<<<<<< START_TITLE (.*?) >>>>>>> END_TITLE/408 )?.filter(Boolean);409 410 const pagePath = htmlChunks[0]?.trim() || "";411 const htmlContent = extractHtmlContent(htmlChunks[1]);412 413 if (!pagePath || !htmlContent) {414 return null;415 }416 417 const page: Page = {418 path: pagePath,419 html: htmlContent,420 };421 422 setPages(prevPages => {423 const existingPageIndex = prevPages.findIndex(p => p.path === currentPagePath || p.path === pagePath);424 425 if (existingPageIndex !== -1) {426 const updatedPages = [...prevPages];427 updatedPages[existingPageIndex] = page;428 return updatedPages;429 } else {430 return [...prevPages, page];431 }432 });433 434 setCurrentPage(pagePath);435 436 if (htmlContent.length > 200) {437 onScrollToBottom?.();438 }439 440 return page;441 };442 443 const extractHtmlContent = (chunk: string): string => {444 if (!chunk) return "";445 const htmlMatch = chunk.trim().match(/<!DOCTYPE html>[\s\S]*/);446 if (!htmlMatch) return "";447 let htmlContent = htmlMatch[0];448 htmlContent = ensureCompleteHtml(htmlContent);449 htmlContent = htmlContent.replace(/```/g, "");450 return htmlContent;451 };452 453 const ensureCompleteHtml = (html: string): string => {454 let completeHtml = html;455 if (completeHtml.includes("<head>") && !completeHtml.includes("</head>")) {456 completeHtml += "\n</head>";457 }458 if (completeHtml.includes("<body") && !completeHtml.includes("</body>")) {459 completeHtml += "\n</body>";460 }461 if (!completeHtml.includes("</html>")) {462 completeHtml += "\n</html>";463 }464 return completeHtml;465 };466 467 const cancelRequest = () => {468 if (controller) {469 controller.abort();470 setController(null);471 }472 setIsAiWorking(false);473 setIsThinking(false);474 };475 476 const selectedModel = useMemo(() => {477 return MODELS.find(m => m.value === model || m.label === model);478 }, [model]);479 480 return {481 isThinking,482 setIsThinking,483 callAiNewProject,484 callAiFollowUp,485 isAiWorking,486 setIsAiWorking,487 selectedElement,488 setSelectedElement,489 selectedFiles,490 setSelectedFiles,491 isEditableModeEnabled,492 setIsEditableModeEnabled,493 globalAiLoading: isThinking || isAiWorking,494 cancelRequest,495 model,496 setModel,497 provider,498 setProvider,499 selectedModel,500 audio,501 };502}