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

support sync ollama embeddings

Browse source
This commit is contained in:
Mehmet Bektas 2024-05-05 19:44:25 -07:00
parent 918367cc7b
commit 3acad270e5
2 changed files with 30 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

20
litellm/llms/ollama.py
View file

@ -417,3 +417,23 @@ async def ollama_aembeddings(
"total_tokens": total_input_tokens,
}
return model_response
def ollama_embeddings(
api_base: str,
model: str,
prompts: list,
optional_params=None,
logging_obj=None,
model_response=None,
encoding=None,
):
return asyncio.run(
ollama_aembeddings(
api_base,
model,
prompts,
optional_params,
logging_obj,
model_response,
encoding)
)

20
litellm/main.py
View file

@ -2946,16 +2946,16 @@ def embedding(
model=model, # type: ignore
llm_provider="ollama", # type: ignore
)
if aembedding:
response = ollama.ollama_aembeddings(
api_base=api_base,
model=model,
prompts=input,
encoding=encoding,
logging_obj=logging,
optional_params=optional_params,
model_response=EmbeddingResponse(),
)
ollama_embeddings_fn = ollama.ollama_aembeddings if aembedding else ollama.ollama_embeddings
response = ollama_embeddings_fn(
api_base=api_base,
model=model,
prompts=input,
encoding=encoding,
logging_obj=logging,
optional_params=optional_params,
model_response=EmbeddingResponse(),
)
elif custom_llm_provider == "sagemaker":
response = sagemaker.embedding(
model=model,