phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 19:24:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat: cohore_v2 support

Browse source
This commit is contained in:
sajdkabir 2025-04-08 23:17:20 +05:30
parent ff3a6830a4
commit ac9c9bb729
10 changed files with 1480 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

39
litellm/main.py
View file

@ -2108,6 +2108,45 @@ def completion( # type: ignore # noqa: PLR0915
api_key=cohere_key,
logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements
)
elif custom_llm_provider == "cohere_chat_v2":
cohere_key = (
api_key
or litellm.cohere_key
or get_secret_str("COHERE_API_KEY")
or get_secret_str("CO_API_KEY")
or litellm.api_key
)
api_base = (
api_base
or litellm.api_base
or get_secret_str("COHERE_API_BASE")
or "https://api.cohere.ai/v2/chat"
)
headers = headers or litellm.headers or {}
if headers is None:
headers = {}
if extra_headers is not None:
headers.update(extra_headers)
response = base_llm_http_handler.completion(
model=model,
stream=stream,
messages=messages,
acompletion=acompletion,
api_base=api_base,
model_response=model_response,
optional_params=optional_params,
litellm_params=litellm_params,
custom_llm_provider="cohere_chat_v2",
timeout=timeout,
headers=headers,
encoding=encoding,
api_key=cohere_key,
logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements
)
elif custom_llm_provider == "maritalk":
maritalk_key = (
api_key