Contextual Retrieval: scalability and usability issues

[LukaszCmielowski]

System Info

llama-stack 0.6.0

Information

• The official example scripts
• My own modified scripts

🐛 Describe the bug

The current contextual retrieval implementation (VectorStoreChunkingStrategyContextual) in llama-stack has two main issues that make it unsuitable for real-world use:

1. Files API upload is mandatory — no external storage referenceContextual retrieval is only available via vector_stores.files.create, which requires a file_id from the Files API upload endpoint. There is no way to use a document that already exists in external storage (S3) without uploading it through the llama-stack server (duplication of input documents and the burden of syncing them).

2. One LLM call per chunk — no batchingFor each chunk, a separate chat_completion call is made. With default chunking (~700 tokens, 400 overlap), a 50-page doc yields ~200 chunks and 200 LLM calls; 500 such documents yield ~100,000 sequential requests. Concurrency is limited by a semaphore (default 3). There is no multi-chunk batching, so round-trip count and latency scale linearly with chunk count. Retry logic (up to 3 retries with exponential backoff per chunk) can make rate-limited runs much slower.

Combined impact: Indexing a 500-document corpus (50 pages each) requires uploading ~2.5 GB through the HTTP server (even if data is already in S3), ~100k LLM calls at low concurrency (hours of runtime), and there is no progress visibility or resumability.

---

Steps to reproduce

Problem 1 (no external reference):

1. Put documents in S3 (or another object store) that llama-stack can access.
2. Try to attach those documents to a vector store using contextual retrieval without uploading them via the Files API (e.g. by providing an S3 URI or presigned URL as the document source).

Problem 2 (one LLM call per chunk):

1. Create a vector store with contextual chunking (e.g. VectorStoreChunkingStrategyContextual with default or custom chunk settings).
2. Attach a multi-page document (e.g. 50+ pages) via the Files API.
3. Inspect server logs or add logging; observe one chat_completion (or equivalent) request per chunk.

---

Actual results

• Problem 1: Contextual retrieval is only usable with documents uploaded through the Files API.
• Problem 2: Each chunk triggers exactly one LLM request. For ~200 chunks per 50-page doc and 500 docs, that is ~100,000 LLM calls.

---

Workarounds

Problem 1 & 2: Using vector_io.insert_chunks with pre-computed embeddings avoids the Files API and allows a custom batching implementation (users must implement contextualization and chunking themselves).

---

Code references

• src/llama_stack/providers/utils/memory/openai_vector_store_mixin.py: _execute_contextual_chunk_transformation (~line 1572), openai_attach_file_to_vector_store (~line 983); contextual chunk loop ~1619–1697 (one openai_chat_completion per chunk, semaphore, retries).
• Files API upload (full read into server):
  • src/llama_stack/providers/inline/files/localfs/files.py (~line 111)
  • src/llama_stack/providers/remote/files/s3/files.py (~line 236)
  • src/llama_stack/providers/remote/files/openai/files.py (~line 146)
• src/llama_stack_api/vector_io/models.py: VectorStoreChunkingStrategyContextualConfig (~line 371), InsertChunksRequest (~line 733).

Error logs

N/A

Expected behavior

• Problem 1: Support for external document references that can be used as the source for contextual retrieval without uploading files through the llama-stack server.
• Problem 2: Batch multiple chunks per LLM call (e.g. structured output) and/or expose an optimized path (e.g. insert_chunks with contextual strategy) so call count and runtime don’t scale one-to-one with chunks.

[LukaszCmielowski]

Are there any benchmarks run ? if yes, can the numbers be shared (correctness and performance) ?

@franciscojavierarceo fyi.

[gyliu513]

@LukaszCmielowski llamastack do not have API-level / component-level benchmarks yet, I think it is a good time to add some benchmark test.

@leseb @mattf @franciscojavierarceo @cdoern thoughts?

[franciscojavierarceo]

Thanks for the detailed write-up — you've identified some real scaling considerations. Let me address each point with some additional context from the implementation.

On the 1:1 Chunk-to-LLM Scaling

You're right that each chunk currently gets its own chat_completion call (openai_vector_store_mixin.py:1621-1697). However, a couple of things worth noting:

1. Prefix caching is already built in. The implementation intentionally splits the prompt so the full document goes into the system message and only the chunk content goes into the user message (lines 1611-1617). This means providers that support KV-cache reuse (OpenAI, vLLM, Anthropic) compute the document prefix once and reuse it across all chunk requests — significantly reducing the actual per-chunk compute cost.

2. Embeddings are already batched. All chunks for a file are embedded in a single openai_embeddings() call (lines 1006-1011), not one-at-a-time.

3. We do have a Batch API (/v1/batches) that supports batched chat/completions requests. Today it isn't wired into the contextual chunking pipeline, but integrating it is a natural improvement we'd welcome contributions on.

That said, inference will always be the bottleneck for contextual retrieval — that's inherent to the approach (per Anthropic's original Contextual Retrieval design). The LLM must reason about each chunk in the context of the full document.

On the Concurrency Defaults

The default max_concurrency of 3 (datatypes.py:575) is intentionally conservative — it's designed to avoid overwhelming inference providers. This is configurable both at the server level (contextual_retrieval_params.default_max_concurrency) and per-request via chunking_strategy.contextual.max_concurrency. If you're running against a provider that can handle higher throughput, you can increase it.

On Mandatory File Uploads

This is intentional. Llama Stack enforces Attribute-Based Access Control (ABAC), and file uploads go through our access control layer. Allowing arbitrary external references (S3 URIs, etc.) would bypass these controls. Files need to be registered through the Files API so they can be properly governed.

On Scale Expectations

For the scale you're describing (500 documents, ~100k chunks), Llama Stack's contextual retrieval API is probably not the right tool. This is an API server, not a distributed compute engine. For large-scale preprocessing we'd recommend:

1. Use a compute framework (Spark, Ray, Flink) for chunking + contextualization at scale
2. Upload pre-computed results via POST /v1/vector-io/insert, which accepts EmbeddedChunk objects with pre-computed embeddings, content, and metadata
3. Treat Llama Stack as the serving/query layer

Here's a simplified example of the insert_chunks request body:

{
  "vector_store_id": "my-store",
  "chunks": [
    {
      "content": "contextualized chunk text...",
      "chunk_id": "chunk-001",
      "metadata": {"document_id": "doc1", "source": "my-pipeline"},
      "chunk_metadata": {
        "document_id": "doc1",
        "chunk_id": "chunk-001",
        "source": "my-pipeline"
      },
      "embedding": [0.1, 0.2, ...],
      "embedding_model": "your-embedding-model",
      "embedding_dimension": 384
    }
  ]
}

On Horizontal Scaling

Llama Stack has a K8s operator with a LlamaStackDistribution CRD that supports a replicas field, so horizontal scaling via HPA is possible. However, for contextual retrieval specifically, scaling the API server horizontally won't help much — the bottleneck is inference throughput, not request handling. Scaling the inference backend (e.g., more vLLM replicas) would be more impactful than scaling the Llama Stack server itself.

What We'd Welcome

• Batch API integration for contextual chunking: The existing /v1/batches endpoint already processes chat/completions requests from JSONL input with configurable concurrency. The contextual chunking pipeline (_execute_contextual_chunk_transformation) already constructs the same OpenAIChatCompletionRequestWithExtraBody objects — the integration path would be to serialize those as JSONL batch input rather than calling inference inline, then collect results from the batch output file. This would decouple contextualization from the synchronous file-attach flow and enable better throughput.
• File Processor API improvements: Better support for reprocessing files that are already uploaded without needing to re-upload
• Documentation: A guide on recommended patterns for large-scale ingestion (preprocessing externally → batch upsert via POST /v1/vector-io/insert)

Happy to discuss further or review PRs in any of these areas.

• co-authored by Claude 😅 🤖