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 { useTRPCClient } from "../../../trpc/client";

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 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 factTriggers = useStore((state) => state.factTriggers);
  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 setFactTriggers = useStore((state) => state.setFactTriggers);
  const removeFact = useStore((state) => state.removeFact);
  const removeFactTrigger = useStore((state) => state.removeFactTrigger);
  const setLoading = useStore((state) => state.setLoading);
  const trpc = useTRPCClient();

  // 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]);

  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]);

  useEffect(() => {
    setFactTriggers(initialFactTriggers);
  }, [initialFactTriggers, setFactTriggers]);

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

  async function handleUpdateFact(factId: string, content: string) {
    // Update the local state first
    setFacts(
      facts.map((fact) => (fact.id === factId ? { ...fact, content } : fact))
    );

    // Then update the database
    await trpc.chat.facts.update.mutate({ factId, content });
  }

  async function handleDeleteFactTrigger(factTriggerId: string) {
    removeFactTrigger(factTriggerId);
    await trpc.chat.factTriggers.deleteOne.mutate({ factTriggerId });
  }

  async function handleUpdateFactTrigger(
    factTriggerId: string,
    content: string
  ) {
    // Update the local state first
    setFactTriggers(
      factTriggers.map((factTrigger) =>
        factTrigger.id === factTriggerId
          ? { ...factTrigger, content }
          : factTrigger
      )
    );

    // Then update the database
    await trpc.chat.factTriggers.update.mutate({ factTriggerId, content });
  }

  return (
    <>
      <div>
        <span>Conversation #{conversationId} - </span>
        <input
          type="text"
          defaultValue={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.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();
                const messagesWithNewUserMessage = [
                  ...messages,
                  {
                    role: "user" as const,
                    parts: [{ type: "text", text: 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,
                    parts: [{ type: "text", text: message }],
                    index: response.insertedUserMessage?.index,
                    runningSummary: undefined,
                  } as CommittedMessage,
                  {
                    id: response.insertedAssistantMessage?.id,
                    conversationId,
                    role: "assistant" as const,
                    // content: response.insertedAssistantMessage?.content,
                    // parts: [{ type: "text", text: response.insertedAssistantMessage?.content }],
                    parts: response.insertedAssistantMessage?.parts,
                    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}>
                {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.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>
  );
}