From 8b78d975e6bef924b20dbe290033c85eccb13247 Mon Sep 17 00:00:00 2001
From: mantagen <contact.mantagen@gmail.com>
Date: Wed, 28 Aug 2024 17:46:46 +0200
Subject: [PATCH] add other id generate functions -- with role prefix

---
 .../src/components/ContextProviders/ChatProvider.tsx   |  2 +-
 packages/aila/src/core/chat/AilaChat.ts                |  8 ++++----
 packages/aila/src/helpers/chat/generateMessageId.ts    |  8 ++++++--
 packages/api/src/router/appSessions.ts                 | 10 +++++-----
 4 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/apps/nextjs/src/components/ContextProviders/ChatProvider.tsx b/apps/nextjs/src/components/ContextProviders/ChatProvider.tsx
index 840d8a1bb..2b68a9eba 100644
--- a/apps/nextjs/src/components/ContextProviders/ChatProvider.tsx
+++ b/apps/nextjs/src/components/ContextProviders/ChatProvider.tsx
@@ -155,7 +155,7 @@ export function ChatProvider({
   } = useChat({
     sendExtraMessageFields: true,
     initialMessages,
-    generateId: generateMessageId,
+    generateId: () => generateMessageId({ role: "user" }),
     id,
     body: {
       id,
diff --git a/packages/aila/src/core/chat/AilaChat.ts b/packages/aila/src/core/chat/AilaChat.ts
index f0e671720..e320040cb 100644
--- a/packages/aila/src/core/chat/AilaChat.ts
+++ b/packages/aila/src/core/chat/AilaChat.ts
@@ -3,7 +3,6 @@ import {
   unsupportedSubjects,
   subjectWarnings,
 } from "@oakai/core/src/utils/subjects";
-import { nanoid } from "nanoid";
 import invariant from "tiny-invariant";
 
 import { AilaChatService, AilaServices } from "../..";
@@ -12,6 +11,7 @@ import {
   AilaGeneration,
   AilaGenerationStatus,
 } from "../../features/generation";
+import { generateMessageId } from "../../helpers/chat/generateMessageId";
 import { JsonPatchDocumentOptional } from "../../protocol/jsonPatchProtocol";
 import { LLMService } from "../llm/LLMService";
 import { OpenAIService } from "../llm/OpenAIService";
@@ -123,7 +123,7 @@ export class AilaChat implements AilaChatService {
   public async systemMessage() {
     invariant(this._generation?.systemPrompt, "System prompt not initialised");
     return {
-      id: nanoid(16),
+      id: generateMessageId({ role: "system" }),
       content: this._generation?.systemPrompt,
       role: "system" as const,
     };
@@ -139,7 +139,7 @@ export class AilaChat implements AilaChatService {
 
     if (this._aila?.lesson.hasSetInitialState) {
       applicableMessages.push({
-        id: nanoid(16),
+        id: generateMessageId({ role: "user" }),
         role: "user",
         content:
           "Now that you have the title, key stage and subject, let's start planning the lesson. Could you tell me if there are Oak lessons I could base my lesson on, or if there are none available let's get going with the first step of the lesson plan creation process!",
@@ -256,7 +256,7 @@ export class AilaChat implements AilaChatService {
       return;
     }
     const assistantMessage: Message = {
-      id: nanoid(16),
+      id: generateMessageId({ role: "assistant" }),
       role: "assistant",
       content,
     };
diff --git a/packages/aila/src/helpers/chat/generateMessageId.ts b/packages/aila/src/helpers/chat/generateMessageId.ts
index 6f7302850..79340d5b4 100644
--- a/packages/aila/src/helpers/chat/generateMessageId.ts
+++ b/packages/aila/src/helpers/chat/generateMessageId.ts
@@ -1,5 +1,9 @@
 import { nanoid } from "nanoid";
 
-export function generateMessageId() {
-  return `${nanoid(16)}`;
+export function generateMessageId({
+  role,
+}: {
+  role: "user" | "assistant" | "system";
+}) {
+  return `${role[0]}-${nanoid(16)}`;
 }
diff --git a/packages/api/src/router/appSessions.ts b/packages/api/src/router/appSessions.ts
index bac04479e..1c48d67ed 100644
--- a/packages/api/src/router/appSessions.ts
+++ b/packages/api/src/router/appSessions.ts
@@ -6,11 +6,11 @@ import { RateLimitExceededError } from "@oakai/core/src/utils/rateLimiting/userB
 import { Prisma, PrismaClientWithAccelerate } from "@oakai/db";
 import * as Sentry from "@sentry/nextjs";
 import { TRPCError } from "@trpc/server";
-import { nanoid } from "nanoid";
 import { isTruthy } from "remeda";
 import { z } from "zod";
 
 import { getSessionModerations } from "../../../aila/src/features/moderation/getSessionModerations";
+import { generateChatId } from "../../../aila/src/helpers/chat/generateChatId";
 import {
   AilaPersistedChat,
   chatSchema,
@@ -135,11 +135,11 @@ export const appSessionsRouter = router({
 
       await checkMutationPermissions(userId);
 
-      const id = nanoid(16);
+      const chatId = generateChatId();
 
       const output: AilaPersistedChat = {
-        id,
-        path: `/aila/${id}`,
+        id: chatId,
+        path: `/aila/${chatId}`,
         title: "",
         topic: "",
         userId,
@@ -154,7 +154,7 @@ export const appSessionsRouter = router({
 
       const appSession = {
         data: {
-          id: nanoid(16),
+          id: chatId,
           appId,
           userId,
           output,