gitdataai/src/hooks/useAiChatQuery.ts

import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query";
import type { CreateConversationParams, CreateMessageParams, MessageResponse } from "@/client/model";
import {
  listConversations,
  createConversation as apiCreateConversation,
  getConversation as apiGetConversation,
  deleteConversation as apiDeleteConversation,
  updateConversation as apiUpdateConversation,
  listMessages,
  createMessage as apiCreateMessage,
  stopMessage as apiStopMessage,
  streamChat,
  editMessage as apiEditMessage,
  listMessageVersions as apiListMessageVersions,
  switchMessageVersion as apiSwitchMessageVersion,
  listMessageForks as apiListMessageForks,
} from "@/client/aiChatApi";

const CONVERSATIONS_KEY = "ai-conversations";
const MESSAGES_KEY = "ai-messages";
const VERSIONS_KEY = "ai-message-versions";

export function useConversationsQuery(projectId?: string | null) {
  return useQuery({
    queryKey: [CONVERSATIONS_KEY, projectId],
    queryFn: () => listConversations({ project_id: projectId }),
    staleTime: 5 * 60 * 1000,
  });
}

export function useConversationQuery(conversationId: string) {
  return useQuery({
    queryKey: [CONVERSATIONS_KEY, conversationId],
    queryFn: () => apiGetConversation(conversationId),
    enabled: !!conversationId,
    staleTime: 5 * 60 * 1000,
  });
}

export function useCreateConversationMutation() {
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: (params: CreateConversationParams) => apiCreateConversation(params),
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: [CONVERSATIONS_KEY] });
    },
  });
}

export function useUpdateConversationMutation() {
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: ({ id, params }: { id: string; params: Parameters<typeof apiUpdateConversation>[1] }) =>
      apiUpdateConversation(id, params),
    onSuccess: (_, vars) => {
      queryClient.invalidateQueries({ queryKey: [CONVERSATIONS_KEY, vars.id] });
      queryClient.invalidateQueries({ queryKey: [CONVERSATIONS_KEY] });
    },
  });
}

export function useDeleteConversationMutation() {
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: (id: string) => apiDeleteConversation(id),
    onSuccess: () => {
      queryClient.invalidateQueries({ queryKey: [CONVERSATIONS_KEY] });
    },
  });
}

export function useMessagesQuery(conversationId: string) {
  return useQuery({
    queryKey: [MESSAGES_KEY, conversationId],
    queryFn: () => listMessages(conversationId),
    enabled: !!conversationId,
    staleTime: 5 * 60 * 1000,
  });
}

export function useCreateMessageMutation() {
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: ({ conversationId, params }: { conversationId: string; params: CreateMessageParams }) =>
      apiCreateMessage(conversationId, params),
    onSuccess: (_, vars) => {
      queryClient.invalidateQueries({ queryKey: [MESSAGES_KEY, vars.conversationId] });
      queryClient.invalidateQueries({ queryKey: [CONVERSATIONS_KEY, vars.conversationId] });
    },
  });
}

export function useStopMessageMutation() {
  return useMutation({
    mutationFn: ({ conversationId, messageId }: { conversationId: string; messageId: string }) =>
      apiStopMessage(conversationId, messageId),
  });
}

export { streamChat };
export type { MessageResponse };

export function useEditMessageMutation() {
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: ({ conversationId, messageId, content }: { conversationId: string; messageId: string; content: string }) =>
      apiEditMessage(conversationId, messageId, content),
    onSuccess: (_, vars) => {
      queryClient.invalidateQueries({ queryKey: [MESSAGES_KEY, vars.conversationId] });
      queryClient.invalidateQueries({ queryKey: [VERSIONS_KEY, vars.conversationId, vars.messageId] });
    },
  });
}

export function useMessageVersionsQuery(conversationId: string, messageId: string) {
  return useQuery({
    queryKey: [VERSIONS_KEY, conversationId, messageId],
    queryFn: () => apiListMessageVersions(conversationId, messageId),
    enabled: !!conversationId && !!messageId,
    staleTime: 30 * 1000,
  });
}

export function useSwitchMessageVersionMutation() {
  const queryClient = useQueryClient();
  return useMutation({
    mutationFn: ({ conversationId, messageId, versionNumber }: { conversationId: string; messageId: string; versionNumber: number }) =>
      apiSwitchMessageVersion(conversationId, messageId, versionNumber),
    onSuccess: (_, vars) => {
      queryClient.invalidateQueries({ queryKey: [MESSAGES_KEY, vars.conversationId] });
      queryClient.invalidateQueries({ queryKey: [VERSIONS_KEY, vars.conversationId, vars.messageId] });
    },
  });
}

export function useMessageForksQuery(conversationId: string, messageId: string) {
  return useQuery({
    queryKey: ["ai-message-forks", conversationId, messageId],
    queryFn: () => apiListMessageForks(conversationId, messageId),
    enabled: !!conversationId && !!messageId,
    staleTime: 30 * 1000,
  });
}