Skip to content

feat(rerank): forward max_tokens_per_doc to providers that support it - #3002

Open
chethanuk wants to merge 2 commits into
volcengine:mainfrom
chethanuk:feat/rerank-forward-token-truncate
Open

feat(rerank): forward max_tokens_per_doc to providers that support it#3002
chethanuk wants to merge 2 commits into
volcengine:mainfrom
chethanuk:feat/rerank-forward-token-truncate

Conversation

@chethanuk

Copy link
Copy Markdown
Contributor

Summary

Follow-up to the max_chars_per_doc cap (#2880). A character cap is provider-agnostic but cannot bound tokens: for CJK text, N characters can be far more than N tokens, so a char cap can still overflow a token-limited reranker. Providers that expose native token truncation should use it.

Adds RerankConfig.max_tokens_per_doc (default 0) and forwards it as max_tokens_per_doc to the providers whose APIs accept it. VikingDB/doubao has no truncation field and is intentionally left untouched (use max_chars_per_doc there).

How it works

flowchart TD
    A["RerankConfig.max_tokens_per_doc"] --> B{"value > 0 ?"}
    B -- "0 (default) → omit field" --> Z["every provider: byte-identical request"]
    B -- "> 0" --> C["from_config forwards value"]
    C --> D["Cohere v2 → /v2/rerank request body"]
    C --> E["OpenAI-compatible (DashScope / vLLM / TEI) → request body"]
    C --> F["LiteLLM → litellm.rerank(...) kwarg"]
    C -. "not wired (no token field)" .-> G["VikingDB / doubao: unchanged"]
Loading

Notes

  • The parameter is sent only when > 0, so default requests are byte-identical to before — asserted per-provider.
  • No change to response parsing, token tracking, or fallback paths (the Cohere edit only lifts the inline JSON into a req_body dict so the field can be added conditionally).
  • Field(default=0, ge=0, strict=True); extra="forbid" keeps typos fatal.

Tests / docs

Each client asserts it forwards the value when set and omits it when 0, plus from_config wiring, config validation, and a guard that VikingDB ignores the field. en + zh doc rows added. Unit suite green; ruff clean.

Review

The mirror PR received no actionable inline review-bot comments (automated review returned "no issues found").

Stacked on the max_chars_per_doc PR (#2999) — it depends on the char cap. Please review/merge #2999 first; until then this PR's diff includes the v0 commits. Closes the CJK char-cap gap noted in #2880.

@chethanuk
chethanuk force-pushed the feat/rerank-forward-token-truncate branch from d4f3784 to aad66f7 Compare July 4, 2026 02:57
@chethanuk

Copy link
Copy Markdown
Contributor Author

Please review cc: @qin-ctx @ZaynJarvis @zhoujh01 @yufeng201 @chenjw - Please let me know if you require any changes, or if the project is currently accepting contributions

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: Backlog

Development

Successfully merging this pull request may close these issues.

1 participant