getChat()をキャッシュ

Author: na-trium-144

app/[docs_id]/chatHistory.tsx

@@ -1,13 +1,14 @@
 "use client";
 
-import { ChatWithMessages } from "@/lib/chatHistory";
+import { ChatWithMessages, getChat } from "@/lib/chatHistory";
 import {
   createContext,
   ReactNode,
   useContext,
   useEffect,
   useState,
 } from "react";
+import useSWR from "swr";
 
 export interface IChatHistoryContext {
   chatHistories: ChatWithMessages[];
@@ -27,19 +28,32 @@ export function useChatHistoryContext() {
 
 export function ChatHistoryProvider({
   children,
+  docs_id,
   initialChatHistories,
 }: {
   children: ReactNode;
+  docs_id: string;
   initialChatHistories: ChatWithMessages[];
 }) {
   const [chatHistories, setChatHistories] =
     useState<ChatWithMessages[]>(initialChatHistories);
+  // 最初はSSRで取得したinitialChatHistoriesを使用(キャッシュされているので古い可能性がある)
   useEffect(() => {
     setChatHistories(initialChatHistories);
   }, [initialChatHistories]);
+  // その後、クライアント側で最新のchatHistoriesを改めて取得して更新する
+  const { data: fetchedChatHistories } = useSWR<ChatWithMessages[]>(
+    docs_id,
+    getChat
+  );
+  useEffect(() => {
+    if (fetchedChatHistories) {
+      setChatHistories(fetchedChatHistories);
+    }
+  }, [fetchedChatHistories]);
 
+  // チャットを更新した際にはクライアントサイドでchatHistoryに反映する
   const addChat = (chat: ChatWithMessages) => {
-    // サーバー側で追加された新しいchatをクライアント側にも反映する
     setChatHistories([...chatHistories, chat]);
   };
 

app/[docs_id]/page.tsx

@@ -6,7 +6,7 @@ import { MarkdownSection, splitMarkdown } from "./splitMarkdown";
 import pyodideLock from "pyodide/pyodide-lock.json";
 import { PageContent } from "./pageContent";
 import { ChatHistoryProvider } from "./chatHistory";
-import { getChat } from "@/lib/chatHistory";
+import { getChat, getChatFromCache } from "@/lib/chatHistory";
 
 export default async function Page({
   params,
@@ -50,10 +50,10 @@ export default async function Page({
     splitMarkdown(text)
   );
 
-  const initialChatHistories = getChat(docs_id);
+  const initialChatHistories = getChatFromCache(docs_id);
 
   return (
-    <ChatHistoryProvider initialChatHistories={await initialChatHistories}>
+    <ChatHistoryProvider initialChatHistories={await initialChatHistories} docs_id={docs_id}>
       <PageContent
         documentContent={await mdContent}
         splitMdContent={await splitMdContent}

app/lib/cache.ts

@@ -0,0 +1,20 @@
+const cacheData: Map<string, ArrayBuffer> = new Map();
+/**
+ * nodejsにcache apiがないので、web標準のcache APIに相当するものの自前実装
+ */
+export const inMemoryCache = {
+  async put(key: string, response: Response): Promise<void> {
+    const arrayBuffer = await response.arrayBuffer();
+    cacheData.set(key, arrayBuffer);
+  },
+  async match(key: string): Promise<Response | undefined> {
+    const arrayBuffer = cacheData.get(key);
+    if (arrayBuffer) {
+      return new Response(arrayBuffer);
+    }
+    return undefined;
+  },
+  async delete(key: string): Promise<boolean> {
+    return cacheData.delete(key);
+  },
+} as const;

app/lib/chatHistory.ts

@@ -1,30 +1,74 @@
+"use server";
+
 import { headers } from "next/headers";
 import { getAuthServer } from "./auth";
 import { getDrizzle } from "./drizzle";
 import { chat, message } from "@/schema/chat";
 import { and, asc, eq } from "drizzle-orm";
+import { Auth } from "better-auth";
+import { inMemoryCache } from "./cache";
 
 export interface CreateChatMessage {
   role: "user" | "ai" | "error";
   content: string;
 }
 
+// cacheに使うキーで、実際のURLではない
+const CACHE_KEY_BASE = "https://my-code.utcode.net/chatHistory";
+
+interface Context {
+  drizzle: Awaited<ReturnType<typeof getDrizzle>>;
+  auth: Auth;
+  userId: string;
+}
+/**
+ * drizzleが初期化されてなければ初期化し、
+ * authが初期化されてなければ初期化し、
+ * userIdがなければセッションから取得してセットする。
+ */
+async function initAll(ctx?: Partial<Context>): Promise<Context> {
+  if (!ctx) {
+    ctx = {};
+  }
+  if (!ctx.drizzle) {
+    ctx.drizzle = await getDrizzle();
+  }
+  if (!ctx.auth) {
+    ctx.auth = await getAuthServer(ctx.drizzle);
+  }
+  if (!ctx.userId) {
+    const session = await ctx.auth.api.getSession({
+      headers: await headers(),
+    });
+    if (!session) {
+      throw new Error("Not authenticated");
+    }
+    ctx.userId = session.user.id;
+  }
+  return ctx as Context;
+}
+async function getCache() {
+  if ("caches" in globalThis) {
+    // worker
+    return await caches.open("chatHistory");
+  } else {
+    // nodejs
+    return inMemoryCache;
+  }
+}
+
 export async function addChat(
   docsId: string,
   sectionId: string,
-  messages: CreateChatMessage[]
+  messages: CreateChatMessage[],
+  context?: Partial<Context>
 ) {
-  const drizzle = await getDrizzle();
-  const auth = await getAuthServer(drizzle);
-  const session = await auth.api.getSession({ headers: await headers() });
-  if (!session) {
-    throw new Error("Not authenticated");
-  }
+  const { drizzle, userId } = await initAll(context);
 
   const [newChat] = await drizzle
     .insert(chat)
     .values({
-      userId: session.user.id,
+      userId,
       docsId,
       sectionId,
     })
@@ -41,6 +85,14 @@ export async function addChat(
     )
     .returning();
 
+  console.log(
+    `deleting cache for chatHistory/getChat for user ${userId} and docs ${docsId}`
+  );
+  const cache = await getCache();
+  await cache.delete(
+    `${CACHE_KEY_BASE}/getChat?docsId=${docsId}&userId=${userId}`
+  );
+
   return {
     ...newChat,
     messages: chatMessages,
@@ -49,16 +101,14 @@ export async function addChat(
 
 export type ChatWithMessages = Awaited<ReturnType<typeof addChat>>;
 
-export async function getChat(docsId: string) {
-  const drizzle = await getDrizzle();
-  const auth = await getAuthServer(drizzle);
-  const session = await auth.api.getSession({ headers: await headers() });
-  if (!session) {
-    return [];
-  }
+export async function getChat(
+  docsId: string,
+  context?: Partial<Context>
+): Promise<ChatWithMessages[]> {
+  const { drizzle, userId } = await initAll(context);
 
   const chats = await drizzle.query.chat.findMany({
-    where: and(eq(chat.userId, session.user.id), eq(chat.docsId, docsId)),
+    where: and(eq(chat.userId, userId), eq(chat.docsId, docsId)),
     with: {
       messages: {
         orderBy: [asc(message.createdAt)],
@@ -67,8 +117,33 @@ export async function getChat(docsId: string) {
     orderBy: [asc(chat.createdAt)],
   });
 
+  const cache = await getCache();
+  await cache.put(
+    `${CACHE_KEY_BASE}/getChat?docsId=${docsId}&userId=${userId}`,
+    new Response(JSON.stringify(chats), {
+      headers: { "Cache-Control": "max-age=86400, s-maxage=86400" },
+    })
+  );
   return chats;
 }
+export async function getChatFromCache(
+  docsId: string,
+  context?: Partial<Context>
+) {
+  const { drizzle, auth, userId } = await initAll(context);
+
+  const cache = await getCache();
+  const cachedResponse = await cache.match(
+    `${CACHE_KEY_BASE}/getChat?docsId=${docsId}&userId=${userId}`
+  );
+  if (cachedResponse) {
+    console.log("Cache hit for chatHistory/getChat");
+    const data = (await cachedResponse.json()) as ChatWithMessages[];
+    return data;
+  }
+  console.log("Cache miss for chatHistory/getChat");
+  return await getChat(docsId, { drizzle, auth, userId });
+}
 
 export async function migrateChatUser(oldUserId: string, newUserId: string) {
   const drizzle = await getDrizzle();