thread-relation.ts 1.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859
  1. import { addDays } from 'date-fns';
  2. import { type Model, type Document, Schema } from 'mongoose';
  3. import { getOrCreateModel } from '~/server/util/mongoose-utils';
  4. import type { IThreadRelation } from '../../interfaces/thread-relation';
  5. const DAYS_UNTIL_EXPIRATION = 3;
  6. const generateExpirationDate = (): Date => {
  7. return addDays(new Date(), DAYS_UNTIL_EXPIRATION);
  8. };
  9. export interface ThreadRelationDocument extends IThreadRelation, Document {
  10. updateThreadExpiration(): Promise<void>;
  11. }
  12. interface ThreadRelationModel extends Model<ThreadRelationDocument> {
  13. getExpiredThreadRelations(limit?: number): Promise<ThreadRelationDocument[] | undefined>;
  14. }
  15. const schema = new Schema<ThreadRelationDocument, ThreadRelationModel>({
  16. userId: {
  17. type: Schema.Types.ObjectId,
  18. ref: 'User',
  19. required: true,
  20. },
  21. vectorStore: {
  22. type: Schema.Types.ObjectId,
  23. ref: 'VectorStore',
  24. required: true,
  25. },
  26. threadId: {
  27. type: String,
  28. required: true,
  29. unique: true,
  30. },
  31. title: {
  32. type: String,
  33. },
  34. expiredAt: {
  35. type: Date,
  36. default: generateExpirationDate,
  37. required: true,
  38. },
  39. });
  40. schema.statics.getExpiredThreadRelations = async function(limit?: number): Promise<ThreadRelationDocument[] | undefined> {
  41. const currentDate = new Date();
  42. const expiredThreadRelations = await this.find({ expiredAt: { $lte: currentDate } }).limit(limit ?? 100).exec();
  43. return expiredThreadRelations;
  44. };
  45. schema.methods.updateThreadExpiration = async function(): Promise<void> {
  46. this.expiredAt = generateExpirationDate();
  47. await this.save();
  48. };
  49. export default getOrCreateModel<ThreadRelationDocument, ThreadRelationModel>('ThreadRelation', schema);