From 7ef86c81ca8a5e955a3528e312b1320f3b3a6ff0 Mon Sep 17 00:00:00 2001 From: asepharyana Date: Tue, 1 Sep 2026 18:01:47 +0700 Subject: [PATCH] feat(ai): audit + harden embedding pipeline - Normalize text before embedding (strip mentions/URLs/emoji/markdown/control chars, lowercase, truncate) on both write and query sides so vectors aren't diluted and tokens aren't wasted - embeddingClient: retry embeddings (maxRetries 2), validate batch dimension consistency, preserve index alignment for empty-normalized texts - archiveEmbedder: store normalized text in archive payload, skip empty-normalized content - backend: normalize search queries, make archive search similarity threshold configurable (AI_LLM_EMBEDDING_ARCHIVE_MIN_SIMILARITY, default 0.6) --- .../backend/src/modules/messages/embed.ts | 24 +++++- .../src/modules/messages/messages.service.ts | 7 +- services/backend/src/shared/config/index.ts | 8 ++ .../modules/ai-moderation/embeddingClient.ts | 82 ++++++++++++++++++- .../message-capture/archiveEmbedder.ts | 14 +++- 5 files changed, 127 insertions(+), 8 deletions(-) diff --git a/services/backend/src/modules/messages/embed.ts b/services/backend/src/modules/messages/embed.ts index 93a6471c..b8cbc003 100644 --- a/services/backend/src/modules/messages/embed.ts +++ b/services/backend/src/modules/messages/embed.ts @@ -3,6 +3,26 @@ import { createChildLogger } from "@/shared/logger/index"; const logger = createChildLogger("messages-embed"); +/** Max chars for a search query fed to the embedding model. */ +const MAX_QUERY_CHARS = 300; + +/** + * Normalize a user search query before embedding so it lands in the same + * vector space as the archived content (which is normalized the same way on + * write). Mirrors the gateway's normalizer: strip control/zero-width chars, + * lowercase, collapse whitespace, cap length. Readable punctuation is kept — + * a search query is already compact. + */ +export function normalizeEmbeddingQuery(raw: string): string { + if (!raw) return ""; + return raw + .replace(/[\p{Cc}\p{Cf}]/gu, " ") + .toLowerCase() + .replace(/\s+/g, " ") + .trim() + .slice(0, MAX_QUERY_CHARS); +} + /** * Embed a search query with the configured OpenAI-compatible embedding model. * Uses raw fetch (the backend has no openai SDK dependency) and returns null @@ -10,8 +30,10 @@ const logger = createChildLogger("messages-embed"); * * encoding_format: "float" is REQUIRED — Nvidia-backed models reject base64. */ -export async function embedQuery(text: string): Promise { +export async function embedQuery(rawQuery: string): Promise { if (!config.AI_LLM_API_KEY || !config.AI_LLM_EMBEDDING_MODEL) return null; + const text = normalizeEmbeddingQuery(rawQuery); + if (!text) return null; try { const res = await fetch(`${config.AI_LLM_BASE_URL}/embeddings`, { method: "POST", diff --git a/services/backend/src/modules/messages/messages.service.ts b/services/backend/src/modules/messages/messages.service.ts index e242a8c8..11cf5309 100644 --- a/services/backend/src/modules/messages/messages.service.ts +++ b/services/backend/src/modules/messages/messages.service.ts @@ -1,3 +1,4 @@ +import { config } from "@/shared/config/index"; import { NotFoundError, ValidationError } from "@/shared/errors/index"; import { createChildLogger } from "@/shared/logger/index"; import { embedQuery } from "./embed.js"; @@ -100,7 +101,11 @@ export class MessagesService { ); return { results: [], nextCursor: null }; } - const hits = await searchArchive(vector, input.limit, 0.6); + const hits = await searchArchive( + vector, + input.limit, + config.AI_LLM_EMBEDDING_ARCHIVE_MIN_SIMILARITY, + ); const results = hits.map((h) => mapSearchHit(h)); return { results, nextCursor: null }; } diff --git a/services/backend/src/shared/config/index.ts b/services/backend/src/shared/config/index.ts index 5254f7f8..4d2502ba 100644 --- a/services/backend/src/shared/config/index.ts +++ b/services/backend/src/shared/config/index.ts @@ -135,6 +135,14 @@ export const configSchema = z AI_LLM_MODEL: z.string().default("text"), AI_LLM_VISION_MODEL: z.string().optional(), AI_LLM_EMBEDDING_MODEL: z.string().optional(), + // Minimum cosine similarity for the public archive semantic search. Lower + // = more (noisier) results; raise it to tighten precision. Tuned for a 1B + // embedding model — re-tune if the model's dimensionality changes. + AI_LLM_EMBEDDING_ARCHIVE_MIN_SIMILARITY: z.coerce + .number() + .min(0) + .max(1) + .default(0.6), AI_LLM_MAX_CONCURRENT: z.coerce.number().int().positive().default(5), AI_LLM_IMAGE_MAX_DIMENSION: z.coerce .number() diff --git a/services/discord-gateway/src/modules/ai-moderation/embeddingClient.ts b/services/discord-gateway/src/modules/ai-moderation/embeddingClient.ts index 29373df6..314eeef4 100644 --- a/services/discord-gateway/src/modules/ai-moderation/embeddingClient.ts +++ b/services/discord-gateway/src/modules/ai-moderation/embeddingClient.ts @@ -15,9 +15,59 @@ import OpenAI from "openai"; import { createChildLogger } from "@/shared/logger/index"; import { config } from "../../shared/config/config.js"; +import { cleanContent } from "./textSignals.js"; const log = createChildLogger("embedding-client"); +// --------------------------------------------------------------------------- +// Text normalization (shared by moderation + archive embedding) +// --------------------------------------------------------------------------- + +/** + * Max characters fed to the embedding model for a single document. Embedding + * models have a hard token ceiling; embedding past it throws / wastes tokens. + * Content messages are truncated; search queries have their own (smaller) cap. + */ +export const EMBEDDING_MAX_CHARS = 1200; +export const EMBEDDING_MAX_QUERY_CHARS = 300; + +/** + * Normalize free-form Discord text before embedding. + * + * Raw messages are full of signal-hostile noise: @mentions, channels, custom + * emoji, URLs, markdown and control chars. Embedding that noise directly + * dilutes the vector (sentences that differ only in an @mention or a link + * land far apart) and inflates token cost. The same cleanup is applied to the + * user's search query so archived vectors and the query share one space. + * + * - `cleanContent` (from textSignals) strips URLs/@mentions/emoji/markdown and + * collapses whitespace — good for embeddings, not just term extraction. + * - Control characters / zero-width joiners are removed (Discord pastes these). + * - Lowercase is applied so "Discord" and "discord" embed identically (embed + * models are case-sensitive; this measurably improves near-duplicate recall). + * + * Falls back to the raw input if normalization empties a string (e.g. a + * message that was only a URL) so index alignment is preserved by callers. + */ +export function normalizeEmbeddingText(raw: string, maxChars: number): string { + if (!raw) return ""; + const cleaned = cleanContent( + raw.replace(/[\p{Cc}\p{Cf}]/gu, " ").toLowerCase(), + ); + if (!cleaned) return raw.slice(0, maxChars); // preserve original if stripped + return cleaned.slice(0, maxChars); +} + +/** Normalize a single embedded document/message. */ +export function normalizeEmbeddingContent(raw: string): string { + return normalizeEmbeddingText(raw, EMBEDDING_MAX_CHARS); +} + +/** Normalize a user-provided search query before embedding it. */ +export function normalizeEmbeddingQuery(raw: string): string { + return normalizeEmbeddingText(raw, EMBEDDING_MAX_QUERY_CHARS); +} + // --------------------------------------------------------------------------- // Client (lazy singleton — same base URL as the chat client) // --------------------------------------------------------------------------- @@ -30,7 +80,10 @@ function getClient(): OpenAI | null { openaiClient = new OpenAI({ apiKey: config.AI_LLM_API_KEY, baseURL: config.AI_LLM_BASE_URL, - maxRetries: 0, + // Embeddings are cheap and idempotent — a transient network blip should + // NOT silently disable the whole semantic cache for a batch. Let the SDK + // retry (2 retries, jittered) instead of failing open immediately. + maxRetries: 2, timeout: 60_000, }); } @@ -49,6 +102,11 @@ export function isEmbeddingEnabled(): boolean { /** * Embed a batch of texts with the configured model. * Returns null on any failure so callers can skip semantic lookup. + * + * Each input is normalized (noise stripped, lowercased, length-capped) before + * embedding — see normalizeEmbeddingText. Index alignment with `texts` is + * preserved: if a text normalizes to empty we embed the raw original so the + * caller's `embeddings[i] ↔ texts[i]` mapping never shifts. */ export async function embedTexts(texts: string[]): Promise { if (!isEmbeddingEnabled()) return null; @@ -57,15 +115,33 @@ export async function embedTexts(texts: string[]): Promise { const client = getClient(); if (!client) return null; + const normalized = texts.map((t) => normalizeEmbeddingContent(t)); + try { const response = await client.embeddings.create({ model: config.AI_LLM_EMBEDDING_MODEL as string, - input: texts, + input: normalized, // OpenAI SDK v6 defaults to base64; Nvidia-backed embedding models // (e.g. llama-nemotron-embed) reject it with 400. Always float. encoding_format: "float", }); - return response.data.map((item) => item.embedding); + // All vectors in one response must share the model's dimension. If they + // don't (shouldn't happen, but guards against a misconfigured/mismatched + // model), fail the batch rather than feed garbage to cosine + Qdrant. + const vectors = response.data.map((item) => item.embedding); + const firstLen = vectors[0]?.length ?? 0; + const consistent = vectors.every((v) => v.length === firstLen); + if (!consistent || firstLen === 0) { + log.error( + { + model: config.AI_LLM_EMBEDDING_MODEL, + dims: vectors.map((v) => v.length), + }, + "Embedding response had inconsistent/empty dimensions — treating as failure", + ); + return null; + } + return vectors; } catch (error) { log.warn( { error: error instanceof Error ? error.message : String(error) }, diff --git a/services/discord-gateway/src/modules/message-capture/archiveEmbedder.ts b/services/discord-gateway/src/modules/message-capture/archiveEmbedder.ts index e33875d4..460cd0bc 100644 --- a/services/discord-gateway/src/modules/message-capture/archiveEmbedder.ts +++ b/services/discord-gateway/src/modules/message-capture/archiveEmbedder.ts @@ -1,4 +1,7 @@ -import { embedText } from "@/modules/ai-moderation/embeddingClient"; +import { + embedText, + normalizeEmbeddingContent, +} from "@/modules/ai-moderation/embeddingClient"; import { ARCHIVE_COLLECTION, qdrantPointId, @@ -41,14 +44,19 @@ export function archiveMessageEmbedded(message: ArchiveMessage): void { void (async () => { try { - const vector = await embedText(text); + // Normalize once: the vector AND the stored payload both use the clean + // text so the public search returns readable content and the vector + // isn't diluted by @mentions/URLs/emoji (see normalizeEmbeddingContent). + const normalized = normalizeEmbeddingContent(text); + if (!normalized) return; // nothing meaningful left after cleanup + const vector = await embedText(normalized); if (!vector) return; const ok = await upsertQdrantPointV2( ARCHIVE_COLLECTION, qdrantPointId(`archive:${message.id}`), vector, { - text: text.slice(0, 4000), + text: normalized.slice(0, 4000), flags: "", analyzed_at: Date.now(), // 5-year persistent window (archive is NOT a TTL cache).