context-efficient-ai/pages/chat/trpc.ts

import {
  router,
  publicProcedure,
  createCallerFactory,
} from "../../trpc/server";
import { createOpenRouter } from "@openrouter/ai-sdk-provider";
import { generateText } from "ai";
import type { Message as UIMessage } from "ai";
import type {
  OtherParameters,
  CommittedMessage,
  DraftMessage,
} from "../../types.js";
import { env } from "../../server/env.js";
// import { client } from "../../database/milvus";
// import {
//   ConsistencyLevelEnum,
//   type NumberArrayId,
// } from "@zilliz/milvus2-sdk-node";
import { db } from "../../database/lowdb";
import { nanoid } from "nanoid";

const mainSystemPrompt = ({
  systemPrompt,
  previousRunningSummary,
}: { systemPrompt: string; previousRunningSummary: string }) => `${systemPrompt}

This is a summary of the conversation so far, from your point-of-view (so "I" and "me" refer to you):
<running_summary>
${previousRunningSummary}
</running_summary>
`;
const runningSummarySystemPrompt = ({
  previousRunningSummary,
}: {
  previousRunningSummary: string;
}) => `Given the following summary of a conversation, coupled with the messages exchanged since that summary was produced, produce a new summary of the conversation.
<running_summary>
  ${previousRunningSummary}
</running_summary>
`;

const openrouter = createOpenRouter({
  apiKey: env.OPENROUTER_API_KEY,
});

export const chat = router({
  listConversations: publicProcedure.query(async () => {
    const rows = await db.data.conversations;
    return rows;
  }),
  fetchConversation: publicProcedure
    .input((x) => x as { id: string })
    .query(async ({ input: { id } }) => {
      const row = await db.data.conversations.find((c) => c.id === id);
      return row;
    }),
  createConversation: publicProcedure.mutation(async () => {
    const title = "New Conversation";
    const row = {
      id: nanoid(),
      title,
      userId: "1",
    };
    await db.data.conversations.push(row);
    db.write();
    return row;
  }),
  deleteConversation: publicProcedure
    .input((x) => x as { id: string })
    .mutation(async ({ input: { id } }) => {
      await db.data.conversations.splice(
        db.data.conversations.findIndex((c) => c.id === id),
        1,
      );
      db.write();
      return { ok: true };
    }),
  updateConversationTitle: publicProcedure
    .input(
      (x) =>
        x as {
          id: string;
          title: string;
        },
    )
    .mutation(async ({ input: { id, title } }) => {
      const conversation = await db.data.conversations.find((c) => c.id === id);
      if (!conversation) throw new Error("Conversation not found");
      conversation.title = title;
      db.write();
      return { ok: true };
    }),
  fetchMessages: publicProcedure
    .input((x) => x as { conversationId: string })
    .query(async ({ input: { conversationId } }) => {
      const rows = await db.data.messages.filter(
        (m) => m.conversationId === conversationId,
      );
      return rows as Array<CommittedMessage>;
    }),
  sendMessage: publicProcedure
    .input(
      (x) =>
        x as {
          conversationId: string;
          messages: Array<DraftMessage | CommittedMessage>;
          systemPrompt: string;
          parameters: OtherParameters;
        },
    )
    .mutation(
      async ({
        input: { conversationId, messages, systemPrompt, parameters },
      }) => {
        /** TODO: Save all unsaved messages (i.e. those without an `id`) to the
         * database. Is this dangerous? Can an attacker just send a bunch of
         * messages, omitting the ids, causing me to save a bunch of them to the
         * database? I guess it's no worse than starting new converations, which
         * anyone can freely do. */
        const previousRunningSummaryIndex = messages.findLastIndex(
          (message) =>
            typeof (message as CommittedMessage).runningSummary !== "undefined",
        );
        const previousRunningSummary =
          previousRunningSummaryIndex >= 0
            ? ((messages[previousRunningSummaryIndex] as CommittedMessage)
                .runningSummary as string)
            : "";
        /** Save the incoming message to the database. */
        const insertedUserMessage: CommittedMessage = {
          id: nanoid(),
          conversationId,
          content: messages[messages.length - 1].content,
          role: "user" as const,
          index: messages.length - 1,
          createdAt: new Date().toISOString(),
        };
        await db.data.messages.push(insertedUserMessage);
        // do not db.write() until the end

        /** Generate a new message from the model, but hold-off on adding it to
         * the database until we produce the associated running-summary, below.
         * The model should be given the conversation summary thus far, and of
         * course the user's latest message, unmodified. Invite the model to
         * create any tools it needs. The tool needs to be implemented in a
         * language which this system can execute; usually an interpretted
         * language like Python or JavaScript. */
        const mainResponse = await generateText({
          model: openrouter("mistralai/mistral-nemo"),
          messages: [
            previousRunningSummary === ""
              ? { role: "system" as const, content: systemPrompt }
              : {
                  role: "system" as const,
                  content: mainSystemPrompt({
                    systemPrompt,
                    previousRunningSummary,
                  }),
                },
            ...messages.slice(previousRunningSummaryIndex + 1),
          ],
          maxSteps: 3,
          tools: undefined,
          ...parameters,
        });
        /** Extract Facts from the user's message, and add them to the database,
         * linking the Facts with the messages they came from. (Yes, this should
         * be done *after* the model response, not before; because when we run a
         * query to find Facts to inject into the context sent to the model, we
         * don't want Facts from the user's current message to be candidates for
         * injection, because we're sending the user's message unadulterated to
         * the model; there's no reason to inject the same Facts that the model is
         * already using to generate its response.) */
        /** Extract Facts from the model's response, and add them to the database,
         * linking the Facts with the messages they came from. */
        /** For each Fact produced in the two fact-extraction steps, generate
         * FactTriggers and add them to the database, linking the FactTriggers
         * with the Facts they came from. A FactTrigger is a natural language
         * phrase that describes a situation in which it would be useful to invoke
         * the Fact. (e.g., "When food preferences are discussed"). */
        /** Produce a running summary of the conversation, and save that along
         * with the model's response to the database. The new running summary is
         * based on the previous running summary combined with the all messages
         * since that summary was produced. */
        const runningSummaryResponse = previousRunningSummary
          ? await generateText({
              model: openrouter("mistralai/mistral-nemo"),
              messages: [
                {
                  role: "system" as const,
                  content: runningSummarySystemPrompt({
                    previousRunningSummary,
                  }),
                },
                ...messages.slice(previousRunningSummaryIndex + 1),
                {
                  role: "assistant" as const,
                  content: mainResponse.text,
                } as UIMessage,
                /** I might need this next message, because models are trained to
                 * respond when the final message in `messages` is from the `user`,
                 * but in our case it's an `assistant` message, so I'm artificially
                 * adding a `user` message to the end of the conversation. */
                {
                  role: "user" as const,
                  content: "What is the new summary of the conversation?",
                } as UIMessage,
              ],
              maxSteps: 3,
              tools: undefined,
              ...parameters,
            })
          : await generateText({
              model: openrouter("mistralai/mistral-nemo"),
              messages: [
                {
                  role: "system" as const,
                  content:
                    "Given the following messages of a conversation, produce a summary of the conversation.",
                },
                ...messages,
                {
                  role: "assistant" as const,
                  content: mainResponse.text,
                } as UIMessage,
                /** I might need this next message, because models are trained to
                 * respond when the final message in `messages` is from the `user`,
                 * but in our case it's an `assistant` message, so I'm artificially
                 * adding a `user` message to the end of the conversation. */
                {
                  role: "user" as const,
                  content: "What is the new summary of the conversation?",
                } as UIMessage,
              ],
              maxSteps: 3,
              tools: undefined,
              ...parameters,
            });
        const insertedAssistantMessage: CommittedMessage = {
          id: nanoid(),
          conversationId,
          content: mainResponse.text,
          runningSummary: runningSummaryResponse.text,
          role: "assistant" as const,
          index: messages.length,
          createdAt: new Date().toISOString(),
        };
        await db.data.messages.push(insertedAssistantMessage);
        await db.write();
        /** TODO: notify the caller, somehow, that some messages were saved to
         * the database and/or were outfitted with runningSummaries, so the
         * caller can update its UI state. */
        return { insertedAssistantMessage, insertedUserMessage };
      },
    ),
});

export const createCaller = createCallerFactory(chat);