[Feat] Add infinity embedding support (contributor pr) (#10196)

* Feature - infinity support for #8764 (#10009)

* Added support for infinity embeddings

* Added test cases

* Fixed tests and api base

* Updated docs and tests

* Removed unused import

* Updated signature

* Added support for infinity embeddings

* Added test cases

* Fixed tests and api base

* Updated docs and tests

* Removed unused import

* Updated signature

* Updated validate params

---------

Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com>

* fix InfinityEmbeddingConfig

---------

Co-authored-by: Prathamesh Saraf <pratamesh1867@gmail.com>
This commit is contained in:
Ishaan Jaff 2025-04-21 20:01:29 -07:00 committed by GitHub
parent 0c2f705417
commit 104e4cb1bc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 529 additions and 22 deletions

View file

@ -2735,6 +2735,21 @@ def get_optional_params_embeddings( # noqa: PLR0915
)
final_params = {**optional_params, **kwargs}
return final_params
elif custom_llm_provider == "infinity":
supported_params = get_supported_openai_params(
model=model,
custom_llm_provider="infinity",
request_type="embeddings",
)
_check_valid_arg(supported_params=supported_params)
optional_params = litellm.InfinityEmbeddingConfig().map_openai_params(
non_default_params=non_default_params,
optional_params={},
model=model,
drop_params=drop_params if drop_params is not None else False,
)
final_params = {**optional_params, **kwargs}
return final_params
elif custom_llm_provider == "fireworks_ai":
supported_params = get_supported_openai_params(
model=model,
@ -5120,6 +5135,11 @@ def validate_environment( # noqa: PLR0915
keys_in_environment = True
else:
missing_keys.append("VOYAGE_API_KEY")
elif custom_llm_provider == "infinity":
if "INFINITY_API_KEY" in os.environ:
keys_in_environment = True
else:
missing_keys.append("INFINITY_API_KEY")
elif custom_llm_provider == "fireworks_ai":
if (
"FIREWORKS_AI_API_KEY" in os.environ
@ -6554,6 +6574,8 @@ class ProviderConfigManager:
return litellm.TritonEmbeddingConfig()
elif litellm.LlmProviders.WATSONX == provider:
return litellm.IBMWatsonXEmbeddingConfig()
elif litellm.LlmProviders.INFINITY == provider:
return litellm.InfinityEmbeddingConfig()
raise ValueError(f"Provider {provider.value} does not support embedding config")
@staticmethod