Feature request: cache document embeddings

[slivering]

Motivation

In collaborative office suites such as Twake Workplace, documents are updated every hour, such as email threads, notes and chat conversations.

When documents are updated, the embeddings are recomputed by vLLM every time, which can be computationally expensive. This is especially wasteful if the documents are long enough to be split in many chunks.

For more context, see these (LLM-generated) answers on Milvus AI Quick Reference:

• How can caching mechanisms be used in RAG to reduce latency, and what types of data might we cache (embeddings, retrieved results for frequent queries, etc.)?
• How can caching of computed embeddings help improve application performance when using Sentence Transformers repeatedly on the same sentences? (a bit more general)

Solution

The easiest way is to have a chunk -> embedding cache. LangChain provides a simple wrapper around OpenAIEmbeddings: https://python.langchain.com/docs/how_to/caching_embeddings/

The store could be a Redis cache with an eviction strategy such as LRU: https://redis.io/docs/latest/develop/reference/eviction/

The cache has to be cleared on application startup, which is necessary if the embedding API changed in the configuration.

Alternatives

The embedding API could be configured to cache the embeddings. I haven't found how to do this in vLLM and it is not always configurable if using an external API.

Related optimizations

• Prompt caching
• Document retrieval caching

[slivering]

Relationship with other functionalities

With SemanticSplitter

Embedding caching works extremely well with SemanticSplitter because SemanticChunker calls the embedding function on triples of sentences.

With contextual retrieval

However, text hash-based embedding caching would not work that well with contextual retrieval.
Because if a part of the document is modified, the document splitting will change, some chunks will be modified, and the context will change as well. So a significant number of contextualized chunks will be affected, depending on the chunks used.

openrag/openrag/components/indexer/chunker.py

Lines 89 to 93 in 935c134

	self._generate_context(
	first_chunks=first_chunks,
	prev_chunk=prev_chunk,
	chunk=curr_chunk,
	source=source,

At the moment, the context generator uses prev_chunk and first_chunks where first_chunks are the first 4 chunks of the document. Meaning that:

1. If the start of the document is modified, every chunk is likely to be affected.
2. If a section is added or modified, the nearby sections may be affected as well. When using MarkdownSplitter (default), the effect is mostly local. However, with RecursiveSplitter and SemanticSplitter changing one chunk might affect splitting of the whole document.

Performance improvements

Contextualization

Current complexity per file: (cost of LLM prompt on system prompt + (4 + 1 + 1) chunks) x number of chunks

The cost of embedding could be 10x smaller than LLM inference, because we can use smaller models and there's no output tokens. Therefore, contextualization is likely the main bottleneck if enabled.

Solutions

Since the contextualizer uses the LLM provided in the configuration, cost optimization is the user's responsibility if the LLM is self-managed. Thankfully, LLM runners such as Anthropic, Ollama and vLLM do implement prompt caching (on disk) and/or prefix caching (per LLM session).

Not sure if there are useful performance recommendations to give to the user.

Further benchmarking should tell whether contextualizer is still performant when doing incremental context changes.