Explorar o código

Implement getThreadRelationAndUpdateExpiration()

Shun Miyazawa hai 1 ano
pai
achega
58cae3fb8c

+ 19 - 12
apps/app/src/features/openai/server/models/thread-relation.ts

@@ -5,6 +5,12 @@ import { getOrCreateModel } from '~/server/util/mongoose-utils';
 
 const DAYS_UNTIL_EXPIRATION = 30;
 
+const generateExpirationDate = (): Date => {
+  const currentDate = new Date();
+  const expirationDate = new Date(currentDate.setDate(currentDate.getDate() + DAYS_UNTIL_EXPIRATION));
+  return expirationDate;
+};
+
 interface Thread {
   threadId: string;
   expiredAt: Date;
@@ -18,7 +24,7 @@ interface ThreadDocument extends ThreadRelation, Document {}
 
 interface ThreadRelationModel extends Model<ThreadDocument> {
   upsertThreadRelation(userId: string, threadId: string): Promise<void>;
-  getThread(userId: string, threadId: string): Promise<Thread | undefined>;
+  getThreadRelationAndUpdateExpiration(userId: string, threadId: string): Promise<ThreadRelation | null>
 }
 
 const schema = new Schema<ThreadDocument, ThreadRelationModel>({
@@ -43,8 +49,7 @@ const schema = new Schema<ThreadDocument, ThreadRelationModel>({
 
 
 schema.statics.upsertThreadRelation = async function(userId: string, threadId: string) {
-  const currentDate = new Date();
-  const expirationDate = new Date(currentDate.setDate(currentDate.getDate() + DAYS_UNTIL_EXPIRATION));
+  const expirationDate = generateExpirationDate();
 
   await this.updateOne(
     { userId },
@@ -60,18 +65,20 @@ schema.statics.upsertThreadRelation = async function(userId: string, threadId: s
   );
 };
 
-schema.statics.getThread = async function(userId: string, threadId: string): Promise<Thread | undefined> {
-  const result = await this.findOne(
+
+schema.statics.getThreadRelationAndUpdateExpiration = async function(userId: string, threadId: string): Promise<ThreadRelation | null> {
+  const expirationDate = generateExpirationDate();
+
+  const result = await this.findOneAndUpdate(
     { userId, 'threads.threadId': threadId },
-    { threads: { $elemMatch: { threadId } } },
+    {
+      $set: { 'threads.$.expiredAt': expirationDate }, // Extend DAYS_UNTIL_EXPIRATION days from the retrieved time.
+    },
+    { new: true }, // 更新後のドキュメントを返す
   );
 
-  if (result != null && result.threads.length > 0) {
-    return {
-      threadId: result.threads[0].threadId,
-      expiredAt: result.threads[0].expiredAt,
-    };
-  }
+  return result;
 };
 
+
 export default getOrCreateModel<ThreadDocument, ThreadRelationModel>('ThreadRelation', schema);

+ 8 - 5
apps/app/src/features/openai/server/services/openai.ts

@@ -52,12 +52,15 @@ class OpenaiService implements IOpenaiService {
       return thread;
     }
 
-    const threadDocument = await ThreadRelationModel.getThread(userId, threadId);
-    if (threadDocument != null) {
-      // Check if a thread entity exists
-      const thread = await this.client.retrieveThread(threadDocument.threadId);
-      return thread;
+    const threadRelation = await ThreadRelationModel.getThreadRelationAndUpdateExpiration(userId, threadId);
+    const threadDocument = threadRelation?.threads.find(thread => thread.threadId === threadId);
+    if (threadDocument == null) {
+      return;
     }
+
+    // Check if a thread entity exists
+    const thread = await this.client.retrieveThread(threadDocument.threadId);
+    return thread;
   }
 
   public async getOrCreateVectorStoreForPublicScope(): Promise<VectorStoreDocument> {