Implement cache embeddings (resolves #200)

[Pouyanpi]

Summary

This PR introduces a caching mechanism for embeddings in the BasicEmbeddingsIndex class. It uses a decorator cache_embeddings to cache the results of the _get_embeddings method. The caching mechanism includes key generation and cache storage, which are implemented using abstract base classes to allow for extensibility.

Changes

• Added cache_embeddings decorator to cache the results of the _get_embeddings method.
• Made _get_embeddings asynchronous.
• Introduced KeyGenerator and CacheStore abstract base classes for key generation and cache storage. (The base class also plays the role of a Factory)
• Implemented HashKeyGenerator, MD5KeyGenerator, InMemoryCacheStore, FilesystemCacheStore, and RedisCacheStore as concrete implementations of the abstract base classes.
• Updated BasicEmbeddingsIndex to use the caching mechanism.
• Implemented EmbeddingsCacheConfig

Limitations

• Cache eviction policy is not implemented.
• The key is generated based on the text content, so if we use two different embedding engines/models there would be a cache hit if the embeddings were generated before with a different model.
• the OpenAIEmbedding model does not return List[List[float]] but rather

Future Work

• Implement a cache eviction policy.
• Update the key generation to include the model name and engine to avoid cache hits when using different models. (easy)

Remarks

An object that needs to use cache embedding must instantiate these two objects (otherwise default values are used). The cache_embedding decorator could be used on any method that accepts a list of texts and returns a list of embeddings. So one could have applied it to the EmbeddingModel as well, but it seems more relevant to have it in BasicEmbeddingIndex as it is instantiated from the configs within the project. Furthermore, The CacheStore and KeyGenerator could be defined within the embedding config, the constructor could also accept a boolean to define whether to use caching at all.

Issue

This PR resolves issue #200.

Example Usage

Using the openai embedding engine

from nemoguardrails.embeddings.basic import BasicEmbeddingsIndex

# Example usage with OpenAI
embedding_index = BasicEmbeddingsIndex(
    embedding_engine="openai",
    embedding_model="text-embedding-ada-002",
    cache_config={"enabled": True, "store": "filesystem"}

)

embeddings = await embedding_index._get_embeddings(
    [
        "this is a text to test",
        "another text to test",
    ]
)

print(embeddings)

# If we run it for the second time it fetches the embeddings from file
embeddings = embedding_index._get_embeddings(
    [
        "this is a text to test",
        "another text to test",
    ]
)

print(embeddings)

using FastEmbed embedding engine

from nemoguardrails.embeddings.basic import BasicEmbeddingsIndex

embedding_index = BasicEmbeddingsIndex(
    embedding_engine="FastEmbed",
    embedding_model="all-MiniLM-L6-v2",
    cache_config = {"enabled": True},
)


# if we remove the "MODIFIED" we will have cache hits
# as discussed in limitations
embeddings = await embedding_index._get_embeddings(
    [
        "MODIFIED this is a text to test",
        "MODIFIED another text to test",
    ]
)

print(embeddings)

using in memory cache store

embedding_index = BasicEmbeddingsIndex(
    embedding_engine="FastEmbed",
    embedding_model="all-MiniLM-L6-v2",
    cache_config={"enabled": True, "store": "in_memory"}

)

embedding_index._get_embeddings(
    [
        "MODIFIED this is a text to test",
        "MODIFIED another text to test",
    ]
)

print(embeddings)

[drazvan]

Looks good! 👍 You should fix the "bulk computing" problem.
We also need to add documentation for this and potentially enable the caching behavior to be controlled from the config. Let me know if you want me to make a first proposal for that.

[drazvan]

Thanks for implementing this @Pouyanpi! 👍
On the more boring admin side, I see that your first commit is not signed. Feel free to sign it and force-push the branch again.

[Pouyanpi]

Thank you for your feedback, @drazvan . I've made the changes. Please let me know if you need further modifications.

If you'd like, I can also update/add the documentation to reflect these updates. Could you please advise on which file I should update and the level of detail you'd prefer?
Thanks!

[drazvan]

@Pouyanpi: do you have availability to wrap this PR this week? We'd like to include it in the 0.8.0 release.

[Pouyanpi]

Sure @drazvan, I will have it done later today.

[drazvan]

The d dictionary here can contain strings rather than actual instances. This will be the case when called with the parameters data from esp_provider. I think there should be a logic that translates from str to actual instances. e.g.

if key_generator == "md5":
  key_generator = MD5KeyGenerator()

[drazvan]

@Pouyanpi : Also, let's make the configuration of the embeddings cache first-class citizen, like this:

class EmbeddingsCache(BaseModel):
    """Configuration for the caching embeddings."""

    enabled: bool = Field(
        default=False,
        description="Whether caching of the embeddings should be enabled or not.",
    )
    key_generator: str = Field(
        default="md5",
        description="The method to use for generating the cache keys.",
    )
    store: str = Field(
        default="filesystem",
        description="What type of store to use for the cached embeddings.",
    )
    store_config: Dict[str, Any] = Field(
        default_factory=dict,
        description="Any additional configuration options required for the store. "
        "For example, path for `filesystem` or `host`/`port`/`db` for redis.",
    )


class EmbeddingSearchProvider(BaseModel):
    """Configuration of a embedding search provider."""

    name: str = Field(
        default="default",
        description="The name of the embedding search provider. If not specified, default is used.",
    )
    parameters: Dict[str, Any] = Field(default_factory=dict)
    cache: EmbeddingsCache = Field(
        default_factory=EmbeddingsCache,
        description="Configuration option for whether to use a cache for embeddings or not.",
    )

And the configuration will be done as per:
https://github.com/NVIDIA/NeMo-Guardrails/blob/develop/docs/user_guides/advanced/embedding-search-providers.md

core:
  embedding_search_provider:
    name: default
    parameters:
      embedding_engine: FastEmbed
      embedding_model: all-MiniLM-L6-v2
    cache:
      enabled: True
      # The rest of the settings are optional; this will use filesystem + md5

[drazvan]

And last but not least, have a look at this as well: #267. This makes the _get_embeddings method async. So, maybe update this (i.e. make the method async) in your PR as well. Not sure in what order we'll merge.

[Pouyanpi]

@drazvan : I have implemented the changes based on your suggestions and have pushed the updates for review.

Currently, the caching feature is disabled by default. Please let me know if this aligns with your expectations.

Regarding the factory-related behavior of CacheStore and KeyGenerator, I acknowledge it might not be the "clean" approach. I have alternative solutions in mind and can share preliminary drafts if you're interested for future releases. However, I believe the current implementation is satisfactory for the moment.

Feel free to review the modifications at your convenience.