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

import {
  router,
  publicProcedure,
  createCallerFactory,
} from "../../trpc/server";
import { createOpenRouter } from "@openrouter/ai-sdk-provider";
import { generateObject, generateText, jsonSchema } 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, type Fact } 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 factsFromUserMessageSystemPrompt = ({
  previousRunningSummary,
}: {
  previousRunningSummary: string;
}) => `You are an expert at extracting facts from conversations. Given the following summary of a conversation, coupled with the messages exchanged since that summary was produced (which will be provided by the user), extract new facts that can be gleaned from the messages exchanged since the summary was produced.
<running_summary>
  ${previousRunningSummary}
</running_summary>
`;

const factsFromUserMessageUserPrompt = ({
  messagesSincePreviousRunningSummary,
}: {
  messagesSincePreviousRunningSummary: Array<DraftMessage>;
}) =>
  `${messagesSincePreviousRunningSummary.map(
    (message) =>
      `<${message.role}_message>${message.content}</${message.role}_message>`,
  )}

Extract new facts from these messages.`;

const factsFromAssistantMessageSystemPrompt = ({
  previousRunningSummary,
}: {
  previousRunningSummary: string;
}) => `You are an insightful fact extractor. Given the following summary of a conversation, coupled with the messages exchanged since that summary was produced (which will be provided by the user), extract new facts that can be gleaned from the final assistant response.
<running_summary>
  ${previousRunningSummary}
</running_summary>
`;

const factsFromAssistantMessageUserPrompt = ({
  messagesSincePreviousRunningSummary,
  mainResponseContent,
}: {
  messagesSincePreviousRunningSummary: Array<DraftMessage>;
  mainResponseContent: string;
}) =>
  `${messagesSincePreviousRunningSummary.map(
    (message) =>
      `<${message.role}_message>${message.content}</${message.role}_message>`,
  )}
<assistant_response>
${mainResponseContent}
</assistant_response>

Extract facts from the assistant's response.`;

const runningSummarySystemPrompt = ({
  previousRunningSummary,
}: {
  previousRunningSummary: string;
}) => `You are an expert at summarizing conversations. Given the following summary of a conversation, coupled with the messages exchanged since that summary was produced (which will be provided by the user), produce a new summary of the conversation.
<running_summary>
  ${previousRunningSummary}
</running_summary>
`;

const runningSummaryUserPrompt = ({
  messagesSincePreviousRunningSummary,
  mainResponseContent,
}: {
  messagesSincePreviousRunningSummary: Array<DraftMessage>;
  mainResponseContent: string;
}) =>
  `${messagesSincePreviousRunningSummary.map(
    (message) =>
      `<${message.role}_message>${message.content}</${message.role}_message>`,
  )}
<assistant_response>
${mainResponseContent}
</assistant_response>

Generate a new running summary of the conversation.`;

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 } }) => {
      db.data.conversations.splice(
        db.data.conversations.findIndex((c) => c.id === id),
        1,
      );
      const deletedMessageIds = db.data.messages
        .filter((m) => m.conversationId === id)
        .map((m) => m.id);
      db.data.messages = db.data.messages.filter(
        (m) => m.conversationId !== id,
      );
      db.data.facts = db.data.facts.filter(
        (fact) => !deletedMessageIds.includes(fact.sourceMessageId),
      );
      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)
            : "";
        const messagesSincePreviousRunningSummary = messages.slice(
          previousRunningSummaryIndex + 1,
        );
        /** 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(),
        };
        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,
                  }),
                },
            ...messagesSincePreviousRunningSummary,
          ],
          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.) */
        const factsFromUserMessageResponse = await generateObject<{
          facts: Array<string>;
        }>({
          model: openrouter("mistralai/mistral-nemo"),
          messages: [
            {
              role: "system" as const,
              content: factsFromUserMessageSystemPrompt({
                previousRunningSummary,
              }),
            },
            {
              role: "user" as const,
              content: factsFromUserMessageUserPrompt({
                messagesSincePreviousRunningSummary,
              }),
            },
          ],
          schema: jsonSchema({
            type: "object",
            properties: {
              facts: {
                type: "array",
                items: {
                  type: "string",
                },
              },
            },
          }),
          maxSteps: 3,
          tools: undefined,
          ...parameters,
        });
        const insertedFactsFromUserMessage: Array<Fact> =
          factsFromUserMessageResponse.object.facts.map((fact) => ({
            id: nanoid(),
            userId: "1",
            sourceMessageId: insertedUserMessage.id,
            content: fact,
            createdAt: new Date().toISOString(),
          }));
        db.data.facts.push(...insertedFactsFromUserMessage);
        /** Extract Facts from the model's response, and add them to the database,
         * linking the Facts with the messages they came from. */
        const factsFromAssistantMessageResponse = await generateObject<{
          facts: Array<string>;
        }>({
          model: openrouter("mistralai/mistral-nemo"),
          messages: [
            {
              role: "system" as const,
              content: factsFromAssistantMessageSystemPrompt({
                previousRunningSummary,
              }),
            },
            /** Yes, the next message is a `user` message, because models are
             * trained to respond to `user` messages. So we wrap the assistant
             * response in XML tags to show that it's not the user speaking,
             * rather it's input for the model to process. The user is only
             * saying "Extract facts..." */
            {
              role: "user" as const,
              content: factsFromAssistantMessageUserPrompt({
                messagesSincePreviousRunningSummary,
                mainResponseContent: mainResponse.text,
              }),
            },
          ],
          schema: jsonSchema({
            type: "object",
            properties: {
              facts: {
                type: "array",
                items: {
                  type: "string",
                },
              },
            },
          }),
          maxSteps: 3,
          tools: undefined,
          ...parameters,
        });
        const insertedFactsFromAssistantMessage: Array<Fact> =
          factsFromAssistantMessageResponse.object.facts.map((factContent) => ({
            id: nanoid(),
            userId: "1",
            sourceMessageId: insertedAssistantMessage.id,
            content: factContent,
            createdAt: new Date().toISOString(),
          }));
        db.data.facts.push(...insertedFactsFromAssistantMessage);
        /** 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 = await generateText({
          model: openrouter("mistralai/mistral-nemo"),
          messages: [
            {
              role: "system" as const,
              content: runningSummarySystemPrompt({
                previousRunningSummary,
              }),
            },
            {
              role: "user" as const,
              content: runningSummaryUserPrompt({
                messagesSincePreviousRunningSummary,
                mainResponseContent: mainResponse.text,
              }),
            },
          ],
          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(),
        };
        db.data.messages.push(insertedAssistantMessage);

        await db.write();

        return {
          insertedAssistantMessage,
          insertedUserMessage,
          insertedFacts: [
            ...insertedFactsFromUserMessage,
            ...insertedFactsFromAssistantMessage,
          ],
        };
      },
    ),
});

export const createCaller = createCallerFactory(chat);