yujunjun/haystack
1
0
Fork 0
mirror of https://github.com/deepset-ai/haystack.git synced 2025-11-13 16:43:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: Cohere - better handling of COHERE_API_URL (#6407)

Browse Source 
* extract API URL from lazy import

* improve solution
This commit is contained in:
Stefano Fiorucci 2023-11-24 10:58:46 +01:00 committed by GitHub
parent c45d8c39c7
commit b850b36a4b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
haystack/preview/components/generators/cohere.py
View File

@ -33,7 +33,7 @@ class CohereGenerator:
api_key: Optional[str] = None, api_key: Optional[str] = None,
model_name: str = "command", model_name: str = "command",
streaming_callback: Optional[Callable] = None, streaming_callback: Optional[Callable] = None,
api_base_url: str = COHERE_API_URL, api_base_url: Optional[str] = None,
**kwargs, **kwargs,
): ):
""" """
@ -61,6 +61,8 @@ class CohereGenerator:
The format is {token_id: bias} where bias is a float between -10 and 10. The format is {token_id: bias} where bias is a float between -10 and 10.
""" """
cohere_import.check()
if not api_key: if not api_key:
api_key = os.environ.get("COHERE_API_KEY") api_key = os.environ.get("COHERE_API_KEY")
if not api_key: if not api_key:
@ -68,6 +70,9 @@ class CohereGenerator:
"CohereGenerator needs an API key to run. Either provide it as init parameter or set the env var COHERE_API_KEY." "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.api_key = api_key
self.model_name = model_name self.model_name = model_name
self.streaming_callback = streaming_callback self.streaming_callback = streaming_callback