refactor: refactored using Choice and CompletionFinishReason #1635
Conversation
paulhendricks
requested review from
a team,
GuanLuo,
PeaBrane,
alec-flowers,
biswapanda,
grahamking,
jthomson04,
kkranen,
oandreeva-nv,
rmccorm4,
ryanolson and
tmonty12
as code owners
WalkthroughThe changes refactor completion choice handling in the LLM codebase to use the official Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Engine
participant DeltaGen
participant OpenAI_Types
Client->>Engine: NvCreateCompletionRequest
Engine->>DeltaGen: Stream characters, yield responses
DeltaGen->>OpenAI_Types: Create Choice (typed finish reason)
OpenAI_Types-->>DeltaGen: Typed Choice
DeltaGen-->>Engine: Annotated response
Engine-->>Client: Completion with typed finish reason
Poem
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
lib/llm/src/protocols/openai/completions.rs (1)
287-288: TODO tracked for logprobs aggregation.The TODO comment indicates future work needed for aggregating logprobs. This aligns with the PR objective mentioning the use of
async_openai::types::Logprobs.Would you like me to create an issue to track the implementation of logprobs aggregation?
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
lib/llm/src/engines.rs(1 hunks)lib/llm/src/protocols/common.rs(3 hunks)lib/llm/src/protocols/openai/chat_completions/delta.rs(1 hunks)lib/llm/src/protocols/openai/completions.rs(4 hunks)lib/llm/src/protocols/openai/completions/aggregator.rs(9 hunks)lib/llm/src/protocols/openai/completions/delta.rs(4 hunks)
🧰 Additional context used
🧠 Learnings (3)
📓 Common learnings
Learnt from: kthui
PR: ai-dynamo/dynamo#1424
File: lib/runtime/src/pipeline/network/egress/push_router.rs:204-209
Timestamp: 2025-06-13T22:07:24.843Z
Learning: The codebase uses async-nats version 0.40, not the older nats crate. Error handling should use async_nats::error::Error variants, not nats::Error variants.
lib/llm/src/engines.rs (1)
Learnt from: kthui
PR: ai-dynamo/dynamo#1424
File: lib/runtime/src/pipeline/network/egress/push_router.rs:204-209
Timestamp: 2025-06-13T22:07:24.843Z
Learning: The codebase uses async-nats version 0.40, not the older nats crate. Error handling should use async_nats::error::Error variants, not nats::Error variants.
lib/llm/src/protocols/openai/completions.rs (1)
Learnt from: ishandhanani
PR: ai-dynamo/dynamo#1626
File: lib/llm/src/preprocessor.rs:238-239
Timestamp: 2025-06-24T20:59:35.725Z
Learning: In lib/llm/src/preprocessor.rs, the `sampling_options` call in the `preprocess_request` method is placed in the common section after the match statement on `request.prompt_input_type()`, meaning it applies to both `PromptInput::Tokens` and `PromptInput::Text` request types.
🧬 Code Graph Analysis (1)
lib/llm/src/engines.rs (1)
lib/llm/src/protocols/openai/completions/delta.rs (1)
create_choice(81-112)
🔇 Additional comments (16)
lib/llm/src/protocols/openai/chat_completions/delta.rs (1)
206-208: LGTM! Clean addition of ContentFilter mapping.The new mapping for
ContentFilterfollows the established pattern and correctly maps the internal enum variant to the OpenAI API equivalent.lib/llm/src/engines.rs (1)
241-241: Excellent type safety improvement.Replacing the string literal with the strongly typed
CompletionFinishReason::Stopenum variant eliminates potential runtime errors from invalid finish reason strings and aligns with the broader refactoring to useasync_openaitypes.lib/llm/src/protocols/common.rs (3)
68-70: Well-structured enum extension.The addition of the
ContentFiltervariant follows the established pattern with proper serde serialization support.
80-80: Consistent Display implementation.The Display implementation for
ContentFilterfollows the same pattern as other variants.
100-125: Robust bidirectional enum conversions.The
Fromtrait implementations correctly handle all enum variants with sensible mappings:
- Multiple internal variants (EoS, Stop, Cancelled) map to OpenAI's Stop, which is appropriate
- ContentFilter and Length map directly to their OpenAI equivalents
- Error maps to Stop as a reasonable fallback
- The reverse conversion is straightforward since OpenAI has fewer variants
This provides a clean interface between internal and external enum representations.
lib/llm/src/protocols/openai/completions/delta.rs (3)
85-85: Improved method signature with type safety.Changing the parameter type from string to
Option<async_openai::types::CompletionFinishReason>eliminates potential runtime errors from invalid finish reason strings and provides compile-time type checking.
100-105: Consistent use of official async_openai types.Using
async_openai::types::Choiceinstead of a custom struct aligns with the refactoring goals and ensures compatibility with the OpenAI API specification.
127-127: Simplified finish reason conversion.The use of
Into::intoleverages the new trait implementations fromcommon.rs, eliminating manual string conversions and reducing error-prone code.lib/llm/src/protocols/openai/completions/aggregator.rs (4)
101-103: Proper index type handling.The casting between
u32(OpenAI API) andu64(internal storage) is handled consistently for both entry creation and DeltaChoice construction.
114-125: Type-safe finish reason handling.The conversion from
async_openai::types::CompletionFinishReasonto internalFinishReasonnow uses pattern matching on typed enums instead of string parsing, eliminating potential runtime errors and improving maintainability. All OpenAI finish reason variants are properly handled.
160-171: Clean conversion to official OpenAI types.The
From<DeltaChoice>implementation correctly converts toasync_openai::types::Choiceusing the new trait implementations, with proper index type casting back tou32for the OpenAI API.
203-207: Well-updated test code.The test code has been properly updated to:
- Use
async_openai::types::Choiceinstead of custom structs- Parse finish reasons through the internal enum with proper error handling
- Assert against typed enum variants instead of strings
- Maintain test coverage for the new type system
The tests now provide better validation of the typed enum behavior.
Also applies to: 216-221, 276-279, 307-310, 324-336, 360-371
lib/llm/src/protocols/openai/completions.rs (4)
52-52: LGTM! Type migration aligns with PR objectives.The migration from legacy
CompletionChoicetoasync_openai::types::Choiceis consistent with the goal of using official async_openai types throughout the codebase.
79-83: LGTM! Trait implementation updated for the new type.The
ContentProviderimplementation correctly adapts to useasync_openai::types::Choice, maintaining the same content extraction logic.
204-218: LGTM! Method signature updated consistently.The
make_responsemethod signature correctly reflects the newasync_openai::types::Choicetype, maintaining consistency with theCompletionResponsestruct changes.
274-301: LGTM! Conversion implementation properly handles the new type.The implementation correctly converts from the internal streaming response to
async_openai::types::Choice. The safety comment for the u32 cast is appreciated, and the error handling for missing text is appropriate.
Overview:
Refactoring to use
async_openai::types::LogprobsDetails:
This change is to be consistent with previous renaming and refactoring for protocols:
Where should the reviewer start?
Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)
Summary by CodeRabbit
New Features
Refactor
Bug Fixes