fix(rag): size the embedding server batch to its context so large chunks embed (#2836)

[Anai-Guo]

What

ramalama rag fails during the embedding phase on real-world doc sets with:

RuntimeError: llama-server embedding request returned 500:
{"error":{"code":500,"message":"input (596 tokens) is too large to process.
increase the physical batch size (current batch size: 512)","type":"server_error"}}

Root cause

The embedding llama-server is started (in rag_handler) with only --embedding and no batch override, so it keeps llama-server's default 512-token physical batch (-ub). An embedding request must fit entirely within one physical batch.

Meanwhile doc2rag sizes chunks with tiktoken / cl100k_base (--chunk-size 400 default), but the embedding model (embeddinggemma-300m) counts the same text with its own tokenizer. On URL-, code- and YAML-dense content the embedding-model count comes out up to ~1.9x higher, so a chunk measured <=400 tokens can arrive at the server as 500-750 real tokens — over the 512 batch limit. Plain prose rarely triggers it, which is why the failure looked random across doc sets (as reported in #2836).

Fix

Pass --batch-size / --ubatch-size sized to the embedding context (falling back to 2048) when launching the embedding server, so any chunk that fits the context also fits a single physical batch. This mirrors how embedding servers are normally configured (whole input in one batch) and needs no new user-facing flag.

Minimal change, confined to rag_handler._build_serve_args call for the embedding server.

Fixes #2836

🤖 Generated with Claude Code

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: c89b1bf8-41f7-47f1-ab38-f3ebce9662af

📥 Commits

Reviewing files that changed from the base of the PR and between 99e159b and 45b7c7f.

📒 Files selected for processing (1)

• ramalama/plugins/runtimes/inference/rag/handler.py

---

📝 Walkthrough

Summary by CodeRabbit

• Bug Fixes
  • Improved RAG embedding server configuration to prevent “input is too large” failures.
  • Batch sizing now automatically aligns with the configured embedding context size, with a default of 2048.

Walkthrough

The RAG handler now derives an embedding batch size from embed_ctx_size, defaulting to 2048, and passes it as both --batch-size and --ubatch-size when launching the embedding server.

Changes

RAG embedding batch sizing

Layer / File(s)	Summary
Configure embedding server batch sizes ramalama/plugins/runtimes/inference/rag/handler.py	Derives a defaulted embedding batch size and adds matching --batch-size and --ubatch-size runtime arguments while preserving the existing embedding server settings.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: mikebonnet, bmahabirbu

Poem

I’m a rabbit with batches aligned,
No oversized tokens left behind.
Context sets the pace,
Both sizes hold place,
And RAG hops along just fine.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the main change: sizing the embedding server batch to the embedding context to avoid large-chunk failures.
Description check	✅ Passed	The description is clearly about the same RAG embedding batch-size fix and the reported failure mode.
Linked Issues check	✅ Passed	The change satisfies #2836 via the allowed alternative of increasing the embedding server's physical batch size.
Out of Scope Changes check	✅ Passed	The PR is narrowly scoped to the embedding server launch args and related documentation.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[gemini-code-assist]

Code Review

This pull request configures the embedding batch size and micro-batch size based on the embedding context size (defaulting to 2048) to prevent failures when processing large chunks. The reviewer pointed out a potential issue where if embed_ctx_size is not specified, the batch size is set to 2048 but the context size remains None. Since llama.cpp requires the context size to be at least as large as the batch size, they suggested passing embed_batch_size as the ctx_size parameter to avoid startup or processing failures.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If embed_ctx_size is not specified (i.e., it is None), embed_batch_size defaults to 2048. However, ctx_size is still passed as embed_ctx_size (None) to _build_serve_args. In llama.cpp, the batch size (--batch-size / --ubatch-size) cannot exceed the context size (--ctx-size). If the default context size of the server/model is smaller than 2048 (e.g., 512), llama-server may fail to start or fail to process chunks larger than 512 tokens. Passing embed_batch_size as the ctx_size ensures that the context size is always at least as large as the batch size.

Suggested change

	embed_batch_size = embed_ctx_size or 2048
	embed_serve_args = _build_serve_args(
	args,
	embedding_model,
	embed_port,
	runtime_args=["--embedding"],
	runtime_args=[
	"--embedding",
	"--batch-size",
	str(embed_batch_size),
	"--ubatch-size",
	str(embed_batch_size),
	],
	ctx_size=embed_ctx_size,
	cache_reuse=0,
	)
	embed_batch_size = embed_ctx_size or 2048
	embed_serve_args = _build_serve_args(
	args,
	embedding_model,
	embed_port,
	runtime_args=[
	"--embedding",
	"--batch-size",
	str(embed_batch_size),
	"--ubatch-size",
	str(embed_batch_size),
	],
	ctx_size=embed_batch_size,
	cache_reuse=0,
	)

[rhatdan]

@Anai-Guo respond to @gemini-code-assist suggestion.