deep-research / src /hooks /useDeepResearch.ts
Leon4gr45's picture
Deploy app
c16e487 verified
import { useState } from "react";
import {
streamText,
smoothStream,
type JSONValue,
type Tool,
type UserContent,
} from "ai";
import { parsePartialJson } from "@ai-sdk/ui-utils";
import { openai } from "@ai-sdk/openai";
import { type GoogleGenerativeAIProviderMetadata } from "@ai-sdk/google";
import { useTranslation } from "react-i18next";
import Plimit from "p-limit";
import { toast } from "sonner";
import useModelProvider from "@/hooks/useAiProvider";
import useWebSearch from "@/hooks/useWebSearch";
import { useTaskStore } from "@/store/task";
import { useHistoryStore } from "@/store/history";
import { useSettingStore } from "@/store/setting";
import { useKnowledgeStore } from "@/store/knowledge";
import {
parseDeepResearchPromptOverrides,
type DeepResearchPromptOverrides,
} from "@/constants/prompts";
import {
getSystemPrompt,
getOutputGuidelinesPrompt,
generateQuestionsPrompt,
writeReportPlanPrompt,
generateSerpQueriesPrompt,
processResultPrompt,
processSearchResultPrompt,
processSearchKnowledgeResultPrompt,
reviewSerpQueriesPrompt,
writeFinalReportPrompt,
getSERPQuerySchema,
} from "@/utils/deep-research/prompts";
import { isNetworkingModel } from "@/utils/model";
import { ThinkTagStreamProcessor, removeJsonMarkdown } from "@/utils/text";
import { parseError } from "@/utils/error";
import { pick, flat, unique } from "radash";
type ProviderOptions = Record<string, Record<string, JSONValue>>;
type Tools = Record<string, Tool>;
function getResponseLanguagePrompt() {
return `\n\n**Respond in the same language as the user's language**`;
}
function smoothTextStream(type: "character" | "word" | "line") {
return smoothStream({
chunking: type === "character" ? /./ : type,
delayInMs: 0,
});
}
function handleError(error: unknown) {
console.error("DeepResearch execution error:", error);
const errorMessage = parseError(error);
if (errorMessage === "[TypeError]: Failed to fetch" || errorMessage === "Failed to fetch") {
const { mode } = useSettingStore.getState();
toast.error(mode === "local" ? "Network Error: Failed to fetch. This is likely a CORS issue. Please open Settings and switch API Request Mode to Server Proxy." : "Failed to connect to server. Please check your network.");
} else {
toast.error(errorMessage);
}
}
function useDeepResearch() {
const { t } = useTranslation();
const taskStore = useTaskStore();
const { smoothTextStreamType } = useSettingStore();
const { createModelProvider, getModel } = useModelProvider();
const { search } = useWebSearch();
const [status, setStatus] = useState<string>("");
function getPromptOverrides() {
const { deepResearchPromptOverrides } = useSettingStore.getState();
try {
return parseDeepResearchPromptOverrides(deepResearchPromptOverrides);
} catch (error) {
handleError(error);
return {} as DeepResearchPromptOverrides;
}
}
function getMaxCollectionTopics() {
const { maxCollectionTopics } = useSettingStore.getState();
const value = Number(maxCollectionTopics);
if (!Number.isFinite(value)) {
return 5;
}
return Math.max(1, Math.min(20, Math.floor(value)));
}
function getAutoReviewRounds() {
const { autoReviewRounds } = useSettingStore.getState();
const value = Number(autoReviewRounds);
if (!Number.isFinite(value)) {
return 0;
}
return Math.max(0, Math.min(5, Math.floor(value)));
}
function getReportPreferenceRequirement(
reportStyle: "balanced" | "executive" | "technical" | "concise",
reportLength: "brief" | "standard" | "comprehensive"
) {
const stylePrompts: Record<
"balanced" | "executive" | "technical" | "concise",
string
> = {
balanced:
"Keep a balanced writing style with clear explanations, practical examples, and moderate technical depth.",
executive:
"Prioritize decision-ready insights. Begin sections with key findings and focus on business impact, risks, and recommendations.",
technical:
"Prioritize technical depth and precision. Include implementation details, tradeoffs, assumptions, and limitations.",
concise:
"Be concise and direct. Eliminate filler and keep each section tightly focused on essential information.",
};
const lengthPrompts: Record<"brief" | "standard" | "comprehensive", string> =
{
brief:
"Keep the report compact while preserving critical insights and evidence.",
standard:
"Write a standard-length report with good depth and practical detail.",
comprehensive:
"Write a comprehensive report with deep coverage, detailed analysis, and thorough supporting context.",
};
return [
"Additional report preferences:",
`- Style: ${stylePrompts[reportStyle]}`,
`- Length: ${lengthPrompts[reportLength]}`,
].join("\n");
}
async function generateSearchSettings(searchModel: string) {
const { provider, enableSearch, searchProvider, searchMaxResult } =
useSettingStore.getState();
if (enableSearch === "1" && searchProvider === "model") {
const createModel = (model: string) => {
// Enable Gemini's built-in search tool
if (
["google", "google-vertex"].includes(provider) &&
isNetworkingModel(model)
) {
return createModelProvider(model, { useSearchGrounding: true });
} else {
return createModelProvider(model);
}
};
const getTools = (model: string) => {
// Enable OpenAI's built-in search tool
if (
["openai", "azure", "openaicompatible"].includes(provider) &&
(model.startsWith("gpt-4o") ||
model.startsWith("gpt-4.1") ||
model.startsWith("gpt-5"))
) {
return {
web_search_preview: openai.tools.webSearchPreview({
// optional configuration:
searchContextSize: searchMaxResult > 5 ? "high" : "medium",
}),
} as Tools;
}
};
const getProviderOptions = (model: string) => {
// Enable OpenRouter's built-in search tool
if (provider === "openrouter") {
return {
openrouter: {
plugins: [
{
id: "web",
max_results: searchMaxResult, // Defaults to 5
},
],
},
} as ProviderOptions;
} else if (
provider === "xai" &&
model.startsWith("grok-3") &&
!model.includes("mini")
) {
return {
xai: {
search_parameters: {
mode: "auto",
max_search_results: searchMaxResult,
},
},
} as ProviderOptions;
}
};
return {
model: await createModel(searchModel),
tools: getTools(searchModel),
providerOptions: getProviderOptions(searchModel),
};
} else {
return {
model: await createModelProvider(searchModel),
};
}
}
async function askQuestions() {
const { question } = useTaskStore.getState();
const { thinkingModel } = getModel();
setStatus(t("research.common.thinking"));
console.log(`[askQuestions] Starting with question: "${question.substring(0, 50)}...", model: ${thinkingModel}`);
const thinkTagStreamProcessor = new ThinkTagStreamProcessor();
const promptOverrides = getPromptOverrides();
let searchSettings;
try {
searchSettings = await generateSearchSettings(thinkingModel);
console.log(`[askQuestions] Generated search settings successfully`);
} catch (err) {
console.error(`[askQuestions] Failed to generate search settings:`, err);
throw err;
}
try {
console.log(`[askQuestions] Initiating streamText call...`);
const result = streamText({
...searchSettings,
system: getSystemPrompt(promptOverrides),
prompt: [
generateQuestionsPrompt(question, promptOverrides),
getResponseLanguagePrompt(),
].join("\n\n"),
experimental_transform: smoothTextStream(smoothTextStreamType),
onError: handleError,
});
let content = "";
let reasoning = "";
taskStore.setQuestion(question);
console.log(`[askQuestions] Awaiting stream parts...`);
for await (const part of result.fullStream) {
if (part.type === "text-delta") {
thinkTagStreamProcessor.processChunk(
part.textDelta,
(data) => {
content += data;
taskStore.updateQuestions(content);
},
(data) => {
reasoning += data;
}
);
} else if (part.type === "reasoning") {
reasoning += part.textDelta;
}
}
console.log(`[askQuestions] Stream completed.`);
if (reasoning) console.log(`[askQuestions] Reasoning length: ${reasoning.length}`);
} catch (err) {
console.error(`[askQuestions] Error during stream processing:`, err);
throw err;
}
}
async function writeReportPlan() {
const { query } = useTaskStore.getState();
const { thinkingModel } = getModel();
setStatus(t("research.common.thinking"));
const thinkTagStreamProcessor = new ThinkTagStreamProcessor();
const promptOverrides = getPromptOverrides();
const searchSettings = await generateSearchSettings(thinkingModel);
const result = streamText({
...searchSettings,
system: getSystemPrompt(promptOverrides),
prompt: [
writeReportPlanPrompt(query, promptOverrides),
getResponseLanguagePrompt(),
].join("\n\n"),
experimental_transform: smoothTextStream(smoothTextStreamType),
onError: handleError,
});
let content = "";
let reasoning = "";
for await (const part of result.fullStream) {
if (part.type === "text-delta") {
thinkTagStreamProcessor.processChunk(
part.textDelta,
(data) => {
content += data;
taskStore.updateReportPlan(content);
},
(data) => {
reasoning += data;
}
);
} else if (part.type === "reasoning") {
reasoning += part.textDelta;
}
}
if (reasoning) console.log(reasoning);
return content;
}
async function searchLocalKnowledges(
query: string,
researchGoal: string,
promptOverrides: DeepResearchPromptOverrides = {}
) {
const { resources } = useTaskStore.getState();
const knowledgeStore = useKnowledgeStore.getState();
const knowledges: Knowledge[] = [];
for (const item of resources) {
if (item.status === "completed") {
const resource = knowledgeStore.get(item.id);
if (resource) {
knowledges.push(resource);
}
}
}
const { networkingModel } = getModel();
const thinkTagStreamProcessor = new ThinkTagStreamProcessor();
const searchResult = streamText({
model: await createModelProvider(networkingModel),
system: getSystemPrompt(promptOverrides),
prompt: [
processSearchKnowledgeResultPrompt(
query,
researchGoal,
knowledges,
promptOverrides
),
getResponseLanguagePrompt(),
].join("\n\n"),
experimental_transform: smoothTextStream(smoothTextStreamType),
onError: handleError,
});
let content = "";
let reasoning = "";
for await (const part of searchResult.fullStream) {
if (part.type === "text-delta") {
thinkTagStreamProcessor.processChunk(
part.textDelta,
(data) => {
content += data;
taskStore.updateTask(query, { learning: content });
},
(data) => {
reasoning += data;
}
);
} else if (part.type === "reasoning") {
reasoning += part.textDelta;
}
}
if (reasoning) console.log(reasoning);
return content;
}
const [abortController, setAbortController] = useState<AbortController | null>(null);
const abortResearch = () => {
if (abortController) {
abortController.abort();
}
};
async function runSearchTask(queries: SearchTask[]) {
const {
enableSearch,
searchProvider,
parallelSearch,
references,
onlyUseLocalResource,
} = useSettingStore.getState();
const { resources } = useTaskStore.getState();
const { networkingModel } = getModel();
const promptOverrides = getPromptOverrides();
setStatus(t("research.common.research"));
const plimit = Plimit(parallelSearch);
const controller = new AbortController();
setAbortController(controller);
let completedCount = 0;
let timeoutId: NodeJS.Timeout | null = null;
const thinkTagStreamProcessor = new ThinkTagStreamProcessor();
await Promise.all(
queries.map((item) => {
return plimit(async () => {
if (controller.signal.aborted) {
taskStore.updateTask(item.query, { state: "failed", learning: "Aborted by user or timeout" });
return "";
}
let content = "";
let reasoning = "";
let searchResult;
let sources: Source[] = [];
let images: ImageSource[] = [];
taskStore.updateTask(item.query, { state: "processing" });
if (resources.length > 0) {
const knowledges = await searchLocalKnowledges(
item.query,
item.researchGoal,
promptOverrides
);
content += [
knowledges,
`### ${t("research.searchResult.references")}`,
resources.map((item) => `- ${item.name}`).join("\n"),
].join("\n\n");
if (onlyUseLocalResource === "enable") {
taskStore.updateTask(item.query, {
state: "completed",
learning: content,
sources,
images,
});
completedCount++;
if (completedCount / queries.length >= 0.8 && !timeoutId) {
timeoutId = setTimeout(() => {
controller.abort();
}, 600000);
}
return content;
} else {
content += "\n\n---\n\n";
}
}
if (enableSearch === "1") {
if (searchProvider !== "model") {
try {
const results = await search(item.query);
sources = results.sources;
images = results.images;
if (sources.length === 0) {
throw new Error("Invalid Search Results");
}
} catch (err) {
console.error(err);
handleError(
`[${searchProvider}]: ${
err instanceof Error ? err.message : "Search Failed"
}`
);
return plimit.clearQueue();
}
const enableReferences =
sources.length > 0 && references === "enable";
searchResult = streamText({
model: await createModelProvider(networkingModel),
system: getSystemPrompt(promptOverrides),
prompt: [
processSearchResultPrompt(
item.query,
item.researchGoal,
sources,
enableReferences,
promptOverrides
),
getResponseLanguagePrompt(),
].join("\n\n"),
experimental_transform: smoothTextStream(smoothTextStreamType),
abortSignal: controller.signal,
onError: handleError,
});
} else {
const searchSettings = await generateSearchSettings(
networkingModel
);
searchResult = streamText({
...searchSettings,
system: getSystemPrompt(promptOverrides),
prompt: [
processResultPrompt(
item.query,
item.researchGoal,
promptOverrides
),
getResponseLanguagePrompt(),
].join("\n\n"),
experimental_transform: smoothTextStream(smoothTextStreamType),
abortSignal: controller.signal,
onError: handleError,
});
}
} else {
searchResult = streamText({
model: await createModelProvider(networkingModel),
system: getSystemPrompt(promptOverrides),
prompt: [
processResultPrompt(
item.query,
item.researchGoal,
promptOverrides
),
getResponseLanguagePrompt(),
].join("\n\n"),
experimental_transform: smoothTextStream(smoothTextStreamType),
abortSignal: controller.signal,
onError: (err) => {
taskStore.updateTask(item.query, { state: "failed" });
handleError(err);
},
});
}
try {
for await (const part of searchResult.fullStream) {
if (part.type === "text-delta") {
thinkTagStreamProcessor.processChunk(
part.textDelta,
(data) => {
content += data;
taskStore.updateTask(item.query, { learning: content });
},
(data) => {
reasoning += data;
}
);
} else if (part.type === "reasoning") {
reasoning += part.textDelta;
} else if (part.type === "source") {
sources.push(part.source);
} else if (part.type === "finish") {
if (part.providerMetadata?.google) {
const { groundingMetadata } = part.providerMetadata.google;
const googleGroundingMetadata =
groundingMetadata as GoogleGenerativeAIProviderMetadata["groundingMetadata"];
if (googleGroundingMetadata?.groundingSupports) {
googleGroundingMetadata.groundingSupports.forEach(
({ segment, groundingChunkIndices }) => {
if (segment.text && groundingChunkIndices) {
const index = groundingChunkIndices.map(
(idx: number) => `[${idx + 1}]`
);
content = content.replaceAll(
segment.text,
`${segment.text}${index.join("")}`
);
}
}
);
}
} else if (part.providerMetadata?.openai) {
// Fixed the problem that OpenAI cannot generate markdown reference link syntax properly in Chinese context
content = content.replaceAll("【", "[").replaceAll("】", "]");
}
}
}
} catch (err: any) {
if (err.name === 'AbortError') {
console.log(`[runSearchTask] Aborted query: ${item.query}`);
} else {
throw err;
}
}
if (reasoning) console.log(reasoning);
if (sources.length > 0) {
content +=
"\n\n" +
sources
.map(
(item, idx) =>
`[${idx + 1}]: ${item.url}${
item.title ? ` "${item.title.replaceAll('"', " ")}"` : ""
}`
)
.join("\n");
}
completedCount++;
if (completedCount / queries.length >= 0.8 && !timeoutId) {
timeoutId = setTimeout(() => {
controller.abort();
}, 600000);
}
if (content.length > 0) {
taskStore.updateTask(item.query, {
state: "completed",
learning: content,
sources,
images,
});
return content;
} else {
taskStore.updateTask(item.query, {
state: "failed",
learning: "Aborted or failed",
sources: [],
images: [],
});
return "";
}
});
})
);
if (timeoutId) clearTimeout(timeoutId);
setAbortController(null);
}
async function reviewSearchResult() {
const { reportPlan, tasks, suggestion } = useTaskStore.getState();
const { thinkingModel } = getModel();
const promptOverrides = getPromptOverrides();
setStatus(t("research.common.research"));
const learnings = tasks.map((item) => item.learning);
const thinkTagStreamProcessor = new ThinkTagStreamProcessor();
const result = streamText({
model: await createModelProvider(thinkingModel),
system: getSystemPrompt(promptOverrides),
prompt: [
reviewSerpQueriesPrompt(
reportPlan,
learnings,
suggestion,
promptOverrides
),
getResponseLanguagePrompt(),
].join("\n\n"),
experimental_transform: smoothTextStream(smoothTextStreamType),
onError: handleError,
});
const querySchema = getSERPQuerySchema();
let content = "";
let reasoning = "";
let queries: SearchTask[] = [];
for await (const textPart of result.textStream) {
thinkTagStreamProcessor.processChunk(
textPart,
(text) => {
content += text;
const data: PartialJson = parsePartialJson(
removeJsonMarkdown(content)
);
if (
querySchema.safeParse(data.value) &&
data.state === "successful-parse"
) {
if (data.value) {
queries = data.value.map(
(item: { query: string; researchGoal: string }) => ({
state: "unprocessed",
learning: "",
...pick(item, ["query", "researchGoal"]),
})
);
queries = queries.slice(0, getMaxCollectionTopics());
}
}
},
(text) => {
reasoning += text;
}
);
}
if (reasoning) console.log(reasoning);
if (queries.length > 0) {
taskStore.update([...tasks, ...queries]);
await runSearchTask(queries);
return queries.length;
}
return 0;
}
async function forceFinishAndWriteReport() {
abortResearch();
// Use setTimeout to allow state to settle before initiating write
setTimeout(() => {
writeFinalReport();
}, 500);
}
async function writeFinalReport() {
const {
citationImage,
references,
useFileFormatResource,
reportStyle,
reportLength,
} = useSettingStore.getState();
const {
reportPlan,
tasks,
setId,
setTitle,
setSources,
requirement,
updateFinalReport,
} = useTaskStore.getState();
const { save } = useHistoryStore.getState();
const { thinkingModel } = getModel();
const promptOverrides = getPromptOverrides();
setStatus(t("research.common.writing"));
updateFinalReport("");
setTitle("");
setSources([]);
const learnings = tasks.map((item) => item.learning);
const sources: Source[] = unique(
flat(tasks.map((item) => item.sources || [])),
(item) => item.url
);
const images: ImageSource[] = unique(
flat(tasks.map((item) => item.images || [])),
(item) => item.url
);
const enableCitationImage = images.length > 0 && citationImage === "enable";
const enableReferences = sources.length > 0 && references === "enable";
const enableFileFormatResource = useFileFormatResource === "enable";
const mergedRequirement = [
requirement,
getReportPreferenceRequirement(reportStyle, reportLength),
]
.filter((item) => item.trim().length > 0)
.join("\n\n");
const thinkTagStreamProcessor = new ThinkTagStreamProcessor();
const sourceList = enableReferences
? sources.map((item) => pick(item, ["title", "url"]))
: [];
const imageList = enableCitationImage ? images : [];
const file = new File(
[
[
`<LEARNINGS>\n${learnings
.map((detail) => `<learning>\n${detail}\n</learning>`)
.join("\n")}\n</LEARNINGS>`,
`<SOURCES>\n${sourceList
.map(
(item, idx) =>
`<source index="${idx + 1}" url="${item.url}">\n${
item.title
}\n</source>`
)
.join("\n")}\n</SOURCES>`,
`<IMAGES>\n${imageList
.map(
(source, idx) =>
`${idx + 1}. ![${source.description}](${source.url})`
)
.join("\n")}\n</IMAGES>`,
].join("\n\n"),
],
"resources.md",
{ type: "text/markdown" }
);
const fileData = await file.arrayBuffer();
const messageContent: UserContent = [
{
type: "text",
text: [
writeFinalReportPrompt(
reportPlan,
learnings,
sourceList,
imageList,
mergedRequirement,
enableCitationImage,
enableReferences,
enableFileFormatResource,
promptOverrides
),
getResponseLanguagePrompt(),
].join("\n\n"),
},
];
if (enableFileFormatResource) {
messageContent.push({
type: "file",
mimeType: "text/markdown",
filename: "resources.md",
data: fileData,
});
}
const result = streamText({
model: await createModelProvider(thinkingModel),
system: [
getSystemPrompt(promptOverrides),
getOutputGuidelinesPrompt(promptOverrides),
].join("\n\n"),
messages: [
{
role: "user",
content: messageContent,
},
],
temperature: 0.5,
experimental_transform: smoothTextStream(smoothTextStreamType),
onError: handleError,
});
let content = "";
let reasoning = "";
for await (const part of result.fullStream) {
if (part.type === "text-delta") {
thinkTagStreamProcessor.processChunk(
part.textDelta,
(data) => {
content += data;
updateFinalReport(content);
},
(data) => {
reasoning += data;
}
);
} else if (part.type === "reasoning") {
reasoning += part.textDelta;
}
}
if (reasoning) console.log(reasoning);
if (sources.length > 0) {
content +=
"\n\n" +
sources
.map(
(item, idx) =>
`[${idx + 1}]: ${item.url}${
item.title ? ` "${item.title.replaceAll('"', " ")}"` : ""
}`
)
.join("\n");
updateFinalReport(content);
}
if (content.length > 0) {
const title = (content || "")
.split("\n")[0]
.replaceAll("#", "")
.replaceAll("*", "")
.trim();
setTitle(title);
setSources(sources);
const id = save(taskStore.backup());
setId(id);
return content;
} else {
return "";
}
}
async function deepResearch() {
const { reportPlan } = useTaskStore.getState();
const { thinkingModel } = getModel();
const promptOverrides = getPromptOverrides();
setStatus(t("research.common.thinking"));
try {
const thinkTagStreamProcessor = new ThinkTagStreamProcessor();
const result = streamText({
model: await createModelProvider(thinkingModel),
system: getSystemPrompt(promptOverrides),
prompt: [
generateSerpQueriesPrompt(reportPlan, promptOverrides),
getResponseLanguagePrompt(),
].join("\n\n"),
experimental_transform: smoothTextStream(smoothTextStreamType),
onError: handleError,
});
const querySchema = getSERPQuerySchema();
let content = "";
let reasoning = "";
let queries: SearchTask[] = [];
for await (const textPart of result.textStream) {
thinkTagStreamProcessor.processChunk(
textPart,
(text) => {
content += text;
const data: PartialJson = parsePartialJson(
removeJsonMarkdown(content)
);
if (querySchema.safeParse(data.value)) {
if (
data.state === "repaired-parse" ||
data.state === "successful-parse"
) {
if (data.value) {
queries = data.value.map(
(item: { query: string; researchGoal: string }) => ({
state: "unprocessed",
learning: "",
...pick(item, ["query", "researchGoal"]),
})
);
queries = queries.slice(0, getMaxCollectionTopics());
taskStore.update(queries);
}
}
}
},
(text) => {
reasoning += text;
}
);
}
if (reasoning) console.log(reasoning);
if (queries.length > 0) {
await runSearchTask(queries);
let remainingAutoRounds = getAutoReviewRounds();
while (remainingAutoRounds > 0) {
const generatedQueries = await reviewSearchResult();
if (generatedQueries === 0) {
break;
}
remainingAutoRounds -= 1;
}
}
} catch (err) {
console.error(err);
}
}
return {
status,
deepResearch,
askQuestions,
writeReportPlan,
runSearchTask,
reviewSearchResult,
writeFinalReport,
abortResearch,
forceFinishAndWriteReport,
};
}
export default useDeepResearch;