fix: remove stream_usage from text completion#1285
Conversation
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## develop #1285 +/- ##
========================================
Coverage 69.78% 69.79%
========================================
Files 161 161
Lines 16057 16058 +1
========================================
+ Hits 11206 11207 +1
Misses 4851 4851
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Pull Request Overview
This PR removes the unsupported stream_usage parameter from the kwargs before initializing text completion providers to prevent API errors.
- Strip out
stream_usageso text completion calls don’t receive an invalid argument. - Ensures compatibility with providers like OpenAI’s AsyncCompletions.
Comments suppressed due to low confidence (1)
nemoguardrails/llm/models/langchain_initializer.py:258
- Add a unit test to verify that passing
stream_usageinto_init_text_completion_modelno longer raises an unexpected keyword argument error, ensuring this change is covered.
kwargs.pop("stream_usage", None)
| kwargs = _update_model_kwargs(provider_cls, model_name, kwargs) | ||
| # remove stream_usage parameter as it's not supported by text completion APIs | ||
| # (e.g., OpenAI's AsyncCompletions.create() doesn't accept this parameter) | ||
| kwargs.pop("stream_usage", None) |
There was a problem hiding this comment.
[nitpick] Consider extracting the "stream_usage" key into a shared constant or documenting it in the function docstring to avoid magic strings and make the removal rationale more discoverable.
| kwargs.pop("stream_usage", None) | |
| kwargs.pop(STREAM_USAGE_KEY, None) |
trebedea
left a comment
There was a problem hiding this comment.
This is a good fix for now, but we should try to find something more generic for unsupported attributes in kwargs for Langchain LLM providers. Let's merge this and talk in private about a generic fix.
This pull request includes a small but important change to the
_init_text_completion_modelfunction inlangchain_initializer.py. The change ensures compatibility with text completion APIs by removing the unsupportedstream_usageparameter from thekwargsdictionary before initializing the provider class.TODO:
ensure current implementation for token usage tracking does not break for other providers see #1264