deer-flow/frontend/src/core/memory/hooks.ts

import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";

import {
  clearMemory,
  createMemoryFact,
  deleteMemoryFact,
  importMemory,
  loadMemory,
  updateMemoryFact,
} from "./api";
import type {
  MemoryFactInput,
  MemoryFactPatchInput,
  UserMemory,
} from "./types";

export function useMemory() {
  const { data, isLoading, error } = useQuery({
    queryKey: ["memory"],
    queryFn: () => loadMemory(),
  });
  return { memory: data ?? null, isLoading, error };
}

export function useClearMemory() {
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: () => clearMemory(),
    onSuccess: async (memory) => {
      await queryClient.cancelQueries({ queryKey: ["memory"], exact: true });
      queryClient.setQueryData<UserMemory>(["memory"], memory);
    },
  });
}

export function useDeleteMemoryFact() {
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: (factId: string) => deleteMemoryFact(factId),
    onSuccess: async (memory) => {
      await queryClient.cancelQueries({ queryKey: ["memory"], exact: true });
      queryClient.setQueryData<UserMemory>(["memory"], memory);
    },
  });
}

export function useImportMemory() {
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: (memory: UserMemory) => importMemory(memory),
    onSuccess: async (memory) => {
      await queryClient.cancelQueries({ queryKey: ["memory"], exact: true });
      queryClient.setQueryData<UserMemory>(["memory"], memory);
    },
  });
}

export function useCreateMemoryFact() {
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: (input: MemoryFactInput) => createMemoryFact(input),
    onSuccess: async (memory) => {
      await queryClient.cancelQueries({ queryKey: ["memory"], exact: true });
      queryClient.setQueryData<UserMemory>(["memory"], memory);
    },
  });
}

export function useUpdateMemoryFact() {
  const queryClient = useQueryClient();

  return useMutation({
    mutationFn: ({
      factId,
      input,
    }: {
      factId: string;
      input: MemoryFactPatchInput;
    }) => updateMemoryFact(factId, input),
    onSuccess: async (memory) => {
      await queryClient.cancelQueries({ queryKey: ["memory"], exact: true });
      queryClient.setQueryData<UserMemory>(["memory"], memory);
    },
  });
}