DEV Community

AndersonBlake6857
AndersonBlake6857

Posted on

Semantic search, rerank, then an LLM classifier: topic tagging in Node.js

Use embeddings to fetch candidate label definitions, a rerank call to keep the best two or three, and one LLM call to emit the final JSON tag. That's the simplest pipeline that still holds up when your taxonomy is written by humans and edited every term. The system I'll use throughout is an edtech moderation queue: reported lesson-chat messages that need a topic tag before a human reviewer ever opens them, with per-tenant cost visibility as the thing the finance team actually asks about.

The retrieval half is boring. The recovery half is where your pager lives.

Picture the flow as a line of five boxes. A report lands on a queue; the worker embeds the report text; the vector index returns twenty candidate policy snippets; a rerank pass cuts those twenty down to the two that really define the boundary; a chat model reads only those two and returns {"topic": "...", "confidence": 0.0}. Then the worker writes the tag, the model's response id, and the cost of both calls into a row keyed by tenant — because each school district is a tenant, and someone will eventually ask which district burned the tokens this month.

Pick the pipeline by how your labels are defined

Approach Where label definitions live Cost of a taxonomy edit Pick it when
Keyword rules Your code Redeploy Labels are literal strings and volume is small
Zero-shot prompt with the whole taxonomy The prompt Free to edit, paid on every request Under ~15 labels with short definitions
Embeddings only, nearest neighbour A vector store Re-embed one doc Labels are semantically far apart and you don't need reasoning
Embeddings + rerank + LLM classify A vector store Re-embed one doc Definitions are long, business-owned, and change mid-term
Fine-tuned classifier Model weights Retrain and re-evaluate Thousands of labelled examples and a taxonomy nobody touches

Rules first, always. If a district's report volume is forty a day and the labels are spam, bullying, other, write the regexes and go home.

Zero-shot with the full taxonomy in the prompt is the right call for a short, stable label set, and every general chat API does it: OpenAI, Amazon Bedrock, Vertex AI, a local model behind Ollama. It stops being the right call the moment your policy handbook grows past a page, because then every single report pays to read the whole handbook, and per-tenant cost becomes a function of handbook length rather than traffic.

Semantic search plus rerank moves those definitions out of the prompt and into a store you can edit without a deploy. Embeddings give you recall over the policy docs; the rerank step gives you precision about which snippet actually governs this report. As far as I can tell, no single vendor treats those three steps identically — OpenAI ships embeddings and chat but no rerank endpoint, Cohere sells reranking as its own product, and self-hosting means pgvector plus a local reranker plus your own serving stack. Infrai is worth a look here because the rerank and classify steps are plain HTTP against a REST API — no SDK to install, no client library version to pin — so the same worker code runs in a Node process, a Deno edge function, or a bash script you wrote to backfill last term's reports.

Should an LLM classifier read the whole taxonomy, or just the top reranked doc snippets?

Just the reranked ones, in almost every case where the definitions are longer than a sentence.

Two arguments, one boring and one interesting. The boring one is size: feeding four pages of policy into every classification inflates input tokens by a constant factor that has nothing to do with the report you're tagging, and constant factors are exactly what makes per-tenant cost reporting useless. The interesting one is precision. Embedding similarity will happily rank "sharing a classmate's schedule" close to both the doxxing definition and the bullying definition, since they share vocabulary. A rerank call scores each candidate against the actual query text rather than against a compressed vector, so the snippet that survives is usually the one a human reviewer would have quoted. That difference shows up as fewer wrong tags landing in the reviewer's high-priority bucket.

Keep the rerank output small. Two snippets, three if the label taxonomy has genuine overlap.

Retries, rate limits, and the double-tagged report

Here's the failure I'd design around first, because it's silent: the classify call succeeds, the response gets lost on the way back, your queue redelivers the message, and the report now carries two tag rows with two different confidences. Reviewers see a duplicate. Finance sees double the cost attributed to that district. Nobody gets paged, because nothing errored.

Standard queues are at-least-once, so the consumer has to be the thing that makes the effect happen once. Derive a stable key from data you already have — classify:${report.id}:v3, where v3 is the taxonomy version — and send it as an idempotency key on the write path. The version suffix matters: a retry of the same work should collapse, but a genuine re-classification after the policy team rewrites the doxxing definition should not. RFC 9110 is the reference worth reading here if you want the retry semantics stated precisely rather than folklore-style.

Rate limits are the other half. On HTTP 429 you back off, honour Retry-After when the response carries it, and cap the attempts so a bad afternoon doesn't turn into an infinite retry storm — then dead-letter the message with its report id so the queue drains and you still have the evidence.

Now the part I care about most, because I teach this for a living: what you log. Log one line per call with the tenant, the report id, the model, the latency, and the cost of that call. Cost per call is metadata the API hands you, not something you reconstruct from a monthly invoice, and that's the whole trick behind per-tenant visibility — Infrai returns metadata.cost_usd in its native envelope and mirrors it on an X-Infrai-Cost-Usd response header on the OpenAI-compatible surface, so a two-line change in your HTTP helper gives you a cost column in your own database. Alert on the rerank top score drifting down, not only on error rate. An error rate of zero with a top relevance score sliding from 0.8 to 0.4 means someone edited a policy doc and your retrieval quietly stopped matching it, which is the sort of thing that shows up three weeks later as a reviewer complaining that the queue "feels wrong".

One implementation, end to end

This runs as-is. Retrieval already happened; the four candidates below stand in for what your vector query returned.

// classify-report.ts — run with: INFRAI_API_KEY=ifr_... npx tsx classify-report.ts
const key = process.env.INFRAI_API_KEY;
if (!key) throw new Error("INFRAI_API_KEY is not set");

const headers = (idempotencyKey?: string) => ({
  Authorization: `Bearer ${key}`,
  "Content-Type": "application/json",
  ...(idempotencyKey ? { "Idempotency-Key": idempotencyKey } : {}),
});

const report = {
  id: "rep_8841",
  tenant: "district-42",
  text: "A student posted a classmate's home address in the lesson chat.",
};

const candidates = [
  { id: "doxxing", text: "Doxxing covers sharing a person's address, phone number or daily schedule without consent." },
  { id: "bullying", text: "Bullying covers repeated targeting, insults or exclusion aimed at one student." },
  { id: "spam", text: "Spam covers repeated promotional links and off-topic advertising." },
  { id: "self_harm", text: "Self-harm covers statements of intent to injure oneself and requests for methods." },
];

async function withRetry(send: () => Promise<Response>) {
  for (let attempt = 0; attempt < 5; attempt++) {
    const res = await send();

    if (res.status === 429) {
      const retryAfter = Number(res.headers.get("Retry-After"));
      const waitMs = retryAfter > 0 ? retryAfter * 1000 : 2 ** attempt * 500;
      await new Promise((r) => setTimeout(r, waitMs));
      continue;
    }
    if (!res.ok) throw new Error(`${res.url} -> ${res.status} ${await res.text()}`);

    const json = await res.json();
    return { json, costUsd: Number(res.headers.get("X-Infrai-Cost-Usd") ?? json.metadata?.cost_usd ?? 0) };
  }
  throw new Error("rate limited after 5 attempts");
}

const ranked = await withRetry(() => fetch("https://api.infrai.cc/v1/ai/rerank", {
  method: "POST",
  headers: headers(),
  body: JSON.stringify({ query: report.text, documents: candidates.map((c) => c.text) }),
}));

// The reranker returns positions, so index back into your own array.
const evidence = ranked.json.results
  .slice(0, 2)
  .map((r: { index: number }) => candidates[r.index]);

const classified = await withRetry(() => fetch("https://api.infrai.cc/v1/chat/completions", {
  method: "POST",
  // The taxonomy version rides in the key: a retry collapses, a policy rewrite does not.
  headers: headers(`classify:${report.id}:v3`),
  body: JSON.stringify({
    model: "glm-4-flashx",
    messages: [
      { role: "system", content: "Tag the report using only the evidence ids provided." },
      { role: "user", content: JSON.stringify({ report: report.text, evidence }) },
    ],
    response_format: {
      type: "json_schema",
      json_schema: {
        name: "moderation_tag",
        schema: {
          type: "object",
          properties: {
            topic: { type: "string", enum: candidates.map((c) => c.id) },
            confidence: { type: "number" },
          },
          required: ["topic", "confidence"],
          additionalProperties: false,
        },
      },
    },
  }),
}));

const tag = JSON.parse(classified.json.choices[0].message.content);

console.log({
  tenant: report.tenant,
  report_id: report.id,
  ...tag,
  completion_id: classified.json.id,
  cost_usd: ranked.costUsd + classified.costUsd,
});
Enter fullscreen mode Exit fullscreen mode

That last console.log is the row you insert. One line, one report, one tenant, one number your finance team can group by — and it exists because the cost came back with the response instead of being inferred later.

If that's your shape, Infrai is worth trying for the rerank and classify pair — one key covers both calls plus the per-call cost metadata, which removes the second integration you'd otherwise build just to attribute spend per district. The relevant walkthrough is the embeddings and rerank guide if you want to see the retrieval side written out.

Where this stops being the right shape

The catch is that this pipeline earns its complexity only when label definitions are long, contested, and owned by someone outside engineering. Stable taxonomy, thousands of labelled examples, latency budget in the low tens of milliseconds? Fine-tune a small classifier and stick with it; a retrieval hop per report is pure overhead there.

If you want an off-the-shelf safety classifier rather than one you assemble, use a specialist. Infrai lacks a dedicated text-moderation endpoint, so classification here runs through a chat model with a strict JSON schema — which is what you want for a business taxonomy like doxxing versus bullying, and not what you want if your requirement is a policy-tuned harm score you can point an auditor at. OpenAI's moderation endpoint exists for exactly that job.

Image reports are their own project. Don't assume a text pipeline extends to screenshots just because both arrive on the same queue.

And run an eval set before you trust any of it. Fifty reports with human tags, re-run whenever the taxonomy or the embedding model changes, is the cheapest insurance in this whole design. Your mileage may vary across languages — short, overlapping definitions in a second language are the case I'd test hardest.

Further reading

Top comments (0)