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/postgres";
import type { ConversationsId } from "../../database/generated/public/Conversations";
import type { UsersId } from "../../database/generated/public/Users";

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

export const chat = router({
  listConversations: publicProcedure.query(async () => {
    const rows = await db.selectFrom("conversations").selectAll().execute();
    return rows;
  }),
  fetchConversation: publicProcedure
    .input((x) => x as { id: number })
    .query(async ({ input: { id } }) => {
      const row = await db
        .selectFrom("conversations")
        .selectAll()
        .where("id", "=", id as ConversationsId)
        .executeTakeFirst();
      return row;
    }),
  createConversation: publicProcedure.mutation(async () => {
    const title = "New Conversation";
    const row = await db
      .insertInto("conversations")
      .values({
        title,
        user_id: 1 as UsersId,
      })
      .returningAll()
      .executeTakeFirst();
    return row;
  }),
  deleteConversation: publicProcedure
    .input((x) => x as { id: number })
    .mutation(async ({ input: { id } }) => {
      const result = await db
        .deleteFrom("conversations")
        .where("id", "=", Number(id) as ConversationsId)
        .execute();
      return result;
    }),
  updateConversationTitle: publicProcedure
    .input(
      (x) =>
        x as {
          id: number;
          title: string;
        },
    )
    .mutation(async ({ input: { id, title } }) => {
      const result = await db
        .updateTable("conversations")
        .set({ title })
        .where("id", "=", Number(id) as ConversationsId)
        .execute();
      return result[0];
    }),
  fetchMessages: publicProcedure
    .input((x) => x as { conversationId: number })
    .query(async ({ input: { conversationId } }) => {
      const rows = await db
        .selectFrom("messages")
        .selectAll()
        .where("conversation_id", "=", conversationId as ConversationsId)
        .execute();
      return rows.map((row) => ({
        ...row,
        conversationId: conversationId as ConversationsId,
        runningSummary: row.running_summary,
      })) as Array<CommittedMessage>;
    }),
  sendMessage: publicProcedure
    .input(
      (x) =>
        x as {
          conversationId: number;
          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 = await db
          .insertInto("messages")
          .values({
            conversation_id: conversationId as ConversationsId,
            content: messages[messages.length - 1].content,
            role: "user" as const,
            index: messages.length - 1,
            created_at: new Date().toISOString(),
          })
          .returning(["id", "index"])
          .executeTakeFirst();
        /** 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: `${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>
`,
                },
            ...messages.slice(previousRunningSummaryIndex + 1),
          ],
          maxSteps: 3,
          tools: undefined,
          ...parameters,
        });
        console.log("sent", [
          previousRunningSummary === ""
            ? { role: "system" as const, content: systemPrompt }
            : {
                role: "system" as const,
                content: `${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>
`,
              },
          ...messages.slice(previousRunningSummaryIndex + 1),
        ]);
        /** 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: `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>
`,
                },
                ...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 = await db
          .insertInto("messages")
          .values({
            conversation_id: conversationId as ConversationsId,
            content: mainResponse.text,
            running_summary: runningSummaryResponse.text,
            role: "assistant" as const,
            index: messages.length,
            created_at: new Date().toISOString(),
          })
          .returningAll()
          .executeTakeFirst();
        /** 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);