@prostgles/prostgles 2.3.6-dev.1786548258

package/client/src/dashboard/AskLLM/Chat/AskLLMChatMessages/hooks/useLLMChatMessages.tsx

@@ -0,0 +1,157 @@
+import type { DBSSchema } from "@common/publishUtils";
+import type { Message } from "@components/Chat/Chat";
+import ErrorComponent from "@components/ErrorComponent";
+import { FlexCol } from "@components/Flex";
+import { FormFieldDebounced } from "@components/FormField/FormFieldDebounced";
+import React, { useMemo } from "react";
+import { isDefined } from "../../../../../utils/utils";
+import type { UseLLMChatProps } from "../../useLLMChat";
+import { LLMChatMessage } from "../LLMChatMessage/LLMChatMessage";
+import { LLMChatMessageHeader } from "../LLMChatMessage/LLMChatMessageHeader";
+import { useLLMChatMessageGrouper } from "./useLLMChatMessageGrouper";
+
+type P = UseLLMChatProps & {
+  activeChat: DBSSchema["llm_chats"] | undefined;
+};
+
+export const useLLMChatMessages = (props: P) => {
+  const { dbs, user, activeChat, loadedSuggestions, workspaceId } = props;
+  const { status } = activeChat ?? {};
+
+  const { data: llmMessages } = dbs.llm_messages.useSubscribe(
+    { chat_id: activeChat?.id },
+    { orderBy: { created: 1 } },
+    { skip: !activeChat?.id },
+  );
+
+  const isLoadingSince = status?.state === "loading" ? status.since : null;
+  const { llmMessagesWithGroups } = useLLMChatMessageGrouper({
+    llmMessages,
+  });
+
+  const actualMessages: Message[] | undefined = useMemo(
+    () =>
+      !llmMessages ? undefined : (
+        llmMessagesWithGroups
+          ?.map((messageItem, messageOrGroupIndex) => {
+            const isLastMessage =
+              llmMessagesWithGroups.length - 1 === messageOrGroupIndex;
+            const { id, user_id } =
+              messageItem.type === "single_message" ?
+                messageItem.message
+              : messageItem.messages[0].message;
+            const isLoadingSinceDate =
+              !isLastMessage || !isLoadingSince ?
+                undefined
+              : new Date(isLoadingSince);
+            return {
+              id,
+              incoming: user_id !== user?.id,
+              messageTopContent: (
+                <LLMChatMessageHeader dbs={dbs} item={messageItem} />
+              ),
+              message: (
+                <LLMChatMessage
+                  isLoadingSinceDate={isLoadingSinceDate}
+                  messageItem={messageItem}
+                  workspaceId={workspaceId}
+                  loadedSuggestions={loadedSuggestions}
+                  hideLoadingCounter={
+                    user?.options?.hideLlmLoadingCounter ?? false
+                  }
+                />
+              ),
+              sender_id: user_id || "ai",
+            };
+          })
+          .filter(isDefined)
+      ),
+    [
+      llmMessages,
+      llmMessagesWithGroups,
+      isLoadingSince,
+      user?.id,
+      user?.options?.hideLlmLoadingCounter,
+      dbs,
+      workspaceId,
+      loadedSuggestions,
+    ],
+  );
+
+  const messages = useMemo(() => {
+    const lastMessage = llmMessages?.at(-1);
+    const disabled_message =
+      (
+        activeChat?.disabled_until &&
+        new Date(activeChat.disabled_until) > new Date() &&
+        activeChat.disabled_message
+      ) ?
+        activeChat.disabled_message
+      : (
+        // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
+        lastMessage?.meta?.finish_reason === "length" ||
+        // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
+        lastMessage?.meta?.native_finish_reason === "max_output_tokens"
+      ) ?
+        <FlexCol>
+          <ErrorComponent
+            error={
+              "finish_reason = 'length'. Increase max_tokens and try again"
+            }
+          />
+          <FormFieldDebounced
+            label={"Max tokens"}
+            value={
+              // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+              activeChat?.extra_body?.max_tokens ||
+              // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
+              lastMessage.meta?.max_tokens ||
+              6000
+            }
+            onChange={(max_tokens) => {
+              void dbs.llm_chats.update(
+                { id: activeChat!.id },
+                { extra_body: { max_tokens: Number(max_tokens) } },
+              );
+            }}
+          />
+        </FlexCol>
+      : undefined;
+    const messages: Message[] = (
+      actualMessages?.length ? actualMessages : (
+        [
+          {
+            id: "first",
+            message: "Hello, I am the AI assistant. How can I help you?",
+            incoming: true,
+            sender_id: "ai",
+          } as const,
+        ].map((m) => {
+          const incoming = m.sender_id !== user?.id;
+          return {
+            ...m,
+            incoming,
+            message: m.message,
+          };
+        })
+      )).concat(
+      disabled_message ?
+        [
+          {
+            id: "disabled-last",
+            incoming: true,
+            message: disabled_message,
+            sender_id: "ai",
+          },
+        ]
+      : [],
+    );
+
+    return messages;
+  }, [activeChat, actualMessages, dbs.llm_chats, llmMessages, user?.id]);
+
+  return {
+    llmMessages,
+    messages: activeChat && actualMessages ? messages : undefined,
+  };
+};

package/client/src/dashboard/AskLLM/Chat/AskLLMChatMessages/setChatPrompt.ts

@@ -0,0 +1,116 @@
+import { isDefined } from "@common/filterUtils";
+import type { DBSSchema, DBSSchemaForInsert } from "@common/publishUtils";
+import { getEntries } from "@common/utils";
+import type { DBS } from "src/dashboard/Dashboard/DBS";
+import { tout } from "src/utils/utils";
+
+export const setChatPrompt = async ({
+  dbs,
+  chatId,
+  prompt: { id, options },
+  currentPrompt,
+}: {
+  dbs: DBS;
+  chatId: number;
+  prompt: Pick<DBSSchema["llm_prompts"], "id" | "options">;
+  currentPrompt: DBSSchema["llm_prompts"] | undefined;
+}) => {
+  const { max_tokens, mcp_server_tools, temperature } = options || {};
+  const res = await dbs.llm_chats.update(
+    { id: chatId },
+    {
+      llm_prompt_id: id,
+      extra_body:
+        isDefined(max_tokens ?? temperature) ?
+          {
+            temperature,
+            max_tokens,
+          }
+        : undefined,
+    },
+    {
+      returning: { id: 1 },
+    },
+  );
+  if (res?.length !== 1) {
+    throw "Failed to update chat prompt";
+  }
+  const getTools = (opts: DBSSchema["llm_prompts"]["options"]) => {
+    return Promise.all(
+      getEntries(opts?.mcp_server_tools ?? {}).map(
+        async ([server_name, toolList = {}]) => {
+          return dbs.mcp_server_tools.find({
+            server_name,
+            name: { $in: Object.keys(toolList) },
+          });
+        },
+      ),
+    );
+  };
+
+  /** Disable previous tools */
+  if (currentPrompt?.options?.mcp_server_tools && currentPrompt.id !== id) {
+    const currentTools = await getTools(currentPrompt.options);
+    if (currentTools.length) {
+      await dbs.llm_chats_allowed_mcp_tools.delete({
+        $or: currentTools.flat().map(
+          (t) =>
+            ({
+              chat_id: chatId,
+              tool_id: t.id,
+              server_name: t.server_name,
+            }) satisfies DBSSchemaForInsert["llm_chats_allowed_mcp_tools"],
+        ),
+      });
+    }
+  }
+
+  /** Disable previous db access */
+  if (currentPrompt?.options?.database_access && currentPrompt.id !== id) {
+    await dbs.llm_chats.update({ id: chatId }, { db_data_permissions: null });
+  }
+
+  const mcpServersAndTools = getEntries(mcp_server_tools || {});
+  if (mcpServersAndTools.length) {
+    const allowedToolsInsert = await Promise.all(
+      mcpServersAndTools.flatMap(async ([server_name, toolListObj = {}]) => {
+        await dbs.mcp_servers.update({ name: server_name }, { enabled: true });
+        /** Wait for tools to load if missing */
+        while (!(await dbs.mcp_server_tools.count({ server_name }))) {
+          await tout(500);
+        }
+
+        const toolList = Object.keys(toolListObj);
+        const tools = await dbs.mcp_server_tools.find({
+          server_name,
+          name: { $in: toolList },
+        });
+        if (!tools.length || tools.length !== toolList.length) {
+          throw new Error(
+            `Some tools not found for MCP server ${server_name}. Expected: ${toolList}, found: ${tools.map((t) => t.name).join(", ")}`,
+          );
+        }
+        return tools.flatMap(({ id, name }) => {
+          return {
+            chat_id: chatId,
+            server_name,
+            tool_id: id,
+            auto_approve: toolListObj[name] === "auto-approve",
+          } satisfies DBSSchemaForInsert["llm_chats_allowed_mcp_tools"];
+        });
+      }),
+    ).then((arr) => arr.flat());
+    await dbs.llm_chats_allowed_mcp_tools.insertMany(allowedToolsInsert, {
+      onConflict: "DoNothing",
+    });
+  }
+
+  if (options?.database_access) {
+    await dbs.llm_chats.update(
+      { id: chatId },
+      {
+        db_data_permissions: { mode: options.database_access },
+      },
+    );
+  }
+};

package/client/src/dashboard/AskLLM/Chat/AskLLMChatSettings.tsx

@@ -0,0 +1,88 @@
+import type { DBSSchema } from "@common/publishUtils";
+import Btn from "@components/Btn";
+import { FlexCol } from "@components/Flex";
+import Popup from "@components/Popup/Popup";
+import { mdiCogOutline } from "@mdi/js";
+import { usePrgl } from "@pages/ProjectConnection/PrglContextProvider";
+import React, { useMemo, useState } from "react";
+import { t } from "../../../i18n/i18nUtils";
+import { SmartForm, type SmartFormProps } from "../../SmartForm/SmartForm";
+
+export type LLMChatOptionsProps = {
+  prompts: DBSSchema["llm_prompts"][] | undefined;
+  activeChat: DBSSchema["llm_chats"] | undefined;
+  credentials: DBSSchema["llm_credentials"][] | undefined;
+  activeChatId: number | undefined;
+  chatRootDiv: HTMLDivElement;
+};
+
+export const AskLLMChatSettings = (props: LLMChatOptionsProps) => {
+  const { chatRootDiv, activeChatId } = props;
+  const { dbs, dbsTables, dbsMethodSchema, dbsSql } = usePrgl();
+  const [anchorEl, setAnchorEl] = useState<HTMLDivElement>();
+
+  const formProps = useMemo(() => {
+    return {
+      columns: {
+        name: 1,
+        llm_prompt_id: 1,
+        model: 1,
+        db_schema_permissions: 1,
+        db_data_permissions: 1,
+        maximum_consecutive_tool_fails: 1,
+        max_total_cost_usd: 1,
+        extra_body: 1,
+        extra_headers: 1,
+        created: 1,
+        id: 1,
+        options: 1,
+      } as const satisfies Partial<Record<keyof DBSSchema["llm_chats"], 1>>,
+      methods: {},
+      rowFilter: [{ fieldName: "id", value: activeChatId }],
+      jsonbSchemaWithControls: {
+        // noLabels: true,
+      },
+    } satisfies Pick<
+      SmartFormProps,
+      "jsonbSchemaWithControls" | "columns" | "methods" | "rowFilter"
+    >;
+  }, [activeChatId]);
+
+  return (
+    <>
+      <Btn
+        title={t.AskLLMChatHeader["Chat settings"]}
+        variant="faded"
+        iconPath={mdiCogOutline}
+        onClick={() => setAnchorEl(anchorEl ? undefined : chatRootDiv)}
+        data-command="LLMChatOptions.toggle"
+        size="default"
+      />
+      {anchorEl && (
+        <Popup
+          contentStyle={{ padding: 0, maxWidth: "min(100vw, 700px)" }}
+          onClickClose={false}
+          onClose={() => setAnchorEl(undefined)}
+          anchorEl={chatRootDiv}
+          positioning="right-panel"
+          clickCatchStyle={{ opacity: 1 }}
+          title={t.AskLLMChatHeader["Chat settings"]}
+          content={
+            <FlexCol className="f-1 min-s-0">
+              <SmartForm
+                label=""
+                tableName="llm_chats"
+                contentClassname="p-1 pt-1"
+                {...formProps}
+                db={dbs}
+                tables={dbsTables}
+                methods={dbsMethodSchema}
+                sql={dbsSql}
+              />
+            </FlexCol>
+          }
+        />
+      )}
+    </>
+  );
+};

package/client/src/dashboard/AskLLM/Chat/useAskLLMChatSend.ts

@@ -0,0 +1,111 @@
+import type { LLMMessage } from "@common/llmUtils";
+import { MINUTE } from "@common/utils";
+import { useAlert } from "@components/AlertProvider";
+import { type ChatProps } from "@components/Chat/Chat";
+import { useCallback, useMemo } from "react";
+import { isDefined } from "../../../utils/utils";
+import type { AskLLMChatProps } from "./AskLLMChat";
+import type { LLMChatState } from "./useLLMChat";
+import { usePrgl } from "@pages/ProjectConnection/PrglContextProvider";
+import type { LLMMessageContent } from "./AskLLMChatMessages/ToolUseChatMessage/ToolUseChatMessage";
+
+type P = Pick<AskLLMChatProps, "askLLM" | "stopAskLLM"> &
+  Pick<LLMChatState, "activeChatId" | "activeChat"> & {
+    dbSchemaForPrompt: string;
+  };
+
+export const useAskLLMChatSend = ({
+  askLLM,
+  stopAskLLM,
+  activeChat,
+  activeChatId,
+  dbSchemaForPrompt,
+}: P) => {
+  const { connectionId } = usePrgl();
+  const { addAlert } = useAlert();
+  const sendQuery = useCallback(
+    (msg: LLMMessage["message"] | undefined) => {
+      if (!msg || !activeChatId) return;
+      /** TODO: move dbSchemaForPrompt to server-side */
+      void askLLM({
+        connectionId,
+        userMessage: msg,
+        schema: dbSchemaForPrompt,
+        chatId: activeChatId,
+        type: "new-message",
+      }).catch((error) => {
+        const errorText = error?.message || error;
+        const errorTextMessage =
+          typeof errorText === "string" ? errorText : JSON.stringify(errorText);
+
+        addAlert(
+          "Error when when sending AI Assistant query: " + errorTextMessage,
+        );
+      });
+    },
+    [activeChatId, askLLM, connectionId, dbSchemaForPrompt, addAlert],
+  );
+
+  const sendMessage: ChatProps["onSend"] = useCallback(
+    async (userMessage: LLMMessage["message"] | undefined, files) => {
+      const fileMessages = await Promise.all(
+        (files ?? []).map(async (file) => toMediaMessage(file)),
+      );
+      return sendQuery(
+        [...(userMessage ?? []), ...fileMessages].filter(isDefined),
+      );
+    },
+    [sendQuery],
+  );
+
+  const status = activeChat?.status;
+  const isLoading = status?.state === "loading";
+  const chatIsLoading =
+    isLoading && new Date(status.since) > new Date(Date.now() - 1 * MINUTE);
+
+  const onStopSending = useMemo(() => {
+    if (!isLoading || activeChatId === undefined) {
+      return;
+    }
+    return () => stopAskLLM({ chatId: activeChatId });
+  }, [activeChatId, isLoading, stopAskLLM]);
+
+  return {
+    sendMessage,
+    chatIsLoading,
+    onStopSending,
+    sendQuery,
+  };
+};
+
+const toBase64 = (file: File) =>
+  new Promise<string>((resolve, reject) => {
+    const reader = new FileReader();
+    reader.readAsDataURL(file);
+    reader.onload = () => resolve(reader.result as string);
+    reader.onerror = reject;
+  });
+
+const toMediaMessage = async (
+  file: File,
+): Promise<
+  Extract<
+    LLMMessageContent,
+    {
+      source: {
+        type: "base64";
+      };
+    }
+  >
+> => {
+  const base64 = await toBase64(file);
+  const type = file.type.split("/")[0] as "image";
+  return {
+    type,
+    source: {
+      type: "base64",
+      media_type: file.type,
+      data: base64,
+    },
+  };
+};

package/client/src/dashboard/AskLLM/Chat/useLLMChat.tsx

@@ -0,0 +1,126 @@
+import type { DBSSchema } from "@common/publishUtils";
+import { useEffectDeep } from "prostgles-client/dist/prostgles";
+import type { FilterItem } from "prostgles-types";
+import { useCallback, useMemo, useState } from "react";
+import type { Prgl } from "../../../App";
+import type { LoadedSuggestions } from "../../Dashboard/dashboardUtils";
+import type { LLMSetupStateReady } from "../Setup/LLMSetupProvider";
+import { useLLMChatMessages } from "./AskLLMChatMessages/hooks/useLLMChatMessages";
+import { setChatPrompt } from "./AskLLMChatMessages/setChatPrompt";
+import type { AskLLMChatProps } from "./AskLLMChat";
+
+export type UseLLMChatProps = LLMSetupStateReady &
+  Pick<Prgl, "dbs" | "user" | "connectionId"> & {
+    workspaceId: string | undefined;
+    loadedSuggestions: LoadedSuggestions | undefined;
+  } & Pick<AskLLMChatProps, "selectedChat">;
+
+export type LLMChatState = ReturnType<typeof useLLMChat>;
+export const useLLMChat = (props: UseLLMChatProps) => {
+  const {
+    dbs,
+    credentials,
+    firstPromptId,
+    defaultCredential,
+    prompts,
+    selectedChat,
+  } = props;
+  const chatsFilter = useMemo(() => {
+    return {
+      connection_id: { $in: [props.connectionId, null] },
+      parent_chat_message_id:
+        selectedChat?.type === "agent" ? selectedChat.parent_message_id : null,
+    } satisfies FilterItem<DBSSchema["llm_chats"]>;
+  }, [selectedChat, props.connectionId]);
+  const [selectedChatId, setSelectedChat] = useState(selectedChat?.id);
+  const { data: latestChats } = dbs.llm_chats.useSubscribe(chatsFilter, {
+    select: { "*": 1, created_ago: { $ageNow: ["created"] } },
+    orderBy: { created: -1 },
+  });
+  const latestChat = latestChats?.[0];
+  /**
+   * Always show the selected chat if it exists otherwise show latest
+   * If no chats exist, new chat will be created
+   */
+  const activeChat =
+    latestChats?.find((c) => c.id === selectedChatId) ?? latestChat;
+  const activeChatId = activeChat?.id;
+
+  const preferredPromptId = activeChat?.llm_prompt_id ?? firstPromptId;
+  const lastModelId = activeChat?.model;
+  const createNewChat = useCallback(
+    async (promptId: number, ifNoOtherChatsExist = false) => {
+      if (ifNoOtherChatsExist) {
+        const chat = await dbs.llm_chats.findOne(chatsFilter);
+        if (chat) {
+          console.warn("Chat already exists", chat);
+          return;
+        }
+      }
+      const prompt = prompts.find((p) => p.id === promptId);
+      if (!preferredPromptId || !prompt) {
+        console.warn("No prompt found", { prompts });
+        return;
+      }
+      const newChat = await dbs.llm_chats.insert(
+        {
+          name: "New chat",
+          // TODO: add publish rules (forcedData) to DBHandlerClient typings
+          user_id: undefined as unknown as string,
+          connection_id: props.connectionId,
+          llm_prompt_id: promptId,
+          model: lastModelId,
+          // options: {
+          //   /** Anthropic consumes tokens for tools json schema as opposed to openai codex? */
+          //   mcpToolSchemaMode: "hide-schemas-and-descriptions",
+          // },
+        },
+        { returning: "*" },
+      );
+      await setChatPrompt({
+        dbs,
+        chatId: newChat.id,
+        prompt,
+        currentPrompt: undefined,
+      });
+      setSelectedChat(undefined);
+    },
+    [
+      chatsFilter,
+      dbs,
+      lastModelId,
+      preferredPromptId,
+      prompts,
+      props.connectionId,
+    ],
+  );
+
+  useEffectDeep(() => {
+    if (latestChats && !latestChats.length && preferredPromptId) {
+      void createNewChat(preferredPromptId, true);
+    }
+  }, [latestChats, preferredPromptId, defaultCredential, createNewChat]);
+
+  const { llmMessages, messages } = useLLMChatMessages({
+    ...props,
+    activeChat,
+  });
+
+  const prompt = useMemo(() => {
+    return prompts.find((p) => p.id === activeChat?.llm_prompt_id);
+  }, [activeChat?.llm_prompt_id, prompts]);
+
+  return {
+    activeChatId,
+    createNewChat,
+    preferredPromptId,
+    llmMessages,
+    messages,
+    latestChats,
+    setActiveChat: setSelectedChat,
+    credentials,
+    defaultCredential,
+    activeChat,
+    prompt,
+  };
+};