[Community][fix] Fix Azure cosmos db no SQL similarity search with score and mmr

[wassim-mechergui-shift]

Issues:

• similarity search with score not working for VectorStoreRetriever instanciated from AzureCosmosDBNoSqlVectorSearch because it requires an implementation of __select_relevance_score_fn_.
• mmr retrieval is not working because the fetching of metadata field from item to document metadata is currently bugged : (issues: Azure Cosmos DB _similarity_search_with_score fails with KeyError: 'metadata' #26097, [Bug] Azure cosmos db no sql vector store similarity search method "mmr" #28476)

Changes:

• VectorStoreRetriever fix:
  • For the similarity search from VectorStoreRetriever a rather typical implementation was added (example for Qdrant)
• with_embeddings to similarity search:
  • Fix for the with_embedding keyword) due to a change made here (Azure Cosmos DB _similarity_search_with_score fails with KeyError: 'metadata' #26097), item no longer has self._embedding_key and has "$1" (that represents it's order in the parameters). Since Cosmos DB’s SQL API does not support parameterized aliases or dynamic aliasing within the query, we chose to name it 'embeddingKey'
  • optimization : if with_embedding is not True, we don't need to query the embeddings. Querying the embeddings is heavy and makes requests much slower. This was changed
• Changed initialisation of vector policy (from which we get the embedding key and the distance function)
  • if user specifies a dict, we check it has the correct form
  • if user specifies the container creation if it doesn't exist option, we check this is not null
  • for the final user vector policy property
    • if the container exposes one, we use that one (and raise a warning if the user specified policy has differences)
    • if the container doesn't expose one, we use the user exposed one (with an error in case nothing was specified)

These changes were tested locally.

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

1 Skipped Deployment

Name	Status	Preview	Comments	Updated (UTC)
langchain	⬜️ Ignored (Inspect)	Visit Preview		Dec 9, 2024 0:27am

[wassim-mechergui-shift]

values are fetched from here : https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/vector-search#container-vector-policies

[baskaryan]

is this guaranteed to always be specified? seems like we only check that self._vector_embedding_policy isn't empty if self._create_container is True

[wassim-mechergui-shift]

I'm not really sure on what's the intended logic of this code, when I wrote it I followed the embedding_key up above.
But you are right. vector_embedding_policy in this case is a dict. Nothing enforces what this dict needs to be if create_container is not specified. I think there's a minor issue in the code here that can be addressed.

In essence, if we want to use an azure cosmos db no sql vector store, it must have the vector search feature enabled.
To enable and specify it, when we create the container, we must specify (among other things) a Vector Policy (the dict in question). microsoft link
Once created, it's possible to query for these info in the properties of the container. (through container.read()) (vector search may be activated but not exposed, see comment down below)

Additionally, once created, _database.create_container_if_not_exists (the method we use here to create our connection to the container) will not be overriden by the new specified policy (and it's totally valid also in this case to give it None policy).

So here are my suggested changes:

• make vector_embedding_policy optional since we only force the checks if create container is true
• self._vector_embedding_policy shouldn't be acquired from vector_embedding_policy but from the container itself self._vector_embedding_policy should be acquired from container if it is exposed otherwise from user specified options
• optionally add a warning if there's a mismatch between container exposed vector_embedding_policy and used specified vector_embedding_policy

LMK if this is okay with you ! (I added these changes)

[wassim-mechergui-shift]

Actually I was testing and this assumption is not always valid

Once created, it's possible to query for these info in the properties of the container. (through container.read())

It seems there is a bug in azure cosmos db no sql but if the container was the result of a partition_key migration then it seems the information on the vector policy is lost (so not possible to read through container.read()). However, the container still fully supports vector search !

I'll readapt the implementation to account for this case. It seems no longer possible to use container.read() as the one source of truth.

I'm sorry for the long post I wanted to be thorooughly descriptive and argumentative in my implementation choices!

[wassim-mechergui-shift]

this is a minor optimization but its impact can be substantial: embeddings needed only when the with_embedding argument is specified as True
Removing it from default because it adds a heavier weight on the query

When we ran our tests, even with chunk of size 4K tokens, the retrieval is still so much faster when we don't query for the embeddings (we go from 0.3s/chunk on a test to retrieve 100 chunks to ~0.02s/chunk)