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

import {
  ActionIcon,
  Box,
  Group,
  HoverCard,
  JsonInput,
  List,
  Stack,
  Tabs,
  Textarea,
  useMantineTheme,
} from "@mantine/core";
import { trpc } from "../../../trpc/client";
import { useEffect } 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 } from "@tabler/icons-react";

export default function ChatPage() {
  const pageContext = usePageContext();
  const conversationId = pageContext.routeParams.id;
  const conversationTitle = useStore(
    (state) => state.conversations.find((c) => c.id === conversationId)?.title,
  );
  const messages = useStore((state) => state.messages);
  const message = useStore((state) => state.message);
  const systemPrompt = useStore((state) => state.systemPrompt);
  const parameters = useStore((state) => state.parameters);
  const facts = useStore((state) => state.facts);
  const loading = useStore((state) => state.loading);
  const setConversationId = useStore((state) => state.setConversationId);
  const setConversationTitle = useStore((state) => state.setConversationTitle);
  const setMessages = useStore((state) => state.setMessages);
  const setMessage = useStore((state) => state.setMessage);
  const setSystemPrompt = useStore((state) => state.setSystemPrompt);
  const setParameters = useStore((state) => state.setParameters);
  const setFacts = useStore((state) => state.setFacts);
  const removeFact = useStore((state) => state.removeFact);
  const setLoading = useStore((state) => state.setLoading);

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

  useEffect(() => {
    setConversationId(conversationId);
  }, [conversationId, setConversationId]);

  useEffect(() => {
    if (conversation?.id && conversation?.title) {
      setConversationId(conversation.id);
      setConversationTitle(conversation.title);
    }
  }, [
    conversation?.id,
    conversation?.title,
    setConversationId,
    setConversationTitle,
  ]);

  useEffect(() => {
    setMessages(initialMessages);
  }, [initialMessages, setMessages]);

  useEffect(() => {
    setFacts(initialFacts);
  }, [initialFacts, setFacts]);

  async function handleDeleteFact(factId: string) {
    removeFact(factId);
    await trpc.chat.facts.deleteOne.mutate({ factId });
  }

  return (
    <>
      <div>
        <span>Conversation #{conversationId} - </span>
        <input
          type="text"
          value={conversationTitle || ""}
          onChange={(e) => {
            setConversationTitle(e.target.value);
          }}
          onBlur={(e) => {
            trpc.chat.conversations.updateTitle.mutate({
              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.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();
                const messagesWithNewUserMessage = [
                  ...messages,
                  { role: "user" as const, content: message } as DraftMessage,
                ];
                setMessages(messagesWithNewUserMessage);
                setLoading(true);
                const response = await trpc.chat.sendMessage.mutate({
                  conversationId,
                  messages: messagesWithNewUserMessage,
                  systemPrompt,
                  parameters,
                });
                const messagesWithAssistantMessage = [
                  ...messages,
                  {
                    id: response.insertedUserMessage?.id,
                    conversationId,
                    role: "user" as const,
                    content: message,
                    index: response.insertedUserMessage?.index,
                    runningSummary: undefined,
                  } as CommittedMessage,
                  {
                    id: response.insertedAssistantMessage?.id,
                    conversationId,
                    role: "assistant" as const,
                    content: response.insertedAssistantMessage?.content,
                    index: response.insertedAssistantMessage?.index,
                    runningSummary:
                      response.insertedAssistantMessage?.runningSummary ||
                      undefined,
                  } as CommittedMessage,
                ];
                setMessages(messagesWithAssistantMessage);
                setMessage("");
                setFacts(response.insertedFacts);
                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.map((fact) => (
              <List.Item key={fact.id}>
                {fact.content}{" "}
                <IconTrash
                  size={16}
                  stroke={1.5}
                  onClick={(e) => {
                    e.stopPropagation();
                    e.preventDefault();
                    handleDeleteFact(fact.id);
                  }}
                />
              </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.content}</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>
  );
}