fix(cohere.py): support async cohere embedding calls

This commit is contained in:
Krrish Dholakia 2024-07-30 14:49:07 -07:00
parent 185a6857f9
commit 9b2eb1702b
3 changed files with 132 additions and 30 deletions

View file

@ -3114,6 +3114,7 @@ async def aembedding(*args, **kwargs) -> EmbeddingResponse:
or custom_llm_provider == "vertex_ai"
or custom_llm_provider == "databricks"
or custom_llm_provider == "watsonx"
or custom_llm_provider == "cohere"
): # currently implemented aiohttp calls for just azure and openai, soon all.
# Await normally
init_response = await loop.run_in_executor(None, func_with_context)
@ -3440,9 +3441,12 @@ def embedding(
input=input,
optional_params=optional_params,
encoding=encoding,
api_key=cohere_key,
api_key=cohere_key, # type: ignore
logging_obj=logging,
model_response=EmbeddingResponse(),
aembedding=aembedding,
timeout=float(timeout),
client=client,
)
elif custom_llm_provider == "huggingface":
api_key = (