Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Feat] support dimensions OpenAI embedding param #1635

Merged
merged 7 commits into from
Jan 26, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 13 additions & 4 deletions docs/my-website/docs/embedding/supported_embedding.md
Original file line number Diff line number Diff line change
Expand Up @@ -13,16 +13,20 @@ response = embedding(model='text-embedding-ada-002', input=["good morning from l

- `model`: *string* - ID of the model to use. `model='text-embedding-ada-002'`

- `input`: *array* - Input text to embed, encoded as a string or array of tokens. To embed multiple inputs in a single request, pass an array of strings or array of token arrays. The input must not exceed the max input tokens for the model (8192 tokens for text-embedding-ada-002), cannot be an empty string, and any array must be 2048 dimensions or less.
```
- `input`: *string or array* - Input text to embed, encoded as a string or array of tokens. To embed multiple inputs in a single request, pass an array of strings or array of token arrays. The input must not exceed the max input tokens for the model (8192 tokens for text-embedding-ada-002), cannot be an empty string, and any array must be 2048 dimensions or less.
```python
input=["good morning from litellm"]
```

### Optional LiteLLM Fields

- `user`: *string (optional)* A unique identifier representing your end-user,

- `timeout`: *integer* - The maximum time, in seconds, to wait for the API to respond. Defaults to 600 seconds (10 minutes).
- `dimensions`: *integer (Optional)* The number of dimensions the resulting output embeddings should have. Only supported in OpenAI/Azure text-embedding-3 and later models.

- `encoding_format`: *string (Optional)* The format to return the embeddings in. Can be either `"float"` or `"base64"`. Defaults to `encoding_format="float"`

- `timeout`: *integer (Optional)* - The maximum time, in seconds, to wait for the API to respond. Defaults to 600 seconds (10 minutes).

- `api_base`: *string (optional)* - The api endpoint you want to call the model with

Expand Down Expand Up @@ -66,7 +70,12 @@ input=["good morning from litellm"]
from litellm import embedding
import os
os.environ['OPENAI_API_KEY'] = ""
response = embedding('text-embedding-ada-002', input=["good morning from litellm"])
response = embedding(
model="text-embedding-3-small",
input=["good morning from litellm", "this is another item"],
metadata={"anything": "good day"},
dimensions=5 # Only supported in text-embedding-3 and later models.
)
```

| Model Name | Function Call | Required OS Variables |
Expand Down
5 changes: 5 additions & 0 deletions litellm/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -2224,6 +2224,7 @@ def embedding(
model,
input=[],
# Optional params
dimensions: Optional[int] = None,
timeout=600, # default to 10 minutes
# set api_base, api_version, api_key
api_base: Optional[str] = None,
Expand All @@ -2244,6 +2245,7 @@ def embedding(
Parameters:
- model: The embedding model to use.
- input: The input for which embeddings are to be generated.
- dimensions: The number of dimensions the resulting output embeddings should have. Only supported in text-embedding-3 and later models.
- timeout: The timeout value for the API call, default 10 mins
- litellm_call_id: The call ID for litellm logging.
- litellm_logging_obj: The litellm logging object.
Expand Down Expand Up @@ -2277,6 +2279,7 @@ def embedding(
output_cost_per_second = kwargs.get("output_cost_per_second", None)
openai_params = [
"user",
"dimensions",
"request_timeout",
"api_base",
"api_version",
Expand Down Expand Up @@ -2345,7 +2348,9 @@ def embedding(
api_key=api_key,
)
optional_params = get_optional_params_embeddings(
model=model,
user=user,
dimensions=dimensions,
encoding_format=encoding_format,
custom_llm_provider=custom_llm_provider,
**non_default_params,
Expand Down
3 changes: 3 additions & 0 deletions litellm/tests/test_embedding.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,9 @@ def test_openai_embedding_3():
model="text-embedding-3-small",
input=["good morning from litellm", "this is another item"],
metadata={"anything": "good day"},
dimensions=5,
)
print(f"response:", response)
litellm_response = dict(response)
litellm_response_keys = set(litellm_response.keys())
litellm_response_keys.discard("_response_ms")
Expand All @@ -80,6 +82,7 @@ def test_openai_embedding_3():
response = client.embeddings.create(
model="text-embedding-3-small",
input=["good morning from litellm", "this is another item"],
dimensions=5,
)

response = dict(response)
Expand Down
17 changes: 16 additions & 1 deletion litellm/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -3313,8 +3313,10 @@ def get_optional_params_image_gen(

def get_optional_params_embeddings(
# 2 optional params
model=None,
user=None,
encoding_format=None,
dimensions=None,
custom_llm_provider="",
**kwargs,
):
Expand All @@ -3325,14 +3327,27 @@ def get_optional_params_embeddings(
for k, v in special_params.items():
passed_params[k] = v

default_params = {"user": None, "encoding_format": None}
default_params = {"user": None, "encoding_format": None, "dimensions": None}

non_default_params = {
k: v
for k, v in passed_params.items()
if (k in default_params and v != default_params[k])
}
## raise exception if non-default value passed for non-openai/azure embedding calls
if custom_llm_provider == "openai":
# 'dimensions` is only supported in `text-embedding-3` and later models

if (
model is not None
and "text-embedding-3" not in model
and "dimensions" in non_default_params.keys()
):
raise UnsupportedParamsError(
status_code=500,
message=f"Setting dimensions is not supported for OpenAI `text-embedding-3` and later models. To drop it from the call, set `litellm.drop_params = True`.",
)

if (
custom_llm_provider != "openai"
and custom_llm_provider != "azure"
Expand Down
6 changes: 3 additions & 3 deletions poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.