forked from phoenix/litellm-mirror
(feat) add async embeddings
This commit is contained in:
parent
34c75b8c0c
commit
6373f6bddd
2 changed files with 31 additions and 0 deletions
|
@ -1447,6 +1447,21 @@ def batch_completion_models_all_responses(*args, **kwargs):
|
||||||
return responses
|
return responses
|
||||||
|
|
||||||
### EMBEDDING ENDPOINTS ####################
|
### EMBEDDING ENDPOINTS ####################
|
||||||
|
|
||||||
|
async def aembedding(*args, **kwargs):
|
||||||
|
loop = asyncio.get_event_loop()
|
||||||
|
|
||||||
|
# Use a partial function to pass your keyword arguments
|
||||||
|
func = partial(embedding, *args, **kwargs)
|
||||||
|
|
||||||
|
# Add the context to the function
|
||||||
|
ctx = contextvars.copy_context()
|
||||||
|
func_with_context = partial(ctx.run, func)
|
||||||
|
|
||||||
|
# Call the synchronous function using run_in_executor
|
||||||
|
response = await loop.run_in_executor(None, func_with_context)
|
||||||
|
return response
|
||||||
|
|
||||||
@client
|
@client
|
||||||
@timeout( # type: ignore
|
@timeout( # type: ignore
|
||||||
60
|
60
|
||||||
|
|
|
@ -91,4 +91,20 @@ def test_cohere_embedding():
|
||||||
# pytest.fail(f"Error occurred: {e}")
|
# pytest.fail(f"Error occurred: {e}")
|
||||||
# test_hf_embedding()
|
# test_hf_embedding()
|
||||||
|
|
||||||
|
# test async embeddings
|
||||||
|
def test_aembedding():
|
||||||
|
import asyncio
|
||||||
|
async def embedding_call():
|
||||||
|
try:
|
||||||
|
response = await litellm.aembedding(
|
||||||
|
model="text-embedding-ada-002",
|
||||||
|
input=["good morning from litellm", "this is another item"]
|
||||||
|
)
|
||||||
|
print(response)
|
||||||
|
except:
|
||||||
|
print(f"error occurred: {traceback.format_exc()}")
|
||||||
|
pass
|
||||||
|
asyncio.run(embedding_call())
|
||||||
|
|
||||||
|
# test_aembedding()
|
||||||
|
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue