feat(ollama.py): add support for async ollama embeddings

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

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,

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: