Repository navigation
feat(completions): add CompletionRequestBuildingStage and backend sampling params - #915
Conversation
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request is the third in a series aimed at fully implementing the Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
|
Hi @vschandramourya, this PR has merge conflicts that must be resolved before it can be merged. Please rebase your branch: git fetch origin main
git rebase origin/main
# resolve any conflicts, then:
git push --force-with-lease |
📝 WalkthroughWalkthroughAdds /v1/completions request-building: backend-specific builders that convert OpenAI CompletionRequest into proto GenerateRequest for Sglang, vLLM, and TRTLLM, plus a pipeline CompletionRequestBuildingStage and a GrpcClient dispatcher to produce ProtoGenerateRequest instances. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant PipelineStage as CompletionRequestBuildingStage
participant GrpcClient
participant Backend as BackendBuilder
participant Worker as Worker/Engine
Client->>PipelineStage: incoming /v1/completions request
PipelineStage->>PipelineStage: extract preparation (original_text, token_ids)
PipelineStage->>GrpcClient: build_completion_request(request_id, CompletionRequest, original_text, token_ids)
GrpcClient->>Backend: call build_generate_request_from_completion(...)
Backend->>Worker: produce proto GenerateRequest -> forward to backend worker/engine
Worker-->>GrpcClient: ack / response stream
GrpcClient-->>PipelineStage: proto request created
PipelineStage-->>Client: continue pipeline (proto stored)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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. Comment |
There was a problem hiding this comment.
Code Review
This pull request introduces comprehensive support for the /v1/completions API endpoint across Sglang, Trtllm, and Vllm backends. It includes adding CompletionRequest handling to the gRPC client implementations, mapping completion parameters to internal GenerateRequest structures and sampling configurations. New pipeline stages, CompletionPreparationStage and CompletionRequestBuildingStage, have been integrated into the model gateway to process these requests, managing tokenization, stop decoder creation, and constructing backend-specific gRPC requests. Feedback suggests refactoring for improved code conciseness, efficiency, and clarity, particularly by utilizing StringOrArray::to_vec() for stop sequence initialization and optimizing constraint-building functions to reduce allocations and cloning.
f834d92 to
2674efa
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@crates/grpc_client/src/sglang_scheduler.rs`:
- Around line 717-735: CompletionRequest.min_tokens is not being propagated into
proto::SamplingParams, so completions can end before the caller requested;
update the SamplingParams construction in sglang_scheduler to map
request.min_tokens into the sampling params (set
proto::SamplingParams.min_new_tokens to request.min_tokens.unwrap_or_default()
or an explicit 0 where appropriate) alongside the existing max_new_tokens
mapping, using the same pattern as other optional fields so min_tokens is
honored by the backend.
In `@crates/grpc_client/src/trtllm_service.rs`:
- Around line 750-777: The GenerateRequest currently hardcodes
include_stop_token_in_output to false so CompletionRequest.no_stop_trim is
ignored; update the construction of proto::GenerateRequest in trtllm_service.rs
to read the no_stop_trim flag from the incoming CompletionRequest (e.g.,
body.no_stop_trim) and set include_stop_token_in_output accordingly (true when
no_stop_trim is true, fallback to false when absent) so the
include_stop_token_in_output field is threaded through to TensorRT-LLM.
In `@crates/grpc_client/src/vllm_engine.rs`:
- Around line 648-666: The proto::SamplingParams construction is missing mapping
for CompletionRequest.min_tokens and CompletionRequest.no_stop_trim, so add
sampling.min_tokens and sampling.include_stop_str_in_output to the returned
proto::SamplingParams: set min_tokens from the request's min_tokens (use the
request default when absent) and set include_stop_str_in_output to the logical
inverse of no_stop_trim (i.e., include_stop_str_in_output = !no_stop_trim,
handling the Option properly), ensuring the field names min_tokens and
include_stop_str_in_output appear in the struct literal alongside the existing
fields like temperature, top_p, stop, and ignore_eos.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: df880bfe-fa72-4d6a-8ff0-3880903832da
📒 Files selected for processing (6)
crates/grpc_client/src/sglang_scheduler.rscrates/grpc_client/src/trtllm_service.rscrates/grpc_client/src/vllm_engine.rsmodel_gateway/src/routers/grpc/client.rsmodel_gateway/src/routers/grpc/regular/stages/completion/mod.rsmodel_gateway/src/routers/grpc/regular/stages/completion/request_building.rs
…pling params Signed-off-by: VS Chandra Mourya <msrinivasa@together.ai>
2674efa to
3b8ac76
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@model_gateway/src/routers/grpc/client.rs`:
- Around line 387-423: The three backend builders disagree on how to handle
CompletionRequest.max_tokens == None, so in build_completion_request you must
normalize max_tokens once before calling
client.build_generate_request_from_completion for Sglang, Vllm and Trtllm:
compute an effective_max_tokens (e.g., let normalized =
body.max_tokens.or(Some(<agreed-default>)) or explicitly keep None if that is
the agreed behavior) and pass a modified/temporary CompletionRequest (or the
normalized value) into each client's build_generate_request_from_completion call
so all branches (Sglang, Vllm, Trtllm) receive the same effective max_tokens
value when constructing the ProtoGenerateRequest.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 9f013ae6-9d72-4240-889a-5b786448c7d6
📒 Files selected for processing (6)
crates/grpc_client/src/sglang_scheduler.rscrates/grpc_client/src/trtllm_service.rscrates/grpc_client/src/vllm_engine.rsmodel_gateway/src/routers/grpc/client.rsmodel_gateway/src/routers/grpc/regular/stages/completion/mod.rsmodel_gateway/src/routers/grpc/regular/stages/completion/request_building.rs
…pling params (smg-project#915) Signed-off-by: VS Chandra Mourya <msrinivasa@together.ai>
Summary
CompletionRequestBuildingStage) for the/v1/completionsgRPC pipelinebuild_generate_request_from_completion+ sampling params builders to all 3 backends (SGLang, vLLM, TRT-LLM)build_completion_requestdispatcher toGrpcClientPR 3 in the Completions API gRPC pipeline series. PR 1 was #840 (type scaffolding), PR 2 was #907 (preparation).
What changed
New file
model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs—CompletionRequestBuildingStage, parallel toMessageRequestBuildingStage. Usescmpl_{uuid}request ID prefix, callsbuild_completion_request(), no multimodal, no tools.Backend sampling params (3 files in
crates/grpc_client/src/)sglang_scheduler.rs—build_generate_request_from_completion()+build_grpc_sampling_params_from_completion(). Maps allCompletionRequestsampling fields (temperature, top_p, top_k, min_p, frequency_penalty, presence_penalty, repetition_penalty, n, logprobs, stop, stop_token_ids, ignore_eos, no_stop_trim) + structured output constraints (json_schema, regex, ebnf).vllm_engine.rs— Same pattern for vLLM. Handles vLLM-specific differences (top_k=0 for disabled,Option<f32>temperature).trtllm_service.rs— Same pattern using TRT-LLM'sSamplingConfig+OutputConfig+GuidedDecodingParamsproto types. Includes seed, min_tokens, min_p mapping.GrpcClient dispatcher
model_gateway/src/routers/grpc/client.rs—build_completion_request()dispatches to each backend'sbuild_generate_request_from_completion().Module wiring
model_gateway/src/routers/grpc/regular/stages/completion/mod.rs— Wirerequest_buildingmodule + re-exportCompletionRequestBuildingStage.How
Follows the same architecture as Messages request building (#744):
Key difference from Messages:
CompletionRequesthas richer sampling knobs (frequency_penalty, presence_penalty, repetition_penalty, min_p, n, logprobs, ignore_eos, no_stop_trim) and direct structured output constraints (regex, ebnf, json_schema), but no tools and no multimodal.Test plan
cargo clippy -p smg --all-targets --all-features -- -D warnings— passescargo clippy -p smg-grpc-client --all-targets --all-features -- -D warnings— passescargo fmt --check— passes#![allow(dead_code)]is usedRefs: #840, #907
Summary by CodeRabbit