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

import {
  Box,
  Group,
  JsonInput,
  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 { ConversationsId } from "../../../database/generated/public/Conversations";
import type { CommittedMessage, DraftMessage } from "../../../types";
import Markdown from "react-markdown";

export default function ChatPage() {
  const pageContext = usePageContext();
  const conversationId = Number(pageContext.routeParams.id) as ConversationsId;
  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 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 setLoading = useStore((state) => state.setLoading);

  const { conversation, messages: initialMessages } = 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]);

  return (
    <>
      <div>
        <span>Conversation #{conversationId} - </span>
        <input
          type="text"
          value={conversationTitle || ""}
          onChange={(e) => {
            setConversationTitle(e.target.value);
          }}
          onBlur={(e) => {
            trpc.chat.updateConversationTitle.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.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?.running_summary ||
                      undefined,
                  } as CommittedMessage,
                ];
                setMessages(messagesWithAssistantMessage);
                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>
    </>
  );
}

function Messages({
  messages,
}: {
  messages: Array<DraftMessage | CommittedMessage>;
}) {
  const theme = useMantineTheme();
  console.log("messages", messages);
  return (
    <Stack gap="md" justify="flex-start">
      {messages.map((message, index) => (
        <Group
          // biome-ignore lint/suspicious/noArrayIndexKey: <explanation>
          key={index}
          justify={message.role === "user" ? "flex-end" : "flex-start"}
        >
          <Box
            w="75%"
            bg={
              message.role === "user"
                ? theme.colors.gray[2]
                : theme.colors.blue[2]
            }
            p="md"
            bdrs="md"
          >
            <div>
              {"index" in message ? message.index : ""}
              {message.role}
            </div>
            <Markdown>{message.content}</Markdown>
            {"runningSummary" in message && (
              <div>
                <strong>Running Summary:</strong> {message.runningSummary}
              </div>
            )}
          </Box>
        </Group>
      ))}
    </Stack>
  );
}