Skip to content

Commit

Permalink
fix: Cohere - better handling of COHERE_API_URL (#6407)
Browse files Browse the repository at this point in the history
* extract API URL from lazy import

* improve solution
  • Loading branch information
anakin87 committed Nov 24, 2023
1 parent c45d8c3 commit b850b36
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion haystack/preview/components/generators/cohere.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ def __init__(
api_key: Optional[str] = None,
model_name: str = "command",
streaming_callback: Optional[Callable] = None,
api_base_url: str = COHERE_API_URL,
api_base_url: Optional[str] = None,
**kwargs,
):
"""
Expand Down Expand Up @@ -61,13 +61,18 @@ def __init__(
The format is {token_id: bias} where bias is a float between -10 and 10.
"""
cohere_import.check()

if not api_key:
api_key = os.environ.get("COHERE_API_KEY")
if not api_key:
raise ValueError(
"CohereGenerator needs an API key to run. Either provide it as init parameter or set the env var COHERE_API_KEY."
)

if not api_base_url:
api_base_url = COHERE_API_URL

self.api_key = api_key
self.model_name = model_name
self.streaming_callback = streaming_callback
Expand Down

0 comments on commit b850b36

Please sign in to comment.