fix: make LiteLLM SDK embeddings encoding_format configurable (#928)
* fix: make LiteLLM SDK embeddings encoding_format configurable (#925) The hardcoded encoding_format='float' breaks providers like Voyage AI (only accepts 'base64') and Gemini (doesn't support the parameter at all). Add HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT config option that defaults to 'float' for backwards compatibility. Set to empty string to omit the parameter for incompatible providers. * chore: regenerate docs skill after configuration change
This commit is contained in:
parent
d7c73f4342
commit
cece2c903c
5 changed files with 76 additions and 2 deletions
|
|
@ -211,6 +211,7 @@ ENV_EMBEDDINGS_LITELLM_SDK_API_KEY = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_K
|
|||
ENV_EMBEDDINGS_LITELLM_SDK_MODEL = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MODEL"
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_API_BASE = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_BASE"
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS"
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT"
|
||||
ENV_RERANKER_LITELLM_SDK_API_KEY = "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_KEY"
|
||||
ENV_RERANKER_LITELLM_SDK_MODEL = "HINDSIGHT_API_RERANKER_LITELLM_SDK_MODEL"
|
||||
ENV_RERANKER_LITELLM_SDK_API_BASE = "HINDSIGHT_API_RERANKER_LITELLM_SDK_API_BASE"
|
||||
|
|
@ -460,6 +461,7 @@ DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC: int | None = None
|
|||
|
||||
# LiteLLM SDK defaults
|
||||
DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL = "cohere/embed-english-v3.0"
|
||||
DEFAULT_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT = "float"
|
||||
DEFAULT_RERANKER_LITELLM_SDK_MODEL = "cohere/rerank-english-v3.0"
|
||||
|
||||
DEFAULT_HOST = "0.0.0.0"
|
||||
|
|
@ -729,6 +731,7 @@ class HindsightConfig:
|
|||
embeddings_litellm_sdk_model: str
|
||||
embeddings_litellm_sdk_api_base: str | None
|
||||
embeddings_litellm_sdk_output_dimensions: int | None
|
||||
embeddings_litellm_sdk_encoding_format: str | None
|
||||
# Gemini/Vertex AI embeddings
|
||||
embeddings_gemini_api_key: str | None
|
||||
embeddings_gemini_model: str
|
||||
|
|
@ -1200,6 +1203,9 @@ class HindsightConfig:
|
|||
embeddings_litellm_sdk_output_dimensions=int(v)
|
||||
if (v := os.getenv(ENV_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS))
|
||||
else None,
|
||||
embeddings_litellm_sdk_encoding_format=os.getenv(
|
||||
ENV_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT, DEFAULT_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT
|
||||
),
|
||||
# Gemini/Vertex AI embeddings (with fallback to LLM keys)
|
||||
embeddings_gemini_api_key=os.getenv(ENV_EMBEDDINGS_GEMINI_API_KEY) or os.getenv(ENV_LLM_API_KEY),
|
||||
embeddings_gemini_model=os.getenv(ENV_EMBEDDINGS_GEMINI_MODEL, DEFAULT_EMBEDDINGS_GEMINI_MODEL),
|
||||
|
|
|
|||
|
|
@ -757,6 +757,7 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
|||
output_dimensions: int | None = None,
|
||||
batch_size: int = 100,
|
||||
timeout: float = 60.0,
|
||||
encoding_format: str | None = "float",
|
||||
):
|
||||
"""
|
||||
Initialize LiteLLM SDK embeddings client.
|
||||
|
|
@ -768,6 +769,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
|||
output_dimensions: Optional output embedding dimensions (provider-dependent)
|
||||
batch_size: Maximum batch size for embedding requests (default: 100)
|
||||
timeout: Request timeout in seconds (default: 60.0)
|
||||
encoding_format: Encoding format for embeddings (default: "float").
|
||||
Set to None or empty string to omit (needed for Voyage AI, Gemini).
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
|
|
@ -775,6 +778,7 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
|||
self.output_dimensions = output_dimensions
|
||||
self.batch_size = batch_size
|
||||
self.timeout = timeout
|
||||
self.encoding_format = encoding_format or None
|
||||
self._litellm = None # Will be set during initialization
|
||||
self._dimension: int | None = None
|
||||
|
||||
|
|
@ -810,8 +814,9 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
|||
"model": self.model,
|
||||
"input": ["test"],
|
||||
"api_key": self.api_key,
|
||||
"encoding_format": "float",
|
||||
}
|
||||
if self.encoding_format:
|
||||
embed_kwargs["encoding_format"] = self.encoding_format
|
||||
if self.api_base:
|
||||
embed_kwargs["api_base"] = self.api_base
|
||||
if self.output_dimensions is not None:
|
||||
|
|
@ -859,8 +864,9 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
|||
"model": self.model,
|
||||
"input": batch,
|
||||
"api_key": self.api_key,
|
||||
"encoding_format": "float",
|
||||
}
|
||||
if self.encoding_format:
|
||||
embed_kwargs["encoding_format"] = self.encoding_format
|
||||
if self.api_base:
|
||||
embed_kwargs["api_base"] = self.api_base
|
||||
if self.output_dimensions is not None:
|
||||
|
|
@ -1121,6 +1127,7 @@ def create_embeddings_from_env() -> Embeddings:
|
|||
model=config.embeddings_litellm_sdk_model,
|
||||
api_base=config.embeddings_litellm_sdk_api_base,
|
||||
output_dimensions=config.embeddings_litellm_sdk_output_dimensions,
|
||||
encoding_format=config.embeddings_litellm_sdk_encoding_format,
|
||||
)
|
||||
elif provider == "google":
|
||||
vertexai_project_id = config.embeddings_vertexai_project_id
|
||||
|
|
|
|||
|
|
@ -345,6 +345,65 @@ class TestLiteLLMSDKEmbeddings:
|
|||
assert encode_call_args.kwargs["api_base"] == "https://custom.api.com"
|
||||
assert encode_call_args.kwargs["dimensions"] == 768
|
||||
|
||||
async def test_encoding_format_default_is_float(self, mock_litellm):
|
||||
"""Test that encoding_format defaults to 'float' for backwards compatibility."""
|
||||
with patch(
|
||||
"builtins.__import__",
|
||||
side_effect=lambda name, *args: mock_litellm if name == "litellm" else __import__(name, *args),
|
||||
):
|
||||
emb = LiteLLMSDKEmbeddings(
|
||||
api_key="test_key",
|
||||
model="cohere/embed-english-v3.0",
|
||||
)
|
||||
await emb.initialize()
|
||||
|
||||
init_call_args = mock_litellm.aembedding.call_args
|
||||
assert init_call_args.kwargs["encoding_format"] == "float"
|
||||
|
||||
mock_litellm.embedding.return_value.data = [{"embedding": [0.1] * 768, "index": 0}]
|
||||
emb.encode(["test"])
|
||||
|
||||
encode_call_args = mock_litellm.embedding.call_args
|
||||
assert encode_call_args.kwargs["encoding_format"] == "float"
|
||||
|
||||
async def test_encoding_format_omitted_when_none(self, mock_litellm):
|
||||
"""Test that encoding_format is omitted when set to None (for Voyage AI, Gemini)."""
|
||||
with patch(
|
||||
"builtins.__import__",
|
||||
side_effect=lambda name, *args: mock_litellm if name == "litellm" else __import__(name, *args),
|
||||
):
|
||||
emb = LiteLLMSDKEmbeddings(
|
||||
api_key="test_key",
|
||||
model="voyage/voyage-4-large",
|
||||
encoding_format=None,
|
||||
)
|
||||
await emb.initialize()
|
||||
|
||||
init_call_args = mock_litellm.aembedding.call_args
|
||||
assert "encoding_format" not in init_call_args.kwargs
|
||||
|
||||
mock_litellm.embedding.return_value.data = [{"embedding": [0.1] * 768, "index": 0}]
|
||||
emb.encode(["test"])
|
||||
|
||||
encode_call_args = mock_litellm.embedding.call_args
|
||||
assert "encoding_format" not in encode_call_args.kwargs
|
||||
|
||||
async def test_encoding_format_omitted_when_empty_string(self, mock_litellm):
|
||||
"""Test that encoding_format is omitted when set to empty string."""
|
||||
with patch(
|
||||
"builtins.__import__",
|
||||
side_effect=lambda name, *args: mock_litellm if name == "litellm" else __import__(name, *args),
|
||||
):
|
||||
emb = LiteLLMSDKEmbeddings(
|
||||
api_key="test_key",
|
||||
model="gemini/gemini-embedding-2-preview",
|
||||
encoding_format="",
|
||||
)
|
||||
await emb.initialize()
|
||||
|
||||
init_call_args = mock_litellm.aembedding.call_args
|
||||
assert "encoding_format" not in init_call_args.kwargs
|
||||
|
||||
async def test_openai_invalid_output_dimensions_raises(self, mock_litellm):
|
||||
"""Invalid dimensions fail during initialize() (probe call), not per HTTP request.
|
||||
|
||||
|
|
|
|||
|
|
@ -370,6 +370,7 @@ export HINDSIGHT_API_RETAIN_LLM_MAX_BACKOFF=120.0 # Cap at 2min instead of 1m
|
|||
| `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MODEL` | LiteLLM SDK embedding model (use provider prefix, e.g., `cohere/embed-english-v3.0`) | `cohere/embed-english-v3.0` |
|
||||
| `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_BASE` | Custom base URL for LiteLLM SDK embeddings (optional) | - |
|
||||
| `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS` | Optional output embedding dimensions (provider-dependent, e.g., `768` for Gemini embedding models) | - |
|
||||
| `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT` | Encoding format for embedding responses. Set to empty string to omit the parameter (needed for Voyage AI, Gemini). | `float` |
|
||||
| `HINDSIGHT_API_EMBEDDINGS_GEMINI_API_KEY` | Gemini API key for embeddings (falls back to `HINDSIGHT_API_LLM_API_KEY`) | - |
|
||||
| `HINDSIGHT_API_EMBEDDINGS_GEMINI_MODEL` | Gemini embedding model | `gemini-embedding-001` |
|
||||
| `HINDSIGHT_API_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY` | Output embedding dimensions (Gemini supports configurable dimensionality) | `768` |
|
||||
|
|
|
|||
|
|
@ -370,6 +370,7 @@ export HINDSIGHT_API_RETAIN_LLM_MAX_BACKOFF=120.0 # Cap at 2min instead of 1m
|
|||
| `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MODEL` | LiteLLM SDK embedding model (use provider prefix, e.g., `cohere/embed-english-v3.0`) | `cohere/embed-english-v3.0` |
|
||||
| `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_BASE` | Custom base URL for LiteLLM SDK embeddings (optional) | - |
|
||||
| `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS` | Optional output embedding dimensions (provider-dependent, e.g., `768` for Gemini embedding models) | - |
|
||||
| `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_ENCODING_FORMAT` | Encoding format for embedding responses. Set to empty string to omit the parameter (needed for Voyage AI, Gemini). | `float` |
|
||||
| `HINDSIGHT_API_EMBEDDINGS_GEMINI_API_KEY` | Gemini API key for embeddings (falls back to `HINDSIGHT_API_LLM_API_KEY`) | - |
|
||||
| `HINDSIGHT_API_EMBEDDINGS_GEMINI_MODEL` | Gemini embedding model | `gemini-embedding-001` |
|
||||
| `HINDSIGHT_API_EMBEDDINGS_GEMINI_OUTPUT_DIMENSIONALITY` | Output embedding dimensions (Gemini supports configurable dimensionality) | `768` |
|
||||
|
|
|
|||
Loading…
Reference in a new issue