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

import {
  router,
  publicProcedure,
  createCallerFactory,
} from "../../trpc/server.js";
import { generateObject, generateText, jsonSchema, streamText } from "ai";
import type {
  OtherParameters,
  CommittedMessage,
  DraftMessage,
} from "../../types.js";
// import { client } from "../../database/milvus";
// import {
//   ConsistencyLevelEnum,
//   type NumberArrayId,
// } from "@zilliz/milvus2-sdk-node";
import { db } from "../../database/index.js";
import { conversations } from "./conversations.js";
import { messages } from "./messages.js";
import { facts, createCaller as createCallerFacts } from "./facts.js";
import { createCaller as createCallerMessages } from "./messages.js";
import { createCaller as createCallerFactTriggers } from "./fact-triggers.js";
import { factTriggers } from "./fact-triggers.js";
import { MODEL_NAME, openrouter } from "./provider.js";
import type { Fact, FactTrigger } from "../../database/common.js";

const factsCaller = createCallerFacts({});
const messagesCaller = createCallerMessages({});
const factTriggerCaller = createCallerFactTriggers({});

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>
`;

export const chat = router({
  conversations,
  messages,
  facts,
  factTriggers,
  streamMessage: publicProcedure
    .input(
      (x) =>
        x as {
          message: string;
        }
    )
    .subscription(async function* ({ input, signal }) {
      const result = streamText({
        model: openrouter(MODEL_NAME),
        messages: [{ role: "user" as const, content: input.message }],
        abortSignal: signal,
      });
      for await (const chunk of result.textStream) {
        yield chunk;
      }
    }),
  sendMessage: publicProcedure
    .input(
      (x) =>
        x as {
          conversationId: string;
          messages: Array<DraftMessage | CommittedMessage>;
          systemPrompt: string;
          parameters: OtherParameters;
        }
    )
    .subscription(async function* ({
      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
      );

      // Emit status update
      yield {
        status: "saving_user_message",
        message: "Saving user message...",
      } as const;

      /** Save the incoming message to the database. */
      const insertedUserMessage = await db.messages.create({
        conversationId,
        // content: messages[messages.length - 1].content,
        // role: "user" as const,
        ...messages[messages.length - 1],
        index: messages.length - 1,
        createdAt: new Date().toISOString(),
      });

      // Emit status update
      yield {
        status: "generating_response",
        message: "Generating AI response...",
      } as const;

      /** 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(MODEL_NAME),
        messages: [
          previousRunningSummary === ""
            ? {
                role: "system" as const,
                content: systemPrompt,
              }
            : {
                role: "system" as const,
                content: mainSystemPrompt({
                  systemPrompt,
                  previousRunningSummary,
                }),
              },
          ...messagesSincePreviousRunningSummary.map((m) => ({
            role: m.role,
            content: m.parts
              .filter((p) => p.type === "text")
              .map((p) => p.text)
              .join(""),
          })),
        ],
        tools: undefined,
        ...parameters,
      });

      // Emit status update
      yield {
        status: "extracting_facts_from_user",
        message: "Extracting facts from user message...",
      } as const;

      /** 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 factsCaller.extractFromNewMessages({
          previousRunningSummary,
          messagesSincePreviousRunningSummary: [],
          newMessages: messagesSincePreviousRunningSummary,
        });
      const insertedFactsFromUserMessage = await db.facts.createMany(
        factsFromUserMessageResponse.object.facts.map((fact) => ({
          userId: "019900bb-61b3-7333-b760-b27784dfe33b",
          sourceMessageId: insertedUserMessage.id,
          content: fact,
        }))
      );

      // Emit status update
      yield {
        status: "generating_summary",
        message: "Generating conversation summary...",
      } as const;

      /** 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 messagesCaller.generateRunningSummary({
          messagesSincePreviousRunningSummary,
          mainResponseContent: mainResponse.text,
          previousRunningSummary,
        });
      const insertedAssistantMessage = await db.messages.create({
        conversationId,
        // content: mainResponse.text,
        parts: [{ type: "text", text: mainResponse.text }],
        runningSummary: runningSummaryResponse.text,
        role: "assistant" as const,
        index: messages.length,
        createdAt: new Date().toISOString(),
      });

      // Emit status update
      yield {
        status: "extracting_facts_from_assistant",
        message: "Extracting facts from assistant response...",
      } as const;

      /** 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 factsCaller.extractFromNewMessages({
          previousRunningSummary,
          messagesSincePreviousRunningSummary,
          newMessages: [
            {
              role: "assistant" as const,
              // content: mainResponse.text,
              parts: [{ type: "text", text: mainResponse.text }],
            },
          ],
        });

      const insertedFactsFromAssistantMessage = await db.facts.createMany(
        factsFromAssistantMessageResponse.object.facts.map((factContent) => ({
          userId: "019900bb-61b3-7333-b760-b27784dfe33b",
          sourceMessageId: insertedAssistantMessage.id,
          content: factContent,
          createdAt: new Date().toISOString(),
        }))
      );

      const insertedFacts = [
        ...insertedFactsFromUserMessage,
        ...insertedFactsFromAssistantMessage,
      ];

      // Emit status update
      yield {
        status: "generating_fact_triggers",
        message: "Generating fact triggers...",
      } as const;

      /** 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"). */
      for (const fact of insertedFacts) {
        const factTriggers = await factTriggerCaller.generateFromFact({
          mainResponseContent: mainResponse.text,
          previousRunningSummary,
          messagesSincePreviousRunningSummary,
          fact,
        });
        const insertedFactTriggers: Array<Omit<FactTrigger, "id">> =
          factTriggers.object.factTriggers.map((factTrigger) => ({
            sourceFactId: fact.id,
            content: factTrigger,
            priorityMultiplier: 1,
            priorityMultiplierReason: "",
            scopeConversationId: conversationId,
            createdAt: new Date().toISOString(),
          }));
        await db.factTriggers.createMany(insertedFactTriggers);
      }

      // Emit final result
      yield {
        status: "completed",
        message: "Completed!",
        result: {
          insertedAssistantMessage,
          insertedUserMessage,
          insertedFacts,
        },
      } as const;
    }),
});

export const createCaller = createCallerFactory(chat);