Error using Cosmos DB (MongoDB API vCore) Vector Search with LangChain: 'embedding' KeyError

Question

Hi everyone,

I'm implementing semantic search using LangChain and Azure Cosmos DB (MongoDB API vCore) with vector search (vector-diskann). I'm generating embeddings using Azure OpenAI with the text-embedding-3-large model. Everything works fine until I attempt to retrieve documents with retriever.invoke(question).

langchain==0.3.17

langchain-openai==0.3.3

langchain-core==0.3.33

langchain-community==0.3.16

langchain-chroma==0.2.1

langchain-mongodb==0.5.0

CONFIGURATION VECTOR STORE

from langchain_community.vectorstores.azure_cosmos_db import AzureCosmosDBVectorSearch

def get_vector_store_cosmos_mongodb(mongodb_collection):

return AzureCosmosDBVectorSearch(

    collection=mongodb_collection,

    index_name="vector_index",

    text_key="text",

    embedding_key="embedding",

    embedding=get_azure_embedding_function()

)

EMBEDDING AND DOCUMENT INSERTION

mongo_doc = {

"text": page_content,

"embedding": embedding,  # List[float]

"metadata": doc.metadata,

}

collection.insert_many(docs_to_insert)

**The index is created with the following configuration:
**
index_definition = {

"createIndexes": collection_name,

"indexes": [

    {

        "name": "vector_index",

        "key": {"embedding": "cosmosSearch"},

        "cosmosSearchOptions": {

            "kind": "vector-diskann",

            "dimensions": int(os.environ["NUMBER_DIMENSIONS"]),

            "similarity": "COS",

        },

    }

],

}

db.command(index_definition)

ERROR ON RETRIEVAL

When calling:

retriever = vector_store.as_retriever(

            search_type="mmr",

            search_kwargs={

                "k": 5,

                "fetch_k": 40,

                "lambda_mult": 0.7,

            },

        )

documents = retriever.invoke(question)

I get this error:

Error in general search: 'embedding'

Exception type: KeyError

Error message: 'embedding'

Arguments: ('embedding',)

Logs show that the DB and collection are connected and recognized:

Connected successfully to Cosmos DB (MongoDB API vCore).

Database 'prueba' already exists.

Collection 'prueba_V225' already exists in database 'prueba'.

Has anyone faced the 'embedding' KeyError when using AzureCosmosDBVectorSearch with LangChain?

Is there anything I might be missing in the collection configuration or document structure?

Could this be related to how the embedding vector is stored or indexed?

Any help or insight would be greatly appreciated 🙏

Answer 1

Hi Maria Natalia Benavides Sanabria

Thanks for the update and the confirmation that the embedding field exists in your documents and index.

The KeyError: 'embedding' is most likely caused because Cosmos DB’s $search stage doesn’t automatically include all fields in the search results — including embedding. This means that while the data is stored correctly, it’s not returned by the query, and LangChain fails when trying to access it.

✅ What you can do:

1. Test a raw search query (outside LangChain) with a $project stage that explicitly includes the embedding field. This should work.

2.If it works, then you'll need to either:

• Patch LangChain locally to include $project in its vector search query, or
• Use a custom retriever that performs the aggregation manually (we can share code if needed).

3.Long-Term Fixes / Suggestions

• Open a GitHub issue or PR with LangChain: This is a real gap in Cosmos integration — not automatically projecting required fields.

Alternatively, wrap AzureCosmosDBVectorSearch to override the internal _search method, injecting a custom projection.

This is a known limitation and something LangChain may fix in future updates.

Hope this helps. Do let us know if you any further queries.

If this answers your query, do click Accept Answer and Yes for was this answer helpful. And, if you have any further query do let us know.