feat: python bindings for the entire KvPushRouter

[PeaBrane]

Overview:

Bind the entire KvPushRouter to a python object with the generate method. It will accept token ids as List[int] to avoid the extra (de)serialization cycle. The other request fields like sampling and output options are simply pythonized. Closes #2662. Closes #2697

Others

Fixed the mocker engine always generating one fewer token than expected.

Test

Added an e2e test with mockers served with the Python-binded KvPushRouter

[coderabbitai]

Walkthrough

Registers KvPushRouter and KvPushRouterStream in Python bindings, implements their PyO3 classes enabling async streaming generation, adds a KvRouterConfig::inner accessor for internal use, refactors the mocker engine streaming to unbounded channels, and introduces an end-to-end Python test (duplicated) validating the new push-router bindings.

Changes

Cohort / File(s)	Summary
Binding registration lib/bindings/python/rust/lib.rs	Exposes llm::kv::KvPushRouter and llm::kv::KvPushRouterStream in the Python module init sequence.
KV push-router Python bindings lib/bindings/python/rust/llm/kv.rs	Adds PyO3 classes KvPushRouter and KvPushRouterStream. Implements constructor, generate(...) producing an async Python stream via mpsc channel, and __aiter__/__anext__ for iteration. Incorporates pythonize/depythonize and option types.
Router config accessor lib/bindings/python/rust/llm/entrypoint.rs	Adds KvRouterConfig::inner(&self) -> RsKvRouterConfig inherent method (non-Python).
Mocker engine streaming refactor lib/llm/src/mocker/engine.rs	Switches to unbounded channels and UnboundedReceiverStream, adjusts send/cancellation behavior and completion signaling; no public API change.
E2E test for push-router tests/router/test_router_e2e_with_mockers.py	Adds test_kv_push_router_bindings using mockers to validate async token streaming; asserts 20 tokens; includes setup/teardown. Note: test defined twice identically.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  actor PyApp as Python App
  participant PyBind as KvPushRouter (Py)
  participant RustTask as Rust async task
  participant Router as Kv Router
  participant Engine as Mocker Engine
  participant Stream as KvPushRouterStream (Py)

  PyApp->>PyBind: generate(token_ids, model, opts)
  activate PyBind
  PyBind->>RustTask: spawn task with PreprocessedRequest
  note right of RustTask: Builds request and starts streaming
  RustTask->>Router: route(request)
  Router->>Engine: stream tokens
  Engine-->>Router: token events
  Router-->>RustTask: token events
  RustTask-->>Stream: mpsc send PyObject per event
  deactivate PyBind

  loop until exhausted
    PyApp->>Stream: __anext__()
    Stream-->>PyApp: next token/result
  end
  Stream-->>PyApp: StopAsyncIteration on completion

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• feat: reduce / revert routing overheads, do not consider output tokens #2182 — Modifies KvPushRouter generate/stream logic, overlapping with the newly added Python bindings’ generate pathway.
• feat: expose router configurations to dynamo-run #1259 — Introduces/structures KvRouterConfig, which the new KvRouterConfig::inner accessor and bindings rely on.
• feat: Router replicas with state-sharing #2264 — Adjusts KvRouter/KvPushRouter internals invoked by the new bindings’ streaming path.

Poem

I thump my paws: new streams arrive,
Tokens hop out, alive, alive!
Channels unbounded, carrots aligned,
Kv pushes, mockers timed.
I nibble bytes, then softly gleam—
Async dreams in a rabbit stream. 🥕✨

Tip

🔌 Remote MCP (Model Context Protocol) integration is now available!

Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats.

---

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

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

CodeRabbit Commands (Invoked using PR/Issue comments)

Type @coderabbitai help to get the list of available commands.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

Status, Documentation and Community

• Visit our Status Page to check the current availability of CodeRabbit.
• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.