DEV Community

Devanshu Biswas
Devanshu Biswas

Posted on

Dynamic Few-Shot Prompting: Retrieve the k Nearest Exemplars with Cosine kNN

#rag #ai #llm #promptengineering

Static few-shot prompting pastes the same handful of examples into every prompt. That one set has to serve every possible input at once, so it is a permanent compromise: for any particular request most of its examples are off-topic filler that spends tokens and dilutes the pattern, and any kind of input it didn't anticipate gets no relevant demonstration at all. Dynamic few-shot flips that. You keep a pool of labeled exemplars and, per input, retrieve the k most similar and paste only those. The examples become relevant to this request, so a small k beats a big static block on both quality and tokens.

It's kNN in embedding space

The machinery is simple: embed every exemplar once, embed the query, score each exemplar by cosine similarity, take the top-k, and assemble the prompt on the fly. No training — just distance.

Cosine measures the angle between two vectors, ignoring magnitude, so a short query and a longer exemplar still score high when they point the same way (share the important, rare terms). 1.0 is identical direction; 0 means nothing shared.

function cosine(a, b) {
  let dot = 0, na = 0, nb = 0;
  for (const k in a) { na += a[k] * a[k]; if (k in b) dot += a[k] * b[k]; }
  for (const k in b) { nb += b[k] * b[k]; }
  return (na && nb) ? dot / (Math.sqrt(na) * Math.sqrt(nb)) : 0;
}
Enter fullscreen mode Exit fullscreen mode

The selection itself scores every exemplar against the query, sorts descending, and keeps the k highest:

function selectExemplars(query, pool, vectors, idf, dflt, k) {
  const qv = vectorize(tokenize(query), idf, dflt);
  return pool
    .map((e, i) => ({ exemplar: e, sim: cosine(qv, vectors[i]) }))
    .sort((a, b) => b.sim - a.sim)      // nearest first
    .slice(0, k);                       // top-k
}
Enter fullscreen mode Exit fullscreen mode

The demo embeds text with a lightweight TF-IDF bag-of-words (rare terms weigh more) so it runs deterministically in the browser with no API. In production you swap in a dense embedding model — which also catches synonyms like "cancel" ≈ "unsubscribe" — but the shape is identical: embed once, score by cosine, take top-k, assemble on the fly.

Assemble only the picks

The prompt is built from the winners: a task header, each selected exemplar as an input→label pair, then the query with its label left blank. Only the examples change per request.

function assemblePrompt(picks, query) {
  const header = "Classify the support message into exactly one label:\n" +
    "billing | shipping | technical | account | cancellation | feedback";
  const shots = picks
    .map(p => `Message: "${p.exemplar.text}"\nLabel: ${p.exemplar.label}`)
    .join("\n\n");
  return `${header}\n\nExamples:\n${shots}\n\nMessage: "${query}"\nLabel:`;
}
Enter fullscreen mode Exit fullscreen mode

Watch it on a cancellation query: a fixed static set of billing / shipping / technical / account offers nothing on point, while the retriever surfaces the cancellation exemplars at the top.

Two things to get right

The k tradeoff. k=1 is brittle — the whole prompt rides on one exemplar. Too large reaches past the genuinely similar examples into off-topic ones, re-importing the static problem and paying more tokens and latency. The sweet spot is usually small, 3 to 8: each extra shot must earn its place.

Diversity via MMR. Pure top-k can return k near-duplicates. Maximal Marginal Relevance re-ranks to keep picks relevant to the query and different from each other, penalising redundancy with a tunable λ:

const score = lambda * rel - (1 - lambda) * red;  // λ=1 pure relevance; lower = more diverse
Enter fullscreen mode Exit fullscreen mode

Dynamic few-shot is the adaptive upgrade of static few-shot and the example-retrieval cousin of RAG — RAG retrieves facts to ground content; this retrieves examples to shape behaviour. The pool is the model's knowledge, so curate for coverage, correct labels, and no near-duplicates, and grow it from real logged inputs. At scale, precompute the pool vectors once and store them in an ANN index so retrieval stays O(log n).

Type your own query and watch the retriever rank all 15 exemplars live, beside a static block that can't adapt: https://dev48v.infy.uk/prompt/day58-dynamic-few-shot.html

Top comments (0)