diff --git a/litellm/llms/ollama.py b/litellm/llms/ollama.py index 5a0b450c75f..49e2f3d93be 100644 --- a/litellm/llms/ollama.py +++ b/litellm/llms/ollama.py @@ -253,3 +253,75 @@ async def ollama_acompletion(url, data, model_response, encoding, logging_obj): except Exception as e: traceback.print_exc() raise e + + +async def ollama_aembeddings(api_base="http://localhost:11434", + model="llama2", + prompt="Why is the sky blue?", + optional_params=None, + logging_obj=None, + model_response=None, + encoding=None): + if api_base.endswith("/api/embeddings"): + url = api_base + else: + url = f"{api_base}/api/embeddings" + + ## Load Config + config=litellm.OllamaConfig.get_config() + for k, v in config.items(): + if k not in optional_params: # completion(top_k=3) > cohere_config(top_k=3) <- allows for dynamic variables to be passed in + optional_params[k] = v + + data = { + "model": model, + "prompt": prompt, + } + ## LOGGING + logging_obj.pre_call( + input=None, + api_key=None, + additional_args={"api_base": url, "complete_input_dict": data, "headers": {}}, + ) + timeout = aiohttp.ClientTimeout(total=litellm.request_timeout) # 10 minutes + async with aiohttp.ClientSession(timeout=timeout) as session: + response = await session.post(url, json=data) + + if response.status != 200: + text = await response.text() + raise OllamaError(status_code=response.status, message=text) + + ## LOGGING + logging_obj.post_call( + input=prompt, + api_key="", + original_response=response.text, + additional_args={ + "headers": None, + "api_base": api_base, + }, + ) + + response_json = await response.json() + embeddings = response_json["embedding"] + ## RESPONSE OBJECT + output_data = [] + for idx, embedding in enumerate(embeddings): + output_data.append( + { + "object": "embedding", + "index": idx, + "embedding": embedding + } + ) + model_response["object"] = "list" + model_response["data"] = output_data + model_response["model"] = model + + input_tokens = len(encoding.encode(prompt)) + + model_response["usage"] = { + "prompt_tokens": input_tokens, + "total_tokens": input_tokens, + } + return model_response \ No newline at end of file diff --git a/litellm/main.py b/litellm/main.py index 0e7752e16d5..8d2cabdfb4d 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1749,7 +1749,8 @@ async def aembedding(*args, **kwargs): or custom_llm_provider == "anyscale" or custom_llm_provider == "openrouter" or custom_llm_provider == "deepinfra" - or custom_llm_provider == "perplexity"): # currently implemented aiohttp calls for just azure and openai, soon all. + or custom_llm_provider == "perplexity" + or custom_llm_provider == "ollama"): # currently implemented aiohttp calls for just azure and openai, soon all. # Await normally init_response = await loop.run_in_executor(None, func_with_context) if isinstance(init_response, dict) or isinstance(init_response, ModelResponse): ## CACHING SCENARIO @@ -1949,6 +1950,16 @@ def embedding( optional_params=optional_params, model_response= EmbeddingResponse() ) + elif custom_llm_provider == "ollama": + if aembedding == True: + response = ollama.ollama_aembeddings( + model=model, + prompt=input, + encoding=encoding, + logging_obj=logging, + optional_params=optional_params, + model_response=EmbeddingResponse(), + ) elif custom_llm_provider == "sagemaker": response = sagemaker.embedding( model=model, diff --git a/litellm/tests/test_ollama_local.py b/litellm/tests/test_ollama_local.py index 35be975065c..3eb99697076 100644 --- a/litellm/tests/test_ollama_local.py +++ b/litellm/tests/test_ollama_local.py @@ -16,6 +16,14 @@ # user_message = "respond in 20 words. who are you?" # messages = [{ "content": user_message,"role": "user"}] +# async def test_ollama_aembeddings(): +# litellm.set_verbose = True +# input = "The food was delicious and the waiter..." +# response = await litellm.aembedding(model="ollama/mistral", input=input) +# print(response) + +# asyncio.run(test_ollama_aembeddings()) + # def test_ollama_streaming(): # try: # litellm.set_verbose = False @@ -51,7 +59,7 @@ # except Exception as e: # print(e) -# test_ollama_streaming() +# # test_ollama_streaming() # async def test_async_ollama_streaming(): # try: