When gemini-embedding-2 Retrieval Feels 'Almost Right,' Check task_type First

The fix: split task_type between documents and queries

The correction is simple. Embed indexed documents with RETRIEVAL_DOCUMENT and search queries with RETRIEVAL_QUERY. Keep output_dimensionality identical on both sides.

from google import genai
from google.genai import types
import numpy as np
 
client = genai.Client()
DIM = 1536  # keep this identical for documents and queries
 
def embed_document(text: str) -> np.ndarray:
    resp = client.models.embed_content(
        model="gemini-embedding-2",
        contents=text,
        config=types.EmbedContentConfig(
            task_type="RETRIEVAL_DOCUMENT",
            output_dimensionality=DIM,
        ),
    )
    v = np.array(resp.embeddings[0].values, dtype=np.float32)
    return v / np.linalg.norm(v)  # re-normalize after truncating dimensions
 
def embed_query(text: str) -> np.ndarray:
    resp = client.models.embed_content(
        model="gemini-embedding-2",
        contents=text,
        config=types.EmbedContentConfig(
            task_type="RETRIEVAL_QUERY",
            output_dimensionality=DIM,
        ),
    )
    v = np.array(resp.embeddings[0].values, dtype=np.float32)
    return v / np.linalg.norm(v)
 
# Index side
docs = [
    "A seaside landscape washed in sunset light",
    "Mountain ridges and a sea of clouds at dusk",
    "A bright, crystal-clear tropical beach at noon",
]
index = np.vstack([embed_document(d) for d in docs])
 
# Query side
q = embed_query("sunset ocean")
scores = index @ q  # normalized, so dot product = cosine similarity
ranking = np.argsort(-scores)
for rank, i in enumerate(ranking, start=1):
    print(f"#{rank}: {docs[i]}  score={scores[i]:.3f}")

Three things matter here. Split task_type between documents and queries. If you truncate with output_dimensionality, re-run L2 normalization (truncation makes the norm drift from 1). And compare using the dot product of normalized vectors.

The reason re-normalization is needed is that gemini-embedding-2 supports Matryoshka representation learning, so you can drop the tail to reduce dimensions. A freshly truncated vector no longer has length 1, which breaks the assumption behind approximating cosine similarity with a dot product. This directly affects ranking. For the dimension-versus-storage tradeoff, designing storage reduction by trimming gemini-embedding output to 256 dimensions is a useful companion read.

What I measured (with a caveat on how to read it)

Let me be candid: the numbers below come from a small set of my own (help articles plus wallpaper tags, about 1,200 items total). This is not an official benchmark. How much task_type helps depends on the nature of your corpus, so treat this as a direction, not a promise.

In my environment, going from no task_type to "document = RETRIEVAL_DOCUMENT, query = RETRIEVAL_QUERY" raised top-5 recall from roughly 0.72 to about 0.86. The gain was largest for short queries. The fewer the words, like the two-word "sunset ocean," the bigger the improvement from stating intent.

What surprised me was the middle row: forcing both sides to RETRIEVAL_DOCUMENT barely helped. It is not enough to "just set task_type." It only works once documents and queries play separate roles.

The task_type options, and how to choose

gemini-embedding-2 offers task types beyond retrieval. Here are the ones I reach for most. Since the right task_type changes with the job, running everything through "RETRIEVAL_DOCUMENT by default" is a mistake in some cases.

Here is the line I draw when unsure. If the task is asymmetric (short query pulling long documents, natural language pulling code), split DOCUMENT and QUERY. If it is symmetric (similarity of two sentences, clustering), unify both sides on SEMANTIC_SIMILARITY or CLUSTERING. That single distinction removes most of the guesswork.

I group App Store user reviews with CLUSTERING. Same embedding backend, entirely different job, so of course task_type changes too. For the operational side of handling multilingual reviews, I wrote up building multilingual Google Play review replies with Gemini.

A migration trap: existing indexes sometimes need a rebuild

The easiest place to get burned is when you change a model or a setting. Changing task_type after the fact means re-embedding the document side as well. If RETRIEVAL_QUERY document vectors get mixed into an index built with RETRIEVAL_DOCUMENT, the same document lands in a different place.

The rule is simple. If you change the document-side task_type, dimension, or model, re-embed that whole index. If only the query side changed, leave the index alone. Forget this asymmetry and you enter the maze of "I fixed the query and nothing improved."

A full re-embed costs money, so you want to avoid doing it often. Around 1,200 items finishes in a few minutes, but at hundreds of thousands the calculus changes. For a sense of re-embedding cost and caching, the thinking in cutting Gemini API cost with context caching transfers well. Treat rebuilding an index as a design decision, not something you can do on a whim.

A/B test task_type against a small golden set

Because task_type's effect is corpus-dependent, I always confirm it on my own data. You do not need a heavy evaluation harness. Twenty to thirty pairs of "for this query, this document is correct" is plenty.

import numpy as np
 
# Hand-build 20-30 (query, correct document index) pairs
golden = [
    ("sunset ocean", 0),
    ("city lights at night", 5),
    ("snowy mountain morning", 8),
    # ... list real queries from your own app
]
 
def recall_at_k(embed_q, index, golden, k=5):
    hits = 0
    for query, gold_idx in golden:
        q = embed_q(query)
        topk = np.argsort(-(index @ q))[:k]
        if gold_idx in topk:
            hits += 1
    return hits / len(golden)
 
# index_doc is the matrix built with RETRIEVAL_DOCUMENT
# embed_query is the RETRIEVAL_QUERY version; embed_plain_normalized has no task_type
r_tuned = recall_at_k(embed_query, index_doc, golden)
r_plain = recall_at_k(embed_plain_normalized, index_doc, golden)
print(f"RETRIEVAL_QUERY: {r_tuned:.3f} / no task_type: {r_plain:.3f}")

Those twenty lines shift you from "it feels a bit better" to "recall@5 went from 0.72 to 0.86." As an indie developer it is tempting to postpone evaluation, but embedding-config changes are exactly where this small instrument pays off. Even with a managed layer like File Search, it is worth understanding how task_type is handled inside it via the Gemini File Search API, then confirming with your own queries.

Wrapping up

When search misses in that "almost right" way, check task_type before you touch normalization or chunking. Documents get RETRIEVAL_DOCUMENT, queries get RETRIEVAL_QUERY. That single alignment tightens the ranking more often than you would expect.

For your next step, pick just five live queries, pair each with its correct document, and measure recall@5. The moment a number appears, the next thing to fix becomes visible. Ever since I started carrying this little instrument, my uncertainty around embeddings dropped a great deal. Thank you for reading.