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]

Caution

Review failed

The head commit changed during the review from ac8c4f7 to 9fc3958.

📝 Walkthrough

Walkthrough

The RAG handler derives an embedding batch size from embed_ctx_size, defaults to 2048, and applies it to the embedding server’s context, batch, and ubatch settings.

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 using the same value for ctx_size.

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

Suggested reviewers: jhjaggars, olliewalsh, maxamillion, rhatdan, mikebonnet

Poem

I’m a rabbit with batches aligned,
Context and ubatch now entwined.
Tokens hop with care,
Through settings fair,
While RAG leaves big errors behind.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: sizing the embedding server batch/context to avoid large chunk failures.
Description check	✅ Passed	The description directly matches the code change and explains the token-limit failure and batch-size fix.
Linked Issues check	✅ Passed	The change addresses #2836 by ensuring the embedding server batch/context accommodates token-dense chunks without manual chunk-size reduction.
Out of Scope Changes check	✅ Passed	The diff stays confined to the rag embedding server launch logic and explanatory comments, with no unrelated changes.
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.

[Anai-Guo]

Thanks @rhatdan, @gemini-code-assist — good catch, adopted in 2c65fd2.

When --embed-ctx-size is unset we now drive ctx_size from embed_batch_size (both 2048) instead of leaving ctx_size=None. llama.cpp requires batch <= ctx, so passing a 2048 batch while letting the server fall back to a possibly-smaller default context (e.g. 512) would fail to process larger chunks. Tying the two together keeps the invariant regardless of the model's default context. When the user does pass --embed-ctx-size, embed_batch_size == embed_ctx_size, so behavior is unchanged for that path.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[coderabbitai]

🧹 Nitpick comments (1)

ramalama/plugins/runtimes/inference/rag/handler.py (1)

66-81: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Add regression coverage for the batch/context contract.

Please test that the default configuration produces 2048 for ctx_size, --batch-size, and --ubatch-size, while an explicit --embed-ctx-size preserves that configured value. This protects the batch <= context fix and prevents accidental regressions in explicit-context behavior.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@ramalama/plugins/runtimes/inference/rag/handler.py` around lines 66 - 81, Add
regression tests for the embedding serve-argument construction around
_build_serve_args: verify the default configuration sets ctx_size, --batch-size,
and --ubatch-size to 2048, and verify an explicit --embed-ctx-size sets all
three values to that configured size.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@ramalama/plugins/runtimes/inference/rag/handler.py`:
- Around line 66-81: Add regression tests for the embedding serve-argument
construction around _build_serve_args: verify the default configuration sets
ctx_size, --batch-size, and --ubatch-size to 2048, and verify an explicit
--embed-ctx-size sets all three values to that configured size.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: aeba4d4e-eeb4-4b61-9045-2bb2fda59f5e

📥 Commits

Reviewing files that changed from the base of the PR and between 45b7c7f and 2c65fd2.

📒 Files selected for processing (1)

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

[rhatdan]

@bmahabirbu PTAL

[rhatdan]

@Anai-Guo Fix the lint issue.

[Anai-Guo]

@rhatdan I looked into it — I don't think there's anything left to fix in this branch.

The failing Lint Code run is from 26 Jul (run 29961834097) and it fails in ruff format --check, not on the file this PR touches:

1 file would be reformatted, 179 files already formatted
make: *** [Makefile:129: check-format] Error 1

Checked locally against current main with the CI-resolved formatter (ruff 0.16.0, repo pyproject.toml, line-length = 120, quote-style = "preserve"):

$ git clone --depth 1 containers/ramalama && cp <this PR's handler.py> ramalama/plugins/runtimes/inference/rag/handler.py
$ ruff format --check ramalama scripts test bin/ramalama
182 files already formatted

So main + this PR's only changed file is format-clean; the 179-file count in the stale run also doesn't match today's 182, which is consistent with this being drift on the base at that time rather than something in this branch.

I pushed a signed-off empty commit (c55ff6f) to re-run the checks — DCO is green, but the workflows are sitting at action_required, so they need an approval from a maintainer to actually execute. Happy to fix anything the fresh run turns up.