From 3acad270e57d446673137edf9409ab5b773661d5 Mon Sep 17 00:00:00 2001 From: Mehmet Bektas Date: Sun, 5 May 2024 19:44:25 -0700 Subject: [PATCH 1/4] support sync ollama embeddings --- litellm/llms/ollama.py | 20 ++++++++++++++++++++ litellm/main.py | 20 ++++++++++---------- 2 files changed, 30 insertions(+), 10 deletions(-) diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py index 5972d9e8c24..f4c06dbe640 100644 --- a/litellm/llms/ollama.py +++ b/litellm/llms/ollama.py @@ -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) + ) diff --git a/litellm/main.py b/litellm/main.py index 8717af57042..9aaf30f0514 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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, From 64a64c68c5d2e511cd1c8db1b7840f09cffca12d Mon Sep 17 00:00:00 2001 From: Mehmet Bektas Date: Sun, 5 May 2024 20:25:43 -0700 Subject: [PATCH 2/4] update ollama test file --- litellm/tests/test_ollama_local.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/litellm/tests/test_ollama_local.py b/litellm/tests/test_ollama_local.py index d4dbc7341b0..f5d629140e4 100644 --- a/litellm/tests/test_ollama_local.py +++ b/litellm/tests/test_ollama_local.py @@ -24,6 +24,14 @@ # asyncio.run(test_ollama_aembeddings()) +# def test_ollama_embeddings(): +# litellm.set_verbose = True +# input = "The food was delicious and the waiter..." +# response = litellm.embedding(model="ollama/mistral", input=input) +# print(response) + +# test_ollama_embeddings() + # def test_ollama_streaming(): # try: # litellm.set_verbose = False From 157d7a7c28caac70917c8d101636bba0d76afdd5 Mon Sep 17 00:00:00 2001 From: Mehmet Bektas Date: Mon, 6 May 2024 20:11:45 -0700 Subject: [PATCH 3/4] add ollama embeddings unit tests --- litellm/tests/test_ollama.py | 52 ++++++++++++++++++++++++++++++++++-- 1 file changed, 50 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_ollama.py b/litellm/tests/test_ollama.py index 82ec16f0edf..cc33ea9939b 100644 --- a/litellm/tests/test_ollama.py +++ b/litellm/tests/test_ollama.py @@ -1,3 +1,4 @@ +import asyncio import sys, os import traceback from dotenv import load_dotenv @@ -10,10 +11,10 @@ sys.path.insert( ) # Adds the parent directory to the system path import pytest import litellm - +from unittest import mock ## for ollama we can't test making the completion call -from litellm.utils import get_optional_params, get_llm_provider +from litellm.utils import EmbeddingResponse, get_optional_params, get_llm_provider def test_get_ollama_params(): @@ -58,3 +59,50 @@ def test_ollama_json_mode(): except Exception as e: pytest.fail(f"Error occurred: {e}") # test_ollama_json_mode() + + +mock_ollama_embedding_response = EmbeddingResponse(model="ollama/nomic-embed-text") + +@mock.patch( + "litellm.llms.ollama.ollama_embeddings", + return_value=mock_ollama_embedding_response, +) +def test_ollama_embeddings(mock_embeddings): + # assert that ollama_embeddings is called with the right parameters + try: + embeddings = litellm.embedding(model="ollama/nomic-embed-text", input=["hello world"]) + print(embeddings) + mock_embeddings.assert_called_once_with( + api_base="http://localhost:11434", + model="nomic-embed-text", + prompts=["hello world"], + optional_params=mock.ANY, + logging_obj=mock.ANY, + model_response=mock.ANY, + encoding=mock.ANY, + ) + except Exception as e: + pytest.fail(f"Error occurred: {e}") +test_ollama_embeddings() + +@mock.patch( + "litellm.llms.ollama.ollama_aembeddings", + return_value=mock_ollama_embedding_response, +) +def test_ollama_aembeddings(mock_aembeddings): + # assert that ollama_aembeddings is called with the right parameters + try: + embeddings = asyncio.run(litellm.aembedding(model="ollama/nomic-embed-text", input=["hello world"])) + print(embeddings) + mock_aembeddings.assert_called_once_with( + api_base="http://localhost:11434", + model="nomic-embed-text", + prompts=["hello world"], + optional_params=mock.ANY, + logging_obj=mock.ANY, + model_response=mock.ANY, + encoding=mock.ANY, + ) + except Exception as e: + pytest.fail(f"Error occurred: {e}") +test_ollama_aembeddings() From 8e9a4fa9eb97068354d182c2de3add7c2fccf782 Mon Sep 17 00:00:00 2001 From: Mehmet Bektas Date: Mon, 6 May 2024 20:13:11 -0700 Subject: [PATCH 4/4] comment out test method calls, following the pattern --- litellm/tests/test_ollama.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_ollama.py b/litellm/tests/test_ollama.py index cc33ea9939b..77a6c91c3e3 100644 --- a/litellm/tests/test_ollama.py +++ b/litellm/tests/test_ollama.py @@ -83,7 +83,7 @@ def test_ollama_embeddings(mock_embeddings): ) except Exception as e: pytest.fail(f"Error occurred: {e}") -test_ollama_embeddings() +# test_ollama_embeddings() @mock.patch( "litellm.llms.ollama.ollama_aembeddings", @@ -105,4 +105,4 @@ def test_ollama_aembeddings(mock_aembeddings): ) except Exception as e: pytest.fail(f"Error occurred: {e}") -test_ollama_aembeddings() +# test_ollama_aembeddings()