Add platform api wrapper #519
Conversation
guapisolo
requested review from
fzyzcjy,
yueming-yuan and
yushengsu-thu
as code owners
Summary of ChangesHello @guapisolo, 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 significantly enhances the SGLang framework's OpenAI-compatible API by providing more detailed token-level information, particularly Highlights
Ignored Files
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. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
The pull request significantly enhances the OpenAI-compatible chat completion endpoint by introducing support for input_ids and detailed logprobs for prompt tokens. It also refactors the session management within the miles.router module into a more modular structure. New examples and tests are added to demonstrate agentic tool-calling and Fireworks reward calculation, leveraging these new capabilities. While the changes improve flexibility and provide more detailed token information, some areas could benefit from further refinement, particularly regarding code duplication, clarity of experimental features, and consistency in testing.
|
|
||
| input_token_ids = _extract_input_token_ids(choice) | ||
|
|
||
| assert input_token_ids == [] |
There was a problem hiding this comment.
The test test_chat_completions_input_logprobs_prompt_ids_match asserts input_token_ids == []. However, the changes in protocol.py and serving_chat.py introduce input_token_ids to the response. This test seems to contradict the new functionality and might be outdated or incorrectly asserting the expected behavior. It should be updated to reflect that input_token_ids are now expected to be present in the response when logprobs are requested.
| self.router_url = router_url | ||
| self.session_id = session_id | ||
| self.base_url = f"{router_url}/sessions/{session_id}/v1" | ||
| self.base_url = f"{router_url}/sessions/{session_id}" |
There was a problem hiding this comment.
The base_url construction has changed from f"{router_url}/sessions/{session_id}/v1" to f"{router_url}/sessions/{session_id}". This implies a change in the API endpoint structure for session-related operations. While it aligns with the new session management, it's a breaking change for any external clients expecting the /v1 prefix. This should be clearly communicated in release notes or migration guides.
| completion_params["logprobs"] = True | ||
| completion_params["logprob_start_len"] = 0 |
There was a problem hiding this comment.
The build_chat_completion_params function unconditionally sets logprobs=True and logprob_start_len=0. This overrides any logprobs or logprob_start_len values that might be present in the sampling_params provided by the user, which could lead to unexpected behavior or loss of user control over these parameters. Consider merging or prioritizing user-provided values.
| completion_params["logprobs"] = True | |
| completion_params["logprob_start_len"] = 0 | |
| if "logprobs" not in completion_params: | |
| completion_params["logprobs"] = True | |
| if "logprob_start_len" not in completion_params: | |
| completion_params["logprob_start_len"] = 0 |
| # Process keys to match ChatCompletionRequest input | ||
| def build_chat_completion_params(sampling_params: dict[str, Any]) -> dict[str, Any]: | ||
| completion_params = dict(sampling_params) | ||
| key_map = { | ||
| "max_new_tokens": "max_tokens", | ||
| "min_new_tokens": "min_tokens", | ||
| "sampling_seed": "seed", | ||
| } | ||
| for src, dst in key_map.items(): | ||
| if src in completion_params: | ||
| if dst not in completion_params: | ||
| completion_params[dst] = completion_params[src] | ||
| completion_params.pop(src, None) | ||
|
|
||
| # Notice: Here we force the inference backend to return token information and start from 0 | ||
| # The start len should be 0 to make sure prompt token ids and be correctly returned from SGLang. | ||
| completion_params["logprobs"] = True | ||
| completion_params["logprob_start_len"] = 0 | ||
|
|
||
| reserved_keys = {"model", "messages"} |
There was a problem hiding this comment.
The build_chat_completion_params function is duplicated in this file and miles/rollout/generate_hub/agentic_tool_call.py. To improve maintainability and reduce redundancy, this function should be extracted into a shared utility module and imported where needed.
| request_input_ids == input_token_ids | ||
| ), "for prompt part, input_ids return by sglang should match with the request input_ids" |
There was a problem hiding this comment.
The assertion request_input_ids == input_token_ids assumes that the input_ids provided in the request will exactly match the input_token_ids returned by the SGLang backend. This might not always hold true if different tokenizers are used or if there are subtle differences in tokenization logic between the client and the server. A more robust comparison or a warning might be appropriate instead of a hard assertion.
| List of token (IDs, logp, loss_mask) corresponding to the input text | ||
| if return_logprob is False, all logp will be 0.0 |
There was a problem hiding this comment.
The docstring for retrieve_from_text states "if return_logprob is False, all logp will be 0.0". This behavior is somewhat counter-intuitive. If return_logprob is False, it might be more appropriate to return None or an empty list for logp to clearly indicate that log probabilities were not requested or are not available, rather than returning a list of zeros. Alternatively, if logp is always returned, the return_logprob parameter could be removed.
| # This is only a naive trajectory manager to store history session record. | ||
| # Cross turn token input not implemented. |
There was a problem hiding this comment.
The comment "Cross turn token input not implemented." highlights a known limitation. For clarity and future development, it would be beneficial to elaborate on what "cross turn token input" entails and why it's not implemented, or provide a roadmap for its implementation.
| if prompt_token_ids is None: | ||
| return JSONResponse(status_code=404, content={"error": "session not found"}) |
There was a problem hiding this comment.
In the chat_completions endpoint, if manager.calc_prompt_tokens returns None (indicating the session was not found), the function immediately returns a 404 JSON response. While functionally correct, a more specific error message or logging could help in debugging, especially since this is part of an experimental feature.
| if variant.startswith("agentic_tool_call"): | ||
| mock_tools.AGENTIC_MAX_TURNS = args_kwargs.get("generate_max_turns") |
There was a problem hiding this comment.
The generation_env fixture directly modifies the global variable mock_tools.AGENTIC_MAX_TURNS. While this might be acceptable for testing purposes, relying on global state can lead to test isolation issues and make tests harder to reason about. Consider passing this value as a parameter or using a context manager if possible, to limit its scope.
1 participant
Wait for OAI TITO to be merged.