From 58ac2a7e2b7e7d81e0cd69bde98cc5906ac781b8 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 9 Jul 2024 12:38:34 -0700 Subject: [PATCH] docs(supported_embeddings.md): add doc on provider-specific params for embedding models --- .../docs/embedding/supported_embedding.md | 70 +++++++++++++++++++ docs/my-website/docs/providers/cohere.md | 4 +- litellm/tests/test_embedding.py | 41 ++++++----- 3 files changed, 96 insertions(+), 19 deletions(-) diff --git a/docs/my-website/docs/embedding/supported_embedding.md b/docs/my-website/docs/embedding/supported_embedding.md index ebf7a29ebc3..2782bd26091 100644 --- a/docs/my-website/docs/embedding/supported_embedding.md +++ b/docs/my-website/docs/embedding/supported_embedding.md @@ -85,6 +85,17 @@ print(query_result[:5]) ## Input Params for `litellm.embedding()` + + +:::info + +Any non-openai params, will be treated as provider-specific params, and sent in the request body as kwargs to the provider. + +[**See Reserved Params**](https://github.com/BerriAI/litellm/blob/2f5f85cb52f36448d1f8bbfbd3b8af8167d0c4c8/litellm/main.py#L3130) + +[**See Example**](#example) +::: + ### Required Fields - `model`: *string* - ID of the model to use. `model='text-embedding-ada-002'` @@ -363,3 +374,62 @@ All models listed here https://docs.voyageai.com/embeddings/#models-and-specific | voyage-01 | `embedding(model="voyage/voyage-01", input)` | | voyage-lite-01 | `embedding(model="voyage/voyage-lite-01", input)` | | voyage-lite-01-instruct | `embedding(model="voyage/voyage-lite-01-instruct", input)` | + +## Provider-specific Params + + +:::info + +Any non-openai params, will be treated as provider-specific params, and sent in the request body as kwargs to the provider. + +[**See Reserved Params**](https://github.com/BerriAI/litellm/blob/2f5f85cb52f36448d1f8bbfbd3b8af8167d0c4c8/litellm/main.py#L3130) +::: + +### **Example** + +Cohere v3 Models have a required parameter: `input_type`, it can be one of the following four values: + +- `input_type="search_document"`: (default) Use this for texts (documents) you want to store in your vector database +- `input_type="search_query"`: Use this for search queries to find the most relevant documents in your vector database +- `input_type="classification"`: Use this if you use the embeddings as an input for a classification system +- `input_type="clustering"`: Use this if you use the embeddings for text clustering + +https://txt.cohere.com/introducing-embed-v3/ + + + + +```python +from litellm import embedding +os.environ["COHERE_API_KEY"] = "cohere key" + +# cohere call +response = embedding( + model="embed-english-v3.0", + input=["good morning from litellm", "this is another item"], + input_type="search_document" # 👈 PROVIDER-SPECIFIC PARAM +) +``` + + + +**via config** + +```yaml +model_list: + - model_name: "cohere-embed" + litellm_params: + model: embed-english-v3.0 + input_type: search_document # 👈 PROVIDER-SPECIFIC PARAM +``` + +**via request** + +```bash +curl -X POST 'http://0.0.0.0:4000/v1/embeddings' \ +-H 'Authorization: Bearer sk-54d77cd67b9febbb' \ +-H 'Content-Type: application/json' \ +-D '{"input": ["Are you authorized to work in United States of America?"], "model": "cohere-embed", "input_type": "search_document}' +``` + + \ No newline at end of file diff --git a/docs/my-website/docs/providers/cohere.md b/docs/my-website/docs/providers/cohere.md index a37cdb8aad9..42468c30527 100644 --- a/docs/my-website/docs/providers/cohere.md +++ b/docs/my-website/docs/providers/cohere.md @@ -68,7 +68,7 @@ response = embedding( ``` ### Setting - Input Type for v3 models -v3 Models have a required parameter: `input_type`, it can be one of the following four values: +v3 Models have a required parameter: `input_type`. LiteLLM defaults to `search_document`. It can be one of the following four values: - `input_type="search_document"`: (default) Use this for texts (documents) you want to store in your vector database - `input_type="search_query"`: Use this for search queries to find the most relevant documents in your vector database @@ -76,6 +76,8 @@ v3 Models have a required parameter: `input_type`, it can be one of the followin - `input_type="clustering"`: Use this if you use the embeddings for text clustering https://txt.cohere.com/introducing-embed-v3/ + + ```python from litellm import embedding os.environ["COHERE_API_KEY"] = "cohere key" diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index 2bfb6d47496..1e3f5455a92 100644 --- a/litellm/tests/test_embedding.py +++ b/litellm/tests/test_embedding.py @@ -1,18 +1,21 @@ -import sys, os +import os +import sys import traceback + +import openai import pytest from dotenv import load_dotenv -import openai load_dotenv() sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path -import litellm -from litellm import embedding, completion, completion_cost from unittest.mock import MagicMock, patch +import litellm +from litellm import completion, completion_cost, embedding + litellm.set_verbose = False @@ -232,6 +235,7 @@ def test_cohere_embedding(): response = embedding( model="embed-english-v2.0", input=["good morning from litellm", "this is another item"], + input_type="search_query", ) print(f"response:", response) @@ -486,16 +490,16 @@ def test_mistral_embeddings(): def test_watsonx_embeddings(): - - def mock_wx_embed_request(method:str, url:str, **kwargs): + + def mock_wx_embed_request(method: str, url: str, **kwargs): mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = {"Content-Type": "application/json"} mock_response.json.return_value = { - "model_id": "ibm/slate-30m-english-rtrvr", - "created_at": "2024-01-01T00:00:00.00Z", - "results": [ {"embedding": [0.0]*254} ], - "input_token_count": 8 + "model_id": "ibm/slate-30m-english-rtrvr", + "created_at": "2024-01-01T00:00:00.00Z", + "results": [{"embedding": [0.0] * 254}], + "input_token_count": 8, } return mock_response @@ -505,7 +509,7 @@ def test_watsonx_embeddings(): response = litellm.embedding( model="watsonx/ibm/slate-30m-english-rtrvr", input=["good morning from litellm"], - token="secret-token" + token="secret-token", ) print(f"response: {response}") assert isinstance(response.usage, litellm.Usage) @@ -514,26 +518,27 @@ def test_watsonx_embeddings(): except Exception as e: pytest.fail(f"Error occurred: {e}") + @pytest.mark.asyncio async def test_watsonx_aembeddings(): def mock_async_client(*args, **kwargs): mocked_client = MagicMock() - + async def mock_send(request, *args, stream: bool = False, **kwags): mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = {"Content-Type": "application/json"} mock_response.json.return_value = { - "model_id": "ibm/slate-30m-english-rtrvr", - "created_at": "2024-01-01T00:00:00.00Z", - "results": [ {"embedding": [0.0]*254} ], - "input_token_count": 8 + "model_id": "ibm/slate-30m-english-rtrvr", + "created_at": "2024-01-01T00:00:00.00Z", + "results": [{"embedding": [0.0] * 254}], + "input_token_count": 8, } mock_response.is_error = False return mock_response - + mocked_client.send = mock_send return mocked_client @@ -544,7 +549,7 @@ async def test_watsonx_aembeddings(): response = await litellm.aembedding( model="watsonx/ibm/slate-30m-english-rtrvr", input=["good morning from litellm"], - token="secret-token" + token="secret-token", ) print(f"response: {response}") assert isinstance(response.usage, litellm.Usage)