From cc61e2576b6f47c6357949ae9060fc7b2fe1e53d Mon Sep 17 00:00:00 2001 From: MythEclipse Date: Mon, 25 May 2026 23:23:12 +0700 Subject: [PATCH] refactor: optimize AI moderation pipeline, fix OOM risks, token duplication, and add Zod validation --- src/moderation/aiAnalysisWorker.ts | 21 ++- src/moderation/aiAnalyzer.ts | 42 +++-- src/moderation/conversationContext.ts | 63 +++---- src/moderation/llmModerationClient.ts | 242 +++++++++++++------------- src/moderation/messageStore.ts | 6 +- 5 files changed, 198 insertions(+), 176 deletions(-) diff --git a/src/moderation/aiAnalysisWorker.ts b/src/moderation/aiAnalysisWorker.ts index 543ebc4..3e83561 100644 --- a/src/moderation/aiAnalysisWorker.ts +++ b/src/moderation/aiAnalysisWorker.ts @@ -1,6 +1,6 @@ import { config } from "../config.js"; import { initializeDatabase } from "../database/drizzle.js"; -import { buildConversationPromptMessages } from "./conversationContext.js"; +import { buildConversationContext } from "./conversationContext.js"; import { runModerationAnalysis } from "./llmModerationClient.js"; import { getAttachmentsForMessages, @@ -67,7 +67,7 @@ export default async function processAnalysisRequest({ limit: config.AI_ANALYSIS_CONTEXT_MESSAGE_LIMIT, }); - const promptMessages = buildConversationPromptMessages({ + const contextLines = buildConversationContext({ contextBefore, targets: messages, maxTokens: config.AI_ANALYSIS_MAX_CONTEXT_TOKENS, @@ -80,7 +80,7 @@ export default async function processAnalysisRequest({ const result = await runModerationAnalysis({ targets: messages, - contextText: promptMessages.join("\n"), + contextText: contextLines.join("\n"), attachments, }); @@ -94,12 +94,19 @@ export default async function processAnalysisRequest({ analysis: analysisResult.analysis, analyzedAt: Date.now(), error: null, - } + }, })); - - const rows = await updateMessagesAIAnalysisBulk(updates); - return { ok: true, conversationKey, rows }; + try { + const rows = await updateMessagesAIAnalysisBulk(updates); + return { ok: true, conversationKey, rows }; + } catch (dbErr) { + // If bulk update fails, we log it but don't fail the worker completely + // so it can at least retry later without blowing up the circuit breaker if it was an isolated issue + throw new Error( + `Failed to update DB: ${dbErr instanceof Error ? dbErr.message : String(dbErr)}`, + ); + } } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); const errorStack = error instanceof Error ? error.stack : undefined; diff --git a/src/moderation/aiAnalyzer.ts b/src/moderation/aiAnalyzer.ts index 74497ba..d45e78c 100644 --- a/src/moderation/aiAnalyzer.ts +++ b/src/moderation/aiAnalyzer.ts @@ -3,6 +3,10 @@ import { fileURLToPath } from "node:url"; import { Piscina } from "piscina"; import { config } from "../config.js"; import { createChildLogger } from "../logger.js"; +import { + estimateTokens, + formatMessageForPrompt, +} from "./conversationContext.js"; import { getMessageById, getPendingConversationKeys, @@ -88,10 +92,8 @@ export function pickBatchWithinBudget( let usedTokens = 0; for (const msg of messages) { - // Estimate tokens based on actual content length (conservative: 3 chars/token) - const content = msg.edited_content ?? msg.content; - const contentTokens = Math.ceil(content.length / 3); - const msgTokens = contentTokens + tokensPerMessage; + const formatted = formatMessageForPrompt(msg, "target"); + const msgTokens = estimateTokens(formatted) + tokensPerMessage; if (usedTokens + msgTokens <= maxTokens) { batch.push(msg); @@ -118,7 +120,8 @@ async function processBatch( ): Promise { if (messages.length === 0) return; if (Date.now() < globalCooldownUntil) { - return; // Circuit breaker is open + // Should not normally hit here due to checks in scheduleConversationAnalysis, but just in case + return; } activeRequests++; @@ -126,7 +129,10 @@ async function processBatch( const processingStartedAt = Date.now(); conversationProcessing.set(conversationKey, processingStartedAt); try { - const result = (await workerPool.run({ conversationKey, messages })) as AnalysisWorkerResponse; + const result = (await workerPool.run({ + conversationKey, + messages, + })) as AnalysisWorkerResponse; for (const row of result.rows) { getModerationBroadcaster()?.messageAnalyzed(row); @@ -136,9 +142,11 @@ async function processBatch( consecutiveErrors++; if (consecutiveErrors >= MAX_CONSECUTIVE_ERRORS) { globalCooldownUntil = Date.now() + 60000; - logger.warn("Global circuit breaker triggered due to consecutive errors"); + logger.warn( + "Global circuit breaker triggered due to consecutive errors", + ); } - + lastError = result.error ?? "Analysis worker failed"; conversationErrorCooldown.set( conversationKey, @@ -201,7 +209,6 @@ async function processBatch( } } - /** * Debounced analysis trigger for a conversation */ @@ -211,9 +218,20 @@ function scheduleConversationAnalysis(conversationKey: string): void { return; } - // Skip if in error cooldown - const cooldownUntil = conversationErrorCooldown.get(conversationKey); - if (cooldownUntil && Date.now() < cooldownUntil) { + // Check cooldowns + const convoCooldown = conversationErrorCooldown.get(conversationKey) || 0; + const activeCooldown = Math.max(convoCooldown, globalCooldownUntil); + + if (activeCooldown && Date.now() < activeCooldown) { + // Instead of dropping, re-schedule for after cooldown if not already scheduled + if (!conversationDebounceTimers.has(conversationKey)) { + const remaining = activeCooldown - Date.now(); + const timer = setTimeout(() => { + conversationDebounceTimers.delete(conversationKey); + scheduleConversationAnalysis(conversationKey); + }, remaining + 500); // 500ms buffer after cooldown + conversationDebounceTimers.set(conversationKey, timer); + } return; } diff --git a/src/moderation/conversationContext.ts b/src/moderation/conversationContext.ts index fdb4ade..394ea3d 100644 --- a/src/moderation/conversationContext.ts +++ b/src/moderation/conversationContext.ts @@ -14,57 +14,52 @@ function formatTimestamp(ms: number): string { } /** - * Estimates token count for a string (rough approximation: ~4 chars per token) + * Estimates token count for a string (pessimistic approximation for Indonesian slang & JSON overhead) */ -function estimateTokens(text: string): number { - return Math.ceil(text.length / 4); +export function estimateTokens(text: string): number { + return Math.ceil(text.length / 3) + 15; } /** - * Builds conversation prompt messages with context and targets - * - Marks target messages with [target], prior context with [context] - * - Uses edited_content when present, otherwise content - * - Maintains chronological order - * - Respects maxTokens budget, prioritizing targets and most recent context + * Formats a single message for context or target display */ -export function buildConversationPromptMessages( +export function formatMessageForPrompt( + msg: MessageRecord, + label: "context" | "target", +): string { + const content = msg.edited_content ?? msg.content; + const timestamp = formatTimestamp(msg.created_at); + return `[${label}] id=${msg.id} time=${timestamp} user=${msg.username}: ${content}`; +} + +/** + * Builds conversation historical context without including targets. + * Calculates how much token budget targets use, and fills the rest with context. + */ +export function buildConversationContext( input: ConversationContextInput, ): string[] { const { contextBefore, targets, maxTokens } = input; - const formatMessage = (msg: MessageRecord, label: string): string => { - const content = msg.edited_content ?? msg.content; - const timestamp = formatTimestamp(msg.created_at); - return `[${label}] id=${msg.id} time=${timestamp} user=${msg.username}: ${content}`; - }; + // Calculate tokens used by targets + let usedTokens = targets.reduce((sum, msg) => { + return sum + estimateTokens(formatMessageForPrompt(msg, "target")); + }, 0); - const targetEntries = targets.map((msg) => ({ - msg, - label: "target" as const, - line: formatMessage(msg, "target"), - })); + const selectedContextLines: string[] = []; - let usedTokens = targetEntries.reduce( - (sum, entry) => sum + estimateTokens(entry.line), - 0, - ); - - const selectedContextEntries: Array<{ - msg: MessageRecord; - label: "context"; - line: string; - }> = []; + // Go backwards through context, taking most recent first for (let i = contextBefore.length - 1; i >= 0; i--) { const msg = contextBefore[i]; - const line = formatMessage(msg, "context"); + const line = formatMessageForPrompt(msg, "context"); const lineTokens = estimateTokens(line); + if (usedTokens + lineTokens <= maxTokens) { - selectedContextEntries.push({ msg, label: "context", line }); + // Unshift so oldest context is first in the array + selectedContextLines.unshift(line); usedTokens += lineTokens; } } - return [...selectedContextEntries, ...targetEntries] - .sort((a, b) => a.msg.created_at - b.msg.created_at) - .map((entry) => entry.line); + return selectedContextLines; } diff --git a/src/moderation/llmModerationClient.ts b/src/moderation/llmModerationClient.ts index d63eac3..a7a1020 100644 --- a/src/moderation/llmModerationClient.ts +++ b/src/moderation/llmModerationClient.ts @@ -1,4 +1,5 @@ import OpenAI from "openai"; +import { z } from "zod"; import { config } from "../config.js"; import { createChildLogger } from "../logger.js"; import { retryWithBackoff } from "../retry.js"; @@ -8,45 +9,66 @@ import type { MessageRecord, } from "./types.js"; +const ModerationResponseSchema = z.object({ + results: z.array( + z.object({ + message_id: z.union([z.string(), z.number()]).transform(String), + status: z.enum(["clean", "warn", "flagged"]).catch("clean"), + flags: z.array(z.string()).catch([]), + score: z.number().catch(0), + analysis: z.string().catch(""), + }), + ), +}); + const log = createChildLogger("llmModerationClient"); const openai = new OpenAI({ apiKey: config.AI_LLM_API_KEY, baseURL: config.AI_LLM_BASE_URL, maxRetries: 0, - timeout: 2_147_483_647, + timeout: 30000, fetch: async (url, init) => { - const response = await globalThis.fetch(url, init); - const body = - typeof response.text === "function" - ? await response.text() - : JSON.stringify(await response.json()); + // Add internal timeout for the global fetch as safety + const controller = new AbortController(); + const timeout = setTimeout(() => controller.abort(), 30000); + const fetchInit = { ...init, signal: controller.signal }; - let normalizedBody = body; - if (response.ok !== false) { - try { - JSON.parse(body); - } catch (error) { - log.warn( - { - error: error instanceof Error ? error.message : String(error), - status: response.status ?? 200, - bodyLength: body.length, - body, - }, - "LLM provider returned malformed JSON response body", - ); - normalizedBody = JSON.stringify(extractJson(body)); + try { + const response = await globalThis.fetch(url, fetchInit); + const body = + typeof response.text === "function" + ? await response.text() + : JSON.stringify(await response.json()); + + let normalizedBody = body; + if (response.ok !== false) { + try { + JSON.parse(body); + } catch (error) { + log.warn( + { + error: error instanceof Error ? error.message : String(error), + status: response.status ?? 200, + bodyLength: body.length, + body, + }, + "LLM provider returned malformed JSON response body", + ); + normalizedBody = JSON.stringify(extractJson(body)); + } } + + const headers = new Headers(response.headers ?? undefined); + headers.set("Content-Type", "application/json"); + headers.delete("Content-Length"); + + return new Response(normalizedBody, { + status: response.status ?? 200, + headers, + }); + } finally { + clearTimeout(timeout); } - - const headers = new Headers(response.headers ?? undefined); - headers.set("Content-Type", "application/json"); - headers.delete("Content-Length"); - - return new Response(normalizedBody, { - status: response.status ?? 200, - headers, - }); }, }); @@ -130,8 +152,6 @@ export function extractJson(content: string): any { throw new Error("No JSON object found in response"); } - - export function parseModerationResponse( content: string, targetIds: string[], @@ -156,66 +176,37 @@ export function parseModerationResponse( } } - if (!parsed || typeof parsed !== "object" || !Array.isArray(parsed.results)) { - throw new Error("Response missing 'results' array"); + const parseResult = ModerationResponseSchema.safeParse(parsed); + if (!parseResult.success) { + throw new Error(`Zod validation failed: ${parseResult.error.message}`); } - const response = parsed as RawModerationResponse; + const response = parseResult.data; const foundIds = new Set(); const targetIdSet = new Set(targetIds); - const results: (AnalysisResult | null)[] = response.results.map( - (result, index) => { - const { message_id, status, flags, score, analysis } = result; + const results: (AnalysisResult | null)[] = response.results.map((result) => { + const { message_id, status, flags, score, analysis } = result; + const finalId = message_id.trim(); - if (!message_id) { - throw new Error("Result missing 'message_id'"); - } + if (!targetIdSet.has(finalId)) { + return null; + } - const finalId = String(message_id).trim(); + if (foundIds.has(finalId)) { + return null; // Ignore duplicates safely + } - if (!targetIdSet.has(finalId)) { - log.warn( - { unknownId: finalId, originalId: message_id, targetIds }, - "Skipping moderation result for non-target message_id", - ); - return null; - } + foundIds.add(finalId); - if (foundIds.has(finalId)) { - log.warn({ duplicateId: finalId }, "Duplicate message_id in response"); - throw new Error(`Duplicate message_id: ${finalId}`); - } - - foundIds.add(finalId); - - const validStatuses = ["clean", "warn", "flagged"] as const; - const safeStatus = validStatuses.includes(status as any) ? status : "clean"; - - let numScore = Number(score); - if (!Number.isFinite(numScore)) { - numScore = 0; - } - numScore = Math.max(0, Math.min(1, numScore)); - - let flagsArray: string[] = []; - if (Array.isArray(flags)) { - flagsArray = flags.map((f) => String(f)); - } else if (flags) { - flagsArray = [String(flags)]; - } - - const analysisStr = analysis ? String(analysis) : ""; - - return { - messageId: finalId, - status: safeStatus as "clean" | "warn" | "flagged", - flags: flagsArray, - score: numScore, - analysis: analysisStr, - }; - }, - ); + return { + messageId: finalId, + status: status as "clean" | "warn" | "flagged", + flags, + score: Math.max(0, Math.min(1, score)), + analysis, + }; + }); const filteredResults = results.filter( (r): r is AnalysisResult => r !== null, @@ -362,7 +353,9 @@ export async function runModerationAnalysis( const targetIdSet = new Set(targets.map((t) => t.id)); const candidateAttachments = (attachments ?? []) - .filter((att) => getAttachmentImageUrl(att) && att.type.startsWith("image/")) + .filter( + (att) => getAttachmentImageUrl(att) && att.type.startsWith("image/"), + ) .sort((a, b) => { // Target-message attachments always come first so they consume the cap first const aIsTarget = targetIdSet.has(a.message_id) ? 1 : 0; @@ -380,12 +373,16 @@ export async function runModerationAnalysis( const urlToUse = getAttachmentImageUrl(att); if (!urlToUse) return; + const controller = new AbortController(); + const timeoutId = setTimeout(() => controller.abort(), 15000); + try { log.info( { attachmentId: att.id, messageId: att.message_id, url: urlToUse }, "Downloading attachment for base64 encoding", ); - const res = await fetch(urlToUse); + + const res = await fetch(urlToUse, { signal: controller.signal }); if (!res.ok) { log.warn( { attachmentId: att.id, status: res.status, url: urlToUse }, @@ -394,13 +391,31 @@ export async function runModerationAnalysis( return; } - const contentLength = Number(res.headers.get("content-length") || 0); - if (contentLength > 10 * 1024 * 1024) { - log.warn({ attachmentId: att.id, contentLength }, "Attachment too large, skipping"); - return; + if (!res.body) return; + + let totalBytes = 0; + const chunks: Uint8Array[] = []; + const reader = res.body.getReader(); + + while (true) { + const { done, value } = await reader.read(); + if (done) break; + + if (value) { + totalBytes += value.length; + if (totalBytes > 10 * 1024 * 1024) { + log.warn( + { attachmentId: att.id }, + "Attachment exceeded 10MB limit, aborting stream", + ); + reader.cancel(); + return; + } + chunks.push(value); + } } - const imageBytes = Buffer.from(await res.arrayBuffer()); + const imageBytes = Buffer.concat(chunks); const sniffedMime = sniffImageMimeType(imageBytes); if (!sniffedMime) { log.warn( @@ -417,7 +432,10 @@ export async function runModerationAnalysis( } const dataUrl = `data:${sniffedMime};base64,${imageBytes.toString("base64")}`; - const part: RawImagePart = { type: "image_url", image_url: { url: dataUrl } }; + const part: RawImagePart = { + type: "image_url", + image_url: { url: dataUrl }, + }; const existing = messageImageMap.get(att.message_id) ?? []; existing.push(part); @@ -430,6 +448,8 @@ export async function runModerationAnalysis( }, "Error base64 encoding attachment", ); + } finally { + clearTimeout(timeoutId); } }), ); @@ -524,7 +544,10 @@ CRITICAL: "message_id" HARUS berupa STRING (dibungkus tanda kutip ganda). Jangan const buildMessageContent = (): string | ContentPart[] => { const correction = lastParseError - ? { error: lastParseError, preview: lastInvalidContent?.slice(0, 800) ?? "" } + ? { + error: lastParseError, + preview: lastInvalidContent?.slice(0, 800) ?? "", + } : undefined; const systemText = buildSystemPrompt(correction); @@ -543,7 +566,10 @@ CRITICAL: "message_id" HARUS berupa STRING (dibungkus tanda kutip ganda). Jangan // Multimodal path: interleave text + images per message const parts: ContentPart[] = [ - { type: "text", text: `${systemText}\n\n## Pesan yang Dianalisis (dengan lampiran gambar)\n` }, + { + type: "text", + text: `${systemText}\n\n## Pesan yang Dianalisis (dengan lampiran gambar)\n`, + }, ]; for (const msg of targets) { @@ -586,33 +612,7 @@ CRITICAL: "message_id" HARUS berupa STRING (dibungkus tanda kutip ganda). Jangan top_p: 0.95, max_tokens: 16384, response_format: { - type: "json_schema", - json_schema: { - name: "moderation", - strict: true, - schema: { - type: "object", - properties: { - results: { - type: "array", - items: { - type: "object", - properties: { - message_id: { type: "string" }, - status: { type: "string", enum: ["clean", "warn", "flagged"] }, - flags: { type: "array", items: { type: "string" } }, - score: { type: "number" }, - analysis: { type: "string" } - }, - required: ["message_id", "status", "flags", "score", "analysis"], - additionalProperties: false - } - } - }, - required: ["results"], - additionalProperties: false - } - } + type: "json_object", }, stream: false, chat_template_kwargs: { enable_thinking: false }, @@ -674,7 +674,7 @@ CRITICAL: "message_id" HARUS berupa STRING (dibungkus tanda kutip ganda). Jangan const errorMsg = parseError instanceof Error ? parseError.message : String(parseError); const content: string = lastInvalidContent; - + log.error( { error: errorMsg, diff --git a/src/moderation/messageStore.ts b/src/moderation/messageStore.ts index 460c3fe..6776b78 100644 --- a/src/moderation/messageStore.ts +++ b/src/moderation/messageStore.ts @@ -420,12 +420,14 @@ export async function updateMessageAIAnalysis( } export async function updateMessagesAIAnalysisBulk( - updates: Array<{ messageId: string; result: AIAnalysisUpdate }> + updates: Array<{ messageId: string; result: AIAnalysisUpdate }>, ): Promise { if (updates.length === 0) return []; try { const results = await Promise.all( - updates.map(({ messageId, result }) => updateMessageAIAnalysis(messageId, result)) + updates.map(({ messageId, result }) => + updateMessageAIAnalysis(messageId, result), + ), ); return results.filter((r): r is MessageRecord => r !== null); } catch (error) {