deerflow-code/frontend-web/src/core/studio/hooks.ts

import { useState } from "react";
import { useQuery } from "@tanstack/react-query";
import { useStream } from "@langchain/langgraph-sdk/react";

import { getStoredAuth } from "@/core/auth";
import { getAPIClient } from "@/core/api";
import { fetchNotebookSources, fetchStudioNotebooks, fetchStudioToken } from "./api";
import type { AgentThreadState } from "@/core/threads/types";
import type { SourceItem, StudioNotebook } from "./types";

export function useStudioNotebooks(): {
  notebooks: StudioNotebook[];
  isLoading: boolean;
  isError: boolean;
  token: string | undefined;
} {
  const email = getStoredAuth()?.email ?? "";

  const tokenQuery = useQuery({
    queryKey: ["studio", "token", email],
    queryFn: () => fetchStudioToken(email),
    enabled: Boolean(email),
    staleTime: 5 * 60 * 1000,
    retry: 1,
  });

  const notebooksQuery = useQuery({
    queryKey: ["studio", "notebooks", tokenQuery.data],
    queryFn: () => fetchStudioNotebooks(tokenQuery.data!),
    enabled: Boolean(tokenQuery.data),
    staleTime: 60 * 1000,
    retry: 1,
  });

  return {
    notebooks: (notebooksQuery.data ?? []).filter((nb) => nb.is_owner),
    isLoading: tokenQuery.isLoading || notebooksQuery.isLoading,
    isError: tokenQuery.isError || notebooksQuery.isError,
    token: tokenQuery.data,
  };
}

export function useNotebookThreads(notebookId: number | null) {
  const client = getAPIClient();
  return useQuery({
    queryKey: ["notebook-threads", notebookId],
    queryFn: () =>
      client.threads.search({
        metadata: { thread_type: "notebook", note_id: notebookId },
        limit: 20,
        sortBy: "updated_at",
        sortOrder: "desc",
      }),
    enabled: notebookId !== null,
    staleTime: 30 * 1000,
    refetchOnWindowFocus: false,
  });
}

export function useNotebookChat(
  notebookId: number,
  sources: SourceItem[],
  modelName?: string,
  initialThreadId?: string,
) {
  const [threadId, setThreadId] = useState<string | undefined>(initialThreadId);

  const email = getStoredAuth()?.email ?? "";
  const { data: token } = useQuery({
    queryKey: ["studio", "token", email],
    queryFn: () => fetchStudioToken(email),
    enabled: Boolean(email),
    staleTime: 5 * 60 * 1000,
    retry: 1,
  });

  const thread = useStream<AgentThreadState>({
    client: getAPIClient(),
    assistantId: "lead_agent",
    threadId,
    fetchStateHistory: { limit: 1 },
    onCreated(meta) {
      setThreadId(meta.thread_id);
      void getAPIClient()
        .threads.update(meta.thread_id, {
          metadata: { thread_type: "notebook", note_id: notebookId },
        })
        .catch(() => ({}));
    },
  });

  const send = async (text: string, selectedSourceIds: number[]) => {
    if (thread.isLoading || !token) return;
    // API 需要 document.id，而不是 source.id
    const toDocIds = (ids: number[]) =>
      ids
        .map((id) => sources.find((s) => s.id === id)?.document?.id)
        .filter((id): id is number => id !== undefined);
    const sourceIds =
      selectedSourceIds.length > 0
        ? toDocIds(selectedSourceIds)
        : sources.map((s) => s.document?.id).filter((id): id is number => id !== undefined);

    await thread.submit(
      { messages: [{ type: "human", content: text }] },
      {
        threadId,
        streamSubgraphs: true,
        config: {
          configurable: {
            notebook_search_space_id: notebookId,
            notebook_source_ids: sourceIds,
            notebook_token: token,
            ...(modelName ? { model_name: modelName } : {}),
          },
        },
      },
    );
  };

  return { thread, messages: thread.messages, isStreaming: thread.isLoading, send, threadId, setThreadId };
}

export function useNotebookSources(notebookId: number | null): {
  sources: SourceItem[];
  isLoading: boolean;
  isError: boolean;
} {
  const email = getStoredAuth()?.email ?? "";

  // 复用与 useStudioNotebooks 相同的 token query key，命中缓存不重复请求
  const tokenQuery = useQuery({
    queryKey: ["studio", "token", email],
    queryFn: () => fetchStudioToken(email),
    enabled: Boolean(email),
    staleTime: 5 * 60 * 1000,
    retry: 1,
  });

  const sourcesQuery = useQuery({
    queryKey: ["studio", "sources", notebookId],
    queryFn: () => fetchNotebookSources(notebookId!, tokenQuery.data!),
    enabled: notebookId !== null && Boolean(tokenQuery.data),
    staleTime: 60 * 1000,
    retry: 1,
  });

  return {
    sources: sourcesQuery.data ?? [],
    isLoading: tokenQuery.isLoading || sourcesQuery.isLoading,
    isError: tokenQuery.isError || sourcesQuery.isError,
  };
}