Cache embeddings by content hash

[0x7c13]

Caching the embeddings by content hash of the partitioned file content to avoid/reduce potential cost of invoking the embedding API.

Motivation and Context (Why the change? What's the scenario?)

It is very costly (both time and money wise) to call and invoke the embedding API. So we want to reduce the chance of hitting it by caching the embeddings as much as possible.

High level description (Approach, Design)

Cache the embeddings of a given partitioned content by its content hash to determine if content embedding can be re-used.

[0x7c13]

Btw, this Spell Check action should be configured to scan the content added/updated by the PR only if possible.

[dluc]

This addresses the concern only when the handler runs on a single machine. In a production environment KM should be deployed on multiple VMs, for availability and performance, and for the cache to work as expected it should be centralized in a service like Redis or MemCache, or reusing some of the Content Storage API that leverages Azure Blobs for example.
Duplicate requests seems to be a problem only when parsing the same content several times, which I think is more likely a test scenario. It's fine adding such feature, but I think it will not make a huge difference given these considerations.
Also, the cache is in memory, and any benefit from reusing data is lost as soon as a handler is disposed or the service stops.

There are also a couple of bugs:

• KM allows to use multiple embedding generators, see the loop at line 113. The same content generates different embeddings on different generators, so the cache cannot be shared.
• Embedding generators can be configured to use different models, e.g. Ada002 or Ada003. While this is not a problem in the current implementation holding data in memory, if cache is stored on disk, the cache key should consider also which model is used, considering that embeddings are different in content and length.

[0x7c13]

This addresses the concern only when the handler runs on a single machine. In a production environment KM should be deployed on multiple VMs, for availability and performance, and for the cache to work as expected it should be centralized in a service like Redis or MemCache, or reusing some of the Content Storage API that leverages Azure Blobs for example. Duplicate requests seems to be a problem only when parsing the same content several times, which I think is more likely a test scenario. It's fine adding such feature, but I think it will not make a huge difference given these considerations. Also, the cache is in memory, and any benefit from reusing data is lost as soon as a handler is disposed or the service stops.

There are also a couple of bugs:

• KM allows to use multiple embedding generators, see the loop at line 113. The same content generates different embeddings on different generators, so the cache cannot be shared.
• Embedding generators can be configured to use different models, e.g. Ada002 or Ada003. While this is not a problem in the current implementation holding data in memory, if cache is stored on disk, the cache key should consider also which model is used, considering that embeddings are different in content and length.

Yes, I agree with you, that's why I posted a feature request under issues: #390

[dluc]

Thought a bit after posting in the feature request yesterday. I think a good solution would be developing an interface IEmbeddingCache, injecting it into each embedding generator via DI, so that any client would benefit.

Something like:

public interface IEmbeddingCache {

    public async WriteAsync(....cache key, embedding ...);

    public async Task<Embedding?> ReasAsync(....cache key....);
}

public class OpenAITextEmbeddingGenerator : ITextEmbeddingGenerator
{
    public OpenAITextEmbeddingGenerator(
            OpenAIConfig config,
            OpenAIClient openAIClient,
            ITextTokenizer? textTokenizer = null,
            IEmbeddingCache? embeddingCache = null,
            ILoggerFactory? loggerFactory = null)
        {
            // ....

Then each generator will know how to calculate the cache key properly, without having to leak that logic outside.

[0x7c13]

This addresses the concern only when the handler runs on a single machine. In a production environment KM should be deployed on multiple VMs, for availability and performance, and for the cache to work as expected it should be centralized in a service like Redis or MemCache, or reusing some of the Content Storage API that leverages Azure Blobs for example. Duplicate requests seems to be a problem only when parsing the same content several times, which I think is more likely a test scenario. It's fine adding such feature, but I think it will not make a huge difference given these considerations. Also, the cache is in memory, and any benefit from reusing data is lost as soon as a handler is disposed or the service stops.

There are also a couple of bugs:

• KM allows to use multiple embedding generators, see the loop at line 113. The same content generates different embeddings on different generators, so the cache cannot be shared.
• Embedding generators can be configured to use different models, e.g. Ada002 or Ada003. While this is not a problem in the current implementation holding data in memory, if cache is stored on disk, the cache key should consider also which model is used, considering that embeddings are different in content and length.

What about re-using the embeddings/vectors used by the search vs cache? We could store two copies of the vectors, one for the actual semantic search and one for the embedding cache which is easier to implement. But it is also possible to have one actual storage for both with additional metadata stored in the distributed cache like Redis for the mapping.

[dluc]

What about re-using the embeddings/vectors used by the search vs cache? We could store two copies of the vectors, one for the actual semantic search and one for the embedding cache which is easier to implement. But it is also possible to have one actual storage for both with additional metadata stored in the distributed cache like Redis for the mapping.

I'm not sure I understand, but I would avoid storing twice the same data. The embedding generator is a function generating a vector from some content, but for the same content different generators generate different embeddings, so it's important to take into account that. When caching OpenAI embeddings and then switching to Mixtral for example, the cache of OpenAI embeddings should not be used for Mixtral embeddings, and viceversa.
Rephrasing my previous comment, clients calling any embedding generator should not have to worry about cost or caching. It should be an embedding generator responsibility to cache internally, and when using a shared cache, not to conflict with other embedding generators, e.g. using cache keys that include all the context information needed to avoid cache collisions.

[dluc]

Looks like the PR has become stale, with a few things to address, see also linked issue.

If this is a pressing problem, the approach should be reusable (e.g. not having to add caching logic in every handler - usually caching is a cross-cutting concern solved with generic KV stores decoupled from specific scenarios), scale over multiple VMs (e.g. allow to extend the solution with Redis/Memcache), and being optional via config settings.