@backstageai/common/models/AiRequests.model.js

Common code for Backstage AI services

const mongoose = require("mongoose");
const parentRequestSchema = new mongoose.Schema(
  {
    userId: {
      type: mongoose.Schema.Types.ObjectId,
      ref: "User",
      required: true,
      index: true,
    },
    // You can store the overall prompt or other metadata here
    prompt: {
      type: String,
      default: "",
    },
    artistName: {
      type: String,
    },
    // Store the input data for the
    requestPayload: {
      type: mongoose.Schema.Types.Mixed, // Store JSON for flexible input data
      required: true,
    },
    status: {
      type: String,
      enum: ["pending", "queued", "processing", "completed", "failed"],
      default: "pending",
    },
    // Soft delete or any other fields you like
    deleted: {
      type: Boolean,
      default: false,
      index: true,
    },
    isAlbum: {
      type: Boolean,
      default: false,
    },
    songTitle: {
      type: String,
    },
    themeImage: {
      type: String,
    },
    themeId: {
      type: mongoose.Schema.Types.ObjectId,
      required: false,
      index: true,
    },
  },
  {
    timestamps: true, // manages createdAt/updatedAt
  }
);

// Exclude soft-deleted records by default
parentRequestSchema.pre(/^find/, function (next) {
  this.where({ deleted: false });
  next();
});

const ParentRequest =
  mongoose.models.ParentRequest ||
  mongoose.model("ParentRequest", parentRequestSchema, "parentRequests");

const aiRequestSchema = new mongoose.Schema(
  {
    parentRequestId: {
      type: mongoose.Schema.Types.ObjectId,
      ref: "ParentRequest", // or ref: "AIRequest" if you prefer
      default: null,
      index: true,
    },
    generationId: {
      type: String,
      index: true,
      required: true,
      default: "pending",
    },
    outputUrl: {
      type: String,
    },
    thumbnailUrl: {
      type: String,
    },
    userId: {
      type: mongoose.Schema.Types.ObjectId,
      ref: "User",
      required: true,
      index: true,
    },
    themeId: {
      type: mongoose.Schema.Types.ObjectId,
      required: false,
      index: true,
    },
    taskType: {
      type: String,
      enum: ["image", "video"],
      required: true,
      index: true,
    },
    generationType: {
      type: String,
      enum: ["image", "logo", "video", "avatar", "background-removal"],
      required: true,
      index: true,
      default: "image",
    },
    aiService: {
      type: String,
      enum: ["falAi", "replicateAi", "lumalabs", "comfyDeploy"],
      required: true,
    },
    requestPayload: {
      type: mongoose.Schema.Types.Mixed, // Can store JSON for flexible input data
      required: true,
    },
    artistName: {
      type: String,
    },
    response: {
      type: mongoose.Schema.Types.Mixed, // To store the AI service response
      default: null,
    },
    webhookResponse: {
      type: mongoose.Schema.Types.Mixed, // To store webhook responses, if applicable
      default: null,
    },
    status: {
      type: String,
      enum: ["pending", "queued", "processing", "completed", "failed"],
      default: "pending",
    },
    errorMessage: {
      type: String, // To store any error messages in case of failure
      default: null,
    },
    deleted: {
      type: Boolean,
      default: false, // Soft delete flag
      index: true,
    },
    notificationShown: {
      type: Boolean,
      default: false,
    },
    webhookCallReceived: {
      type: Boolean,
      default: false,
    },
    isAlbum: {
      type: Boolean,
      default: false,
    },
    songTitle: {
      type: String,
    },
    progress: {
      type: Number,
    },
    avatarId: {
      type: mongoose.Schema.Types.ObjectId,
      ref: "Avatar",
    },
    themeImage: {
      type: String,
    },
    excludeMyRequests: {
      type: Boolean,
      default: false,
    },
  },
  {
    timestamps: true, // Automatically manage `createdAt` and `updatedAt`
  }
);

// Ensure queries exclude soft-deleted records by default
aiRequestSchema.pre(/^find/, function (next) {
  this.where({ deleted: false });
  next();
});

const AIRequest =
  mongoose.models.AIRequest ||
  mongoose.model("AIRequest", aiRequestSchema, "aiRequests");

module.exports = {
  AIRequest,
  ParentRequest,
};