fix: modernize Cohere generator (#1384)#1673
Open
HayatoFujihara wants to merge 4 commits intoNVIDIA:mainfrom
Open
fix: modernize Cohere generator (#1384)#1673HayatoFujihara wants to merge 4 commits intoNVIDIA:mainfrom
HayatoFujihara wants to merge 4 commits intoNVIDIA:mainfrom
Conversation
- Extract _load_client() for client initialization, supporting pickle serialization/deserialization via __setstate__ - Remove deprecated 'command' model default; model name is now required - Add 404 (NotFoundError) handling for both v1 and v2 API paths, raising BadGeneratorException to stop execution immediately - Migrate v1 error handling from RuntimeError to cohere.errors - Wrap string results in Message objects in _call_model for type consistency with base Generator contract - Update tests: model name required, _load_client, serialization, 404 handling, Message type verification Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
- Add live chat API test verifying Message type and non-empty text - Add live v1 API test verifying 404 handling for newer models - Tests are skipped when COHERE_API_KEY is not set (CI-safe) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
jmartin-tech
requested changes
Apr 6, 2026
Collaborator
jmartin-tech
left a comment
jmartin-tech
left a comment
There was a problem hiding this comment.
This is a good start on modernization.
A few asks from the issue have drifted from the core design requiring a bit of guideance here.
Also the generator currently accepts configuration to support use of the legacy v1 client via the legacy code in self.cohere.Client. #1673 is also linked to #1323 denoting a desire to have self.cohere.ClientV2 be used for all inference requests. Consider evaluating the scope that additional change to determine if it can be incorporated easily with this PR.
…move __setstate__, use config_root in tests (NVIDIA#1384) - Rename _load_client() to _load_unsafe() to follow Configurable framework convention - Remove redundant __setstate__ override (Configurable.__setstate__ auto-calls _load_unsafe) - Replace direct name= constructor arg with config_root dict pattern in all tests - Add cohere_config_root and cohere_v1_config_root test fixtures Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
, NVIDIA#1384) - Remove legacy cohere.Client and generate() API path entirely - Unify all requests to use cohere.ClientV2 with chat() API - Remove COHERE_GENERATION_LIMIT and sub-batching in _call_model (base class generate() handles iteration for supports_multiple_generations=False) - Remove v1-only parameters (stop, preset, api_version) - Add DeprecationWarning when api_version is set in config - Remove v1-specific tests and fixtures, add deprecation warning test Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Contributor
Author
|
Thanks for the thorough review! I've addressed all feedback in two commits: Commit 1 (
Commit 2 (
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes #1384
Addresses all seven checklist items in #1384.
nameparameter is now required; passing no model name raisesTargetNameMissingError. The deprecatedcommanddefault has been removed.cohere.errors.NotFoundErrorand raiseBadGeneratorExceptionto stop execution immediately, following the pattern innvcf.py.cohere.errors: Replaced the broadRuntimeErrorcatch in the v1 path with fine-grainedcohere.errorstypes (NotFoundError,GatewayTimeoutError,TooManyRequestsError,ServiceUnavailableError,InternalServerError). The v2 path already used these.test_live_chat_api_output_format— verifies that the v2 chat API returns aMessageobject with non-empty text. Skipped whenCOHERE_API_KEYis not set.test_live_v1_api_rejects_new_model— Cohere has retired the generate endpoint and all current models return 404 on v1, so the test verifies that the 404 handling correctly raisesBadGeneratorExceptionrather than silently returningNone.Nones: Error paths continue to emitNoneper the existing protocol._call_modelwrapsstrresults inMessagewhile passingNonethrough unchanged._load_clientpattern: Extracted client initialization into_load_client(), called from__init__and__setstate__. Added_unsafe_attributes = ["generator"]so pickle round-trips correctly restore the client.Test plan
pytest tests/generators/test_cohere.py— 17 tests (15 mock + 2 live)ruff check garak/generators/cohere.py tests/generators/test_cohere.pyCOHERE_API_KEY; safely skipped in CI when absent