返回 AiToEarn
queue.service.ts
根目录 / project / aitoearn-backend / libs / aitoearn-queue / src / queue.service.ts
1 import type { Job, JobsOptions, Queue } from 'bullmq'
2 import type {
3 AgentTaskAnalysisData,
4 AiImageData,
5 DraftGenerationData,
6 EngagementReplyToCommentData,
7 EngagementTaskDistributionData,
8 MaterialGenerateData,
9 PostMediaTaskData,
10 PostPublishData,
11 } from './interfaces'
12 import { InjectQueue } from '@nestjs/bullmq'
13 import { Injectable } from '@nestjs/common'
14 import { QueueName } from './enums'
15 import { ContentGenerationTaskData } from './interfaces/content-generation-task.interface'
16 import { QueueConfig } from './queue.config'
17
18 export interface DraftGenerationQueueInfo {
19 position: number | null
20 waitingCount: number
21 }
22
23 /**
24 * 统一的队列服务
25 * 提供所有队列的操作方法
26 */
27 @Injectable()
28 export class QueueService {
29 private readonly defaultOptions: JobsOptions
30
31 constructor(
32 private config: QueueConfig,
33 @InjectQueue(QueueName.MaterialGenerate)
34 private materialGenerateQueue: Queue,
35 @InjectQueue(QueueName.PostPublish)
36 private postPublishQueue: Queue,
37 @InjectQueue(QueueName.PostMediaTask)
38 private postMediaTaskQueue: Queue,
39 @InjectQueue(QueueName.AiImageAsync)
40 private aiImageAsyncQueue: Queue,
41 @InjectQueue(QueueName.EngagementTaskDistribution)
42 private engagementTaskDistributionQueue: Queue,
43 @InjectQueue(QueueName.EngagementReplyToComment)
44 private engagementReplyToCommentQueue: Queue,
45 @InjectQueue(QueueName.DumpSocialMediaAvatar)
46 private dumpSocialMediaAvatarQueue: Queue,
47 @InjectQueue(QueueName.UpdatePublishedPost)
48 private updatePublishedPostQueue: Queue,
49 @InjectQueue(QueueName.ContentGenerationTask)
50 private contentGenerationTaskQueue: Queue,
51 @InjectQueue(QueueName.AgentTaskAnalysis)
52 private agentTaskAnalysisQueue: Queue,
53 @InjectQueue(QueueName.DraftGeneration)
54 private draftGenerationQueue: Queue,
55 @InjectQueue(QueueName.DraftGenerationLowPriority)
56 private draftGenerationLowPriorityQueue: Queue,
57 ) {
58 // 从配置中读取默认的 job options
59 this.defaultOptions = config.jobOptions || {
60 removeOnComplete: { age: 30, count: 1000 },
61 removeOnFail: { age: 60, count: 1000 },
62 }
63 }
64
65 /**
66 * 添加素材生成任务
67 */
68 async addMaterialGenerateJob(data: MaterialGenerateData, options?: JobsOptions) {
69 return await this.materialGenerateQueue.add('start', data, {
70 ...this.defaultOptions,
71 ...options,
72 })
73 }
74
75 /**
76 * 添加发布任务
77 */
78 async addPostPublishJob(data: PostPublishData, options?: JobsOptions) {
79 return await this.postPublishQueue.add('publish', data, {
80 ...this.defaultOptions,
81 jobId: data.jobId,
82 ...options,
83 })
84 }
85
86 /**
87 * 获取发布任务
88 */
89 async getPostPublishJob(jobId: string): Promise<Job<PostPublishData> | undefined> {
90 return await this.postPublishQueue.getJob(jobId)
91 }
92
93 /**
94 * 添加发布媒体任务
95 */
96 async addPostMediaTaskJob(data: PostMediaTaskData, options?: JobsOptions) {
97 return await this.postMediaTaskQueue.add('media', data, {
98 ...this.defaultOptions,
99 ...options,
100 })
101 }
102
103 /**
104 * 添加AI图片异步生成任务
105 */
106 async addAiImageAsyncJob(data: AiImageData, options?: JobsOptions) {
107 return await this.aiImageAsyncQueue.add('generate', data, {
108 ...this.defaultOptions,
109 jobId: data.logId,
110 ...options,
111 })
112 }
113
114 async isAiImageAsyncJobActive(aiLogId: string): Promise<boolean> {
115 if (await this.isJobActive(await this.aiImageAsyncQueue.getJob(aiLogId))) {
116 return true
117 }
118
119 const activeJobs = await this.aiImageAsyncQueue.getJobs(['active'], 0, -1, true)
120 return activeJobs.some(job => job?.data?.logId === aiLogId)
121 }
122
123 /**
124 * 添加互动任务分发任务
125 */
126 async addEngagementTaskDistributionJob(
127 data: EngagementTaskDistributionData,
128 options?: JobsOptions,
129 ) {
130 return await this.engagementTaskDistributionQueue.add('distribute', data, {
131 ...this.defaultOptions,
132 ...options,
133 })
134 }
135
136 async addUpdatePublishedPostJob(data: { taskId: string, updatedContentType: string }, options?: JobsOptions) {
137 return await this.updatePublishedPostQueue.add('update-published-post', data, {
138 ...this.defaultOptions,
139 ...options,
140 })
141 }
142
143 /**
144 * 添加评论回复任务
145 */
146 async addEngagementReplyToCommentJob(data: EngagementReplyToCommentData, options?: JobsOptions) {
147 return await this.engagementReplyToCommentQueue.add('reply', data, {
148 ...this.defaultOptions,
149 ...options,
150 })
151 }
152
153 async addDumpSocialMediaAvatarJob(data: { accountId: string }, options?: JobsOptions) {
154 return await this.dumpSocialMediaAvatarQueue.add('dump-social-avatar', data, {
155 ...this.defaultOptions,
156 ...options,
157 })
158 }
159
160 async addContentGenerationTaskJob(data: ContentGenerationTaskData, options?: JobsOptions) {
161 return await this.contentGenerationTaskQueue.add('generate', data, {
162 ...this.defaultOptions,
163 ...options,
164 })
165 }
166
167 /**
168 * 添加Agent任务分析任务
169 */
170 async addAgentTaskAnalysisJob(data: AgentTaskAnalysisData, options?: JobsOptions) {
171 return await this.agentTaskAnalysisQueue.add('analyze', data, {
172 ...this.defaultOptions,
173 jobId: data.taskId,
174 attempts: 3,
175 backoff: {
176 type: 'exponential',
177 delay: 5000,
178 },
179 ...options,
180 })
181 }
182
183 /**
184 * 添加 DraftGeneration 生成任务
185 */
186 async addDraftGenerationJob(data: DraftGenerationData, options?: JobsOptions) {
187 return await this.draftGenerationQueue.add('generate', data, {
188 ...this.defaultOptions,
189 ...options,
190 jobId: data.aiLogId,
191 attempts: 3,
192 backoff: {
193 type: 'exponential',
194 delay: 5000,
195 },
196 })
197 }
198
199 /**
200 * 添加低优先级 DraftGeneration 生成任务
201 */
202 async addLowPriorityDraftGenerationJob(data: DraftGenerationData, options?: JobsOptions) {
203 return await this.draftGenerationLowPriorityQueue.add('generate', data, {
204 ...this.defaultOptions,
205 ...options,
206 jobId: data.aiLogId,
207 attempts: 3,
208 backoff: {
209 type: 'exponential',
210 delay: 5000,
211 },
212 })
213 }
214
215 async getDraftGenerationQueueInfo(aiLogId: string): Promise<DraftGenerationQueueInfo | undefined> {
216 const normalQueueInfo = await this.getDraftGenerationQueueInfoFromQueue(this.draftGenerationQueue, aiLogId)
217 if (normalQueueInfo) {
218 return normalQueueInfo
219 }
220
221 return await this.getDraftGenerationQueueInfoFromQueue(this.draftGenerationLowPriorityQueue, aiLogId)
222 }
223
224 async isDraftGenerationJobActive(aiLogId: string): Promise<boolean> {
225 return await this.isJobActive(await this.draftGenerationQueue.getJob(aiLogId))
226 || await this.isJobActive(await this.draftGenerationLowPriorityQueue.getJob(aiLogId))
227 }
228
229 private async getDraftGenerationQueueInfoFromQueue(queue: Queue, jobId: string): Promise<DraftGenerationQueueInfo | undefined> {
230 const job = await queue.getJob(jobId)
231 if (!job) {
232 return undefined
233 }
234
235 const [state, waitingCount] = await Promise.all([
236 job.getState(),
237 queue.count(),
238 ])
239 let position: number | null = null
240
241 if (state === 'waiting' || state === 'prioritized' || state === 'delayed' || state === 'waiting-children') {
242 const jobs = await queue.getJobs(['prioritized', 'waiting', 'delayed', 'waiting-children'], 0, -1, true)
243 const index = jobs.findIndex(item => item?.id === jobId)
244 position = index >= 0 ? index + 1 : null
245 }
246
247 return { position, waitingCount }
248 }
249
250 private async isJobActive(job?: Job | undefined): Promise<boolean> {
251 return !!job && await job.getState() === 'active'
252 }
253 }
254
254 lines TYPESCRIPT