refactor: optimize AI moderation pipeline, fix OOM risks, token duplication, and add Zod validation

This commit is contained in:
MythEclipse
2026-05-25 23:23:12 +07:00
parent cc2ee84c3b
commit cc61e2576b
5 changed files with 198 additions and 176 deletions
+14 -7
View File
@@ -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;
+30 -12
View File
@@ -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<void> {
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;
}
+29 -34
View File
@@ -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;
}
+121 -121
View File
@@ -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<string>();
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) ?? "<empty>" }
? {
error: lastParseError,
preview: lastInvalidContent?.slice(0, 800) ?? "<empty>",
}
: 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,
+4 -2
View File
@@ -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<MessageRecord[]> {
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) {