context-efficient-ai/pages/chat/@id/+Page.tsx

import {
  ActionIcon,
  Box,
  Group,
  HoverCard,
  JsonInput,
  List,
  Stack,
  Tabs,
  Textarea,
  useMantineTheme,
} from "@mantine/core";
import { useEffect, useState } from "react";
import {
  defaultParameters,
  defaultSystemPrompt,
  useStore,
} from "../../../state";
import { usePageContext } from "vike-react/usePageContext";
import { useData } from "vike-react/useData";
import type { Data } from "./+data";
import type { CommittedMessage, DraftMessage } from "../../../types";
import Markdown from "react-markdown";
import { IconTrash, IconEdit, IconCheck, IconX } from "@tabler/icons-react";
import { useTRPC } from "../../../trpc/client";
import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
import { nanoid } from "nanoid";
import type { Conversation } from "../../../database/common";

export default function ChatPage() {
  const pageContext = usePageContext();
  const conversationId = pageContext.routeParams.id;

  const {
    conversation,
    // messages: initialMessages,
    // facts: initialFacts,
    // factTriggers: initialFactTriggers,
  } = useData<Data>();

  const conversationTitle = conversation?.title;

  const message = useStore((state) => state.message);
  const systemPrompt = useStore((state) => state.systemPrompt);
  const parameters = useStore((state) => state.parameters);
  const loading = useStore((state) => state.loading);
  const setMessage = useStore((state) => state.setMessage);
  const setSystemPrompt = useStore((state) => state.setSystemPrompt);
  const setParameters = useStore((state) => state.setParameters);
  const setLoading = useStore((state) => state.setLoading);
  const trpc = useTRPC();
  const queryClient = useQueryClient();

  const messagesResult = useQuery(
    trpc.chat.messages.fetchByConversationId.queryOptions({
      conversationId,
    })
  );
  const messages: Array<CommittedMessage> | undefined =
    messagesResult.data?.map((m) => ({
      ...m,
      parts: m.parts.filter((p) => p.type === "text"),
    })) || [];

  const facts = useQuery(
    trpc.chat.facts.fetchByConversationId.queryOptions({
      conversationId,
    })
  );
  const factTriggers = useQuery(
    trpc.chat.factTriggers.fetchByConversationId.queryOptions({
      conversationId,
    })
  );

  const deleteFact = useMutation(
    trpc.chat.facts.deleteOne.mutationOptions({
      onMutate: async ({ factId: factIdToDelete }) => {
        /** Cancel affected queries that may be in-flight: */
        await queryClient.cancelQueries({
          queryKey: trpc.chat.facts.fetchByConversationId.queryKey({
            conversationId,
          }),
        });

        /** Optimistically update the affected queries in react-query's cache: */
        const previousFacts = await queryClient.getQueryData(
          trpc.chat.facts.fetchByConversationId.queryKey({
            conversationId,
          })
        );
        if (!previousFacts) {
          return {
            previousFacts: [],
            newFacts: [],
          };
        }
        const newFacts = previousFacts.filter((f) => f.id !== factIdToDelete);
        queryClient.setQueryData(
          trpc.chat.facts.fetchByConversationId.queryKey({
            conversationId,
          }),
          newFacts
        );

        return { previousFacts, newFacts };
      },
      onSettled: async (data, variables, context) => {
        await queryClient.invalidateQueries({
          queryKey: trpc.chat.facts.fetchByConversationId.queryKey({
            conversationId,
          }),
        });
      },
      onError: async (error, variables, context) => {
        console.error(error);
        if (!context) return;
        queryClient.setQueryData(
          trpc.chat.facts.fetchByConversationId.queryKey({
            conversationId,
          }),
          context.previousFacts
        );
      },
    })
  );
  const updateFact = useMutation(
    trpc.chat.facts.update.mutationOptions({
      onMutate: async ({ factId, content }) => {
        /** Cancel affected queries that may be in-flight: */
        await queryClient.cancelQueries({
          queryKey: trpc.chat.facts.fetchByConversationId.queryKey({
            conversationId,
          }),
        });

        /** Optimistically update the affected queries in react-query's cache: */
        const previousFacts = await queryClient.getQueryData(
          trpc.chat.facts.fetchByConversationId.queryKey({
            conversationId,
          })
        );
        if (!previousFacts) {
          return {
            previousFacts: [],
            newFacts: [],
          };
        }
        const newFacts = previousFacts.map((f) =>
          f.id === factId ? { ...f, content } : f
        );
        queryClient.setQueryData(
          trpc.chat.facts.fetchByConversationId.queryKey({
            conversationId,
          }),
          newFacts
        );

        return { previousFacts, newFacts };
      },
      onSettled: async (data, variables, context) => {
        await queryClient.invalidateQueries({
          queryKey: trpc.chat.facts.fetchByConversationId.queryKey({
            conversationId,
          }),
        });
      },
      onError: async (error, variables, context) => {
        console.error(error);
        if (!context) return;
        queryClient.setQueryData(
          trpc.chat.facts.fetchByConversationId.queryKey({
            conversationId,
          }),
          context.previousFacts
        );
      },
    })
  );

  const deleteFactTrigger = useMutation(
    trpc.chat.factTriggers.deleteOne.mutationOptions({
      onMutate: async ({ factTriggerId: factTriggerIdToDelete }) => {
        /** Cancel affected queries that may be in-flight: */
        await queryClient.cancelQueries({
          queryKey: trpc.chat.factTriggers.fetchByConversationId.queryKey({
            conversationId,
          }),
        });

        /** Optimistically update the affected queries in react-query's cache: */
        const previousFactTriggers = await queryClient.getQueryData(
          trpc.chat.factTriggers.fetchByConversationId.queryKey({
            conversationId,
          })
        );
        if (!previousFactTriggers) {
          return {
            previousFactTriggers: [],
            newFactTriggers: [],
          };
        }
        const newFactTriggers = previousFactTriggers.filter(
          (ft) => ft.id !== factTriggerIdToDelete
        );
        queryClient.setQueryData(
          trpc.chat.factTriggers.fetchByConversationId.queryKey({
            conversationId,
          }),
          newFactTriggers
        );

        return { previousFactTriggers, newFactTriggers };
      },
      onSettled: async (data, variables, context) => {
        await queryClient.invalidateQueries({
          queryKey: trpc.chat.factTriggers.fetchByConversationId.queryKey({
            conversationId,
          }),
        });
      },
      onError: async (error, variables, context) => {
        console.error(error);
        if (!context) return;
        queryClient.setQueryData(
          trpc.chat.factTriggers.fetchByConversationId.queryKey({
            conversationId,
          }),
          context.previousFactTriggers
        );
      },
    })
  );
  const updateFactTrigger = useMutation(
    trpc.chat.factTriggers.update.mutationOptions({
      onMutate: async ({ factTriggerId, content }) => {
        /** Cancel affected queries that may be in-flight: */
        await queryClient.cancelQueries({
          queryKey: trpc.chat.factTriggers.fetchByConversationId.queryKey({
            conversationId,
          }),
        });
        /** Optimistically update the affected queries in react-query's cache: */
        const previousFactTriggers = await queryClient.getQueryData(
          trpc.chat.factTriggers.fetchByConversationId.queryKey({
            conversationId,
          })
        );
        if (!previousFactTriggers) {
          return {
            previousFactTriggers: [],
            newFactTriggers: [],
          };
        }
        const newFactTriggers = previousFactTriggers.map((ft) =>
          ft.id === factTriggerId ? { ...ft, content } : ft
        );
        queryClient.setQueryData(
          trpc.chat.factTriggers.fetchByConversationId.queryKey({
            conversationId,
          }),
          newFactTriggers
        );
        return { previousFactTriggers, newFactTriggers };
      },
      onSettled: async (data, variables, context) => {
        await queryClient.invalidateQueries({
          queryKey: trpc.chat.factTriggers.fetchByConversationId.queryKey({
            conversationId,
          }),
        });
      },
      onError: async (error, variables, context) => {
        console.error(error);
        if (!context) return;
        queryClient.setQueryData(
          trpc.chat.factTriggers.fetchByConversationId.queryKey({
            conversationId,
          }),
          context.previousFactTriggers
        );
      },
    })
  );

  const updateConversationTitle = useMutation(
    trpc.chat.conversations.updateTitle.mutationOptions({
      onMutate: async ({ id, title }) => {
        /** Cancel affected queries that may be in-flight: */
        await queryClient.cancelQueries({
          queryKey: trpc.chat.conversations.fetchAll.queryKey(),
        });
        /** Optimistically update the affected queries in react-query's cache: */
        const previousConversations = await queryClient.getQueryData(
          trpc.chat.conversations.fetchAll.queryKey()
        );
        if (!previousConversations) {
          return {
            previousConversations: [],
            newConversations: null,
          };
        }
        const newConversations: Array<Conversation> = [
          ...previousConversations,
          {
            ...conversation,
            title,
          } as Conversation,
        ];
        queryClient.setQueryData(
          trpc.chat.conversations.fetchAll.queryKey(),
          newConversations
        );
        return { previousConversations, newConversations };
      },
      onSettled: async (data, variables, context) => {
        await queryClient.invalidateQueries({
          queryKey: trpc.chat.conversations.fetchOne.queryKey({
            id: conversationId,
          }),
        });
        await queryClient.invalidateQueries({
          queryKey: trpc.chat.conversations.fetchAll.queryKey(),
        });
      },
      onError: async (error, variables, context) => {
        console.error(error);
        if (!context) return;
        queryClient.setQueryData(
          trpc.chat.conversations.fetchAll.queryKey(),
          context.previousConversations
        );
      },
    })
  );

  const sendMessage = useMutation(
    trpc.chat.sendMessage.mutationOptions({
      onMutate: async ({
        conversationId,
        messages,
        systemPrompt,
        parameters,
      }) => {
        /** Cancel affected queries that may be in-flight: */
        await queryClient.cancelQueries({
          queryKey: trpc.chat.messages.fetchByConversationId.queryKey({
            conversationId,
          }),
        });
        /** Optimistically update the affected queries in react-query's cache: */
        const previousMessages: Array<CommittedMessage> | undefined =
          await queryClient.getQueryData(
            trpc.chat.messages.fetchByConversationId.queryKey({
              conversationId,
            })
          );
        if (!previousMessages) {
          return {
            previousMessages: [],
            newMessages: [],
          };
        }
        const newMessages: Array<CommittedMessage> = [
          ...previousMessages,
          {
            /** placeholder id; will be overwritten when we get the true id from the backend */
            id: nanoid(),
            conversationId,
            // content: messages[messages.length - 1].content,
            // role: "user" as const,
            ...messages[messages.length - 1],
            index: previousMessages.length,
            createdAt: new Date().toISOString(),
          } as CommittedMessage,
        ];
        queryClient.setQueryData(
          trpc.chat.messages.fetchByConversationId.queryKey({
            conversationId,
          }),
          newMessages
        );
        return { previousMessages, newMessages };
      },
      onSettled: async (data, variables, context) => {
        await queryClient.invalidateQueries({
          queryKey: trpc.chat.messages.fetchByConversationId.queryKey({
            conversationId,
          }),
        });
        await queryClient.invalidateQueries({
          queryKey: trpc.chat.facts.fetchByConversationId.queryKey({
            conversationId,
          }),
        });
        await queryClient.invalidateQueries({
          queryKey: trpc.chat.factTriggers.fetchByConversationId.queryKey({
            conversationId,
          }),
        });
      },
      onError: async (error, variables, context) => {
        console.error(error);
        if (!context) return;
        queryClient.setQueryData(
          trpc.chat.messages.fetchByConversationId.queryKey({
            conversationId,
          }),
          context.previousMessages
        );
      },
    })
  );

  // State for editing facts
  const [editingFactId, setEditingFactId] = useState<string | null>(null);
  const [editingFactContent, setEditingFactContent] = useState("");

  // State for editing fact triggers
  const [editingFactTriggerId, setEditingFactTriggerId] = useState<
    string | null
  >(null);
  const [editingFactTriggerContent, setEditingFactTriggerContent] =
    useState("");

  // Handle clicking outside to cancel editing
  useEffect(() => {
    function handleClickOutside(event: MouseEvent) {
      if (editingFactId && event.target instanceof Element) {
        const editingElement = event.target.closest(".editing-fact");
        if (!editingElement) {
          setEditingFactId(null);
        }
      }

      if (editingFactTriggerId && event.target instanceof Element) {
        const editingElement = event.target.closest(".editing-fact-trigger");
        if (!editingElement) {
          setEditingFactTriggerId(null);
        }
      }
    }

    document.addEventListener("mousedown", handleClickOutside);
    return () => {
      document.removeEventListener("mousedown", handleClickOutside);
    };
  }, [editingFactId, editingFactTriggerId]);

  async function handleDeleteFact(factId: string) {
    await deleteFact.mutateAsync({ factId });
  }

  async function handleUpdateFact(factId: string, content: string) {
    await updateFact.mutateAsync({ factId, content });
  }

  async function handleDeleteFactTrigger(factTriggerId: string) {
    await deleteFactTrigger.mutateAsync({ factTriggerId });
  }

  async function handleUpdateFactTrigger(
    factTriggerId: string,
    content: string
  ) {
    await updateFactTrigger.mutateAsync({ factTriggerId, content });
  }

  return (
    <>
      <div>
        <span>Conversation #{conversationId} - </span>
        <input
          type="text"
          defaultValue={conversationTitle || ""}
          // onChange={(e) => {
          //   setConversationTitle(e.target.value);
          // }}
          onBlur={(e) => {
            updateConversationTitle.mutateAsync({
              id: conversationId,
              title: e.target.value,
            });
          }}
        />
      </div>
      <Tabs defaultValue="message">
        <Tabs.List>
          <Tabs.Tab value="message">Message</Tabs.Tab>
          <Tabs.Tab value="system-prompt">System Prompt</Tabs.Tab>
          <Tabs.Tab value="parameters">Parameters</Tabs.Tab>
          <Tabs.Tab value="facts">Facts</Tabs.Tab>
          <Tabs.Tab value="fact-triggers">Fact Triggers</Tabs.Tab>
        </Tabs.List>
        <Tabs.Panel value="message">
          <Messages messages={messages || []} />
          <Textarea
            resize="vertical"
            placeholder="Type your message here..."
            value={message}
            disabled={loading}
            onChange={(e) => setMessage(e.target.value)}
            onKeyDown={async (e) => {
              if (e.key === "Enter") {
                e.preventDefault();
                setLoading(true);
                await sendMessage.mutateAsync({
                  conversationId,
                  messages: [
                    ...(messages || []),
                    {
                      role: "user" as const,
                      parts: [{ type: "text", text: message }],
                    } as DraftMessage,
                  ],
                  systemPrompt,
                  parameters,
                });
                setMessage("");
                setLoading(false);
              }
            }}
          />
        </Tabs.Panel>
        <Tabs.Panel value="system-prompt">
          <Textarea
            resize="vertical"
            placeholder={defaultSystemPrompt}
            value={systemPrompt}
            onChange={(e) => setSystemPrompt(e.target.value)}
          />
        </Tabs.Panel>
        <Tabs.Panel value="parameters">
          <JsonInput
            resize="vertical"
            formatOnBlur
            placeholder={JSON.stringify(defaultParameters)}
            value={JSON.stringify(parameters)}
            onChange={(value) => setParameters(JSON.parse(value))}
          />
        </Tabs.Panel>
        <Tabs.Panel value="facts">
          <List>
            {facts.data?.map((fact) => (
              <List.Item key={fact.id}>
                {editingFactId === fact.id ? (
                  <Group wrap="nowrap" className="editing-fact">
                    <Textarea
                      value={editingFactContent}
                      onChange={(e) => setEditingFactContent(e.target.value)}
                      onKeyDown={(e) => {
                        if (e.key === "Enter") {
                          e.preventDefault();
                          handleUpdateFact(fact.id, editingFactContent);
                          setEditingFactId(null);
                        } else if (e.key === "Escape") {
                          setEditingFactId(null);
                        }
                      }}
                      autoFocus
                      style={{ flex: 1 }}
                    />
                    <ActionIcon
                      onClick={() => {
                        handleUpdateFact(fact.id, editingFactContent);
                        setEditingFactId(null);
                      }}
                    >
                      <IconCheck size={16} />
                    </ActionIcon>
                    <ActionIcon onClick={() => setEditingFactId(null)}>
                      <IconX size={16} />
                    </ActionIcon>
                  </Group>
                ) : (
                  <Group wrap="nowrap">
                    <span style={{ flex: 1 }}>{fact.content}</span>
                    <ActionIcon
                      onClick={() => {
                        setEditingFactId(fact.id);
                        setEditingFactContent(fact.content);
                      }}
                    >
                      <IconEdit size={16} />
                    </ActionIcon>
                    <IconTrash
                      size={16}
                      stroke={1.5}
                      onClick={(e) => {
                        e.stopPropagation();
                        e.preventDefault();
                        handleDeleteFact(fact.id);
                      }}
                    />
                  </Group>
                )}
              </List.Item>
            ))}
          </List>
        </Tabs.Panel>
        <Tabs.Panel value="fact-triggers">
          <List>
            {factTriggers.data?.map((factTrigger) => (
              <List.Item key={factTrigger.id}>
                {editingFactTriggerId === factTrigger.id ? (
                  <Group wrap="nowrap" className="editing-fact-trigger">
                    <Textarea
                      value={editingFactTriggerContent}
                      onChange={(e) =>
                        setEditingFactTriggerContent(e.target.value)
                      }
                      onKeyDown={(e) => {
                        if (e.key === "Enter") {
                          e.preventDefault();
                          handleUpdateFactTrigger(
                            factTrigger.id,
                            editingFactTriggerContent
                          );
                          setEditingFactTriggerId(null);
                        } else if (e.key === "Escape") {
                          setEditingFactTriggerId(null);
                        }
                      }}
                      autoFocus
                      style={{ flex: 1 }}
                    />
                    <ActionIcon
                      onClick={() => {
                        handleUpdateFactTrigger(
                          factTrigger.id,
                          editingFactTriggerContent
                        );
                        setEditingFactTriggerId(null);
                      }}
                    >
                      <IconCheck size={16} />
                    </ActionIcon>
                    <ActionIcon onClick={() => setEditingFactTriggerId(null)}>
                      <IconX size={16} />
                    </ActionIcon>
                  </Group>
                ) : (
                  <Group wrap="nowrap">
                    <span style={{ flex: 1 }}>{factTrigger.content}</span>
                    <ActionIcon
                      onClick={() => {
                        setEditingFactTriggerId(factTrigger.id);
                        setEditingFactTriggerContent(factTrigger.content);
                      }}
                    >
                      <IconEdit size={16} />
                    </ActionIcon>
                    <IconTrash
                      size={16}
                      stroke={1.5}
                      onClick={(e) => {
                        e.stopPropagation();
                        e.preventDefault();
                        handleDeleteFactTrigger(factTrigger.id);
                      }}
                    />
                  </Group>
                )}
              </List.Item>
            ))}
          </List>
        </Tabs.Panel>
      </Tabs>
    </>
  );
}

function Messages({
  messages,
}: {
  messages: Array<DraftMessage | CommittedMessage>;
}) {
  const theme = useMantineTheme();
  return (
    <Stack gap="md" justify="flex-start">
      {messages.map((message, index) => (
        <Group
          key={"id" in message ? message.id : index}
          justify={message.role === "user" ? "flex-end" : "flex-start"}
        >
          <HoverCard
            shadow="md"
            position={message.role === "user" ? "left" : "right"}
          >
            <HoverCard.Target>
              <Box
                w="75%"
                bg={
                  message.role === "user"
                    ? theme.colors.gray[2]
                    : theme.colors.blue[2]
                }
                p="md"
                bdrs="md"
              >
                <Markdown>
                  {message.parts
                    .filter((p) => p.type === "text")
                    .map((p) => p.text)
                    .join("\n")}
                </Markdown>
              </Box>
            </HoverCard.Target>
            <HoverCard.Dropdown>
              <ActionIcon.Group>
                <ActionIcon size="lg" variant="filled" color="red">
                  <IconTrash
                    size={16}
                    stroke={1.5}
                    onClick={(e) => {
                      e.stopPropagation();
                      e.preventDefault();
                      // handleDeleteMessage(message.id);
                    }}
                  />
                </ActionIcon>
              </ActionIcon.Group>
            </HoverCard.Dropdown>
          </HoverCard>
          {"runningSummary" in message && message.runningSummary && (
            <Box w="75%" bd="dotted" p="md" bdrs="md">
              <div>
                <strong>Running Summary:</strong>
                <Markdown>{message.runningSummary}</Markdown>
              </div>
            </Box>
          )}
        </Group>
      ))}
    </Stack>
  );
}