feat: add optional LiteLLM SDK embedding output dimensions (#809)
* feat: add optional LiteLLM SDK embedding output dimensions Allow configuring an optional output dimension for litellm-sdk embeddings and pass it through only when set, while preserving default behavior. Made-with: Cursor * test: assert wrapped init error for invalid dimensions Add a LiteLLM SDK embeddings test that verifies invalid OpenAI dimensions fail during initialize() and preserve provider error details in the wrapped RuntimeError. Made-with: Cursor
This commit is contained in:
Amin Bolakhrif
GitHub
parent
fa82efc886
commit
f841bcb92d
5 changed files with 131 additions and 0 deletions
|
|
@ -200,6 +200,7 @@ ENV_RERANKER_LITELLM_MAX_TOKENS_PER_DOC = "HINDSIGHT_API_RERANKER_LITELLM_MAX_TO
|
|||
ENV_EMBEDDINGS_LITELLM_SDK_API_KEY = "HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_KEY"
|
||||
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_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"
|
||||
|
|
@ -695,6 +696,7 @@ class HindsightConfig:
|
|||
embeddings_litellm_sdk_api_key: str | None
|
||||
embeddings_litellm_sdk_model: str
|
||||
embeddings_litellm_sdk_api_base: str | None
|
||||
embeddings_litellm_sdk_output_dimensions: int | None
|
||||
|
||||
# Reranker
|
||||
reranker_provider: str
|
||||
|
|
@ -1143,6 +1145,9 @@ class HindsightConfig:
|
|||
ENV_EMBEDDINGS_LITELLM_SDK_MODEL, DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL
|
||||
),
|
||||
embeddings_litellm_sdk_api_base=os.getenv(ENV_EMBEDDINGS_LITELLM_SDK_API_BASE) or None,
|
||||
embeddings_litellm_sdk_output_dimensions=int(v)
|
||||
if (v := os.getenv(ENV_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS))
|
||||
else None,
|
||||
# Reranker
|
||||
reranker_provider=os.getenv(ENV_RERANKER_PROVIDER, DEFAULT_RERANKER_PROVIDER),
|
||||
reranker_local_model=os.getenv(ENV_RERANKER_LOCAL_MODEL, DEFAULT_RERANKER_LOCAL_MODEL),
|
||||
|
|
|
|||
|
|
@ -752,6 +752,7 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
|||
api_key: str,
|
||||
model: str = DEFAULT_EMBEDDINGS_LITELLM_SDK_MODEL,
|
||||
api_base: str | None = None,
|
||||
output_dimensions: int | None = None,
|
||||
batch_size: int = 100,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
|
|
@ -762,12 +763,14 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
|||
api_key: API key for the embedding provider
|
||||
model: Model name with provider prefix (e.g., "cohere/embed-english-v3.0")
|
||||
api_base: Custom base URL for API (optional)
|
||||
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)
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.api_base = api_base
|
||||
self.output_dimensions = output_dimensions
|
||||
self.batch_size = batch_size
|
||||
self.timeout = timeout
|
||||
self._litellm = None # Will be set during initialization
|
||||
|
|
@ -809,6 +812,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
|||
}
|
||||
if self.api_base:
|
||||
embed_kwargs["api_base"] = self.api_base
|
||||
if self.output_dimensions is not None:
|
||||
embed_kwargs["dimensions"] = self.output_dimensions
|
||||
|
||||
# Use async embedding method (standard in litellm)
|
||||
response = await self._litellm.aembedding(**embed_kwargs)
|
||||
|
|
@ -856,6 +861,8 @@ class LiteLLMSDKEmbeddings(Embeddings):
|
|||
}
|
||||
if self.api_base:
|
||||
embed_kwargs["api_base"] = self.api_base
|
||||
if self.output_dimensions is not None:
|
||||
embed_kwargs["dimensions"] = self.output_dimensions
|
||||
|
||||
# Use sync embedding (litellm doesn't have async in thread-safe way)
|
||||
response = self._litellm.embedding(**embed_kwargs)
|
||||
|
|
@ -938,6 +945,7 @@ def create_embeddings_from_env() -> Embeddings:
|
|||
api_key=api_key,
|
||||
model=config.embeddings_litellm_sdk_model,
|
||||
api_base=config.embeddings_litellm_sdk_api_base,
|
||||
output_dimensions=config.embeddings_litellm_sdk_output_dimensions,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
|
|
|
|||
|
|
@ -277,6 +277,99 @@ class TestLiteLLMSDKEmbeddings:
|
|||
call_args = mock_litellm.embedding.call_args
|
||||
assert call_args.kwargs["api_base"] == "https://custom.api.com"
|
||||
|
||||
async def test_output_dimensions_passed_when_set(self, mock_litellm):
|
||||
"""Test output dimensions are passed to LiteLLM when configured."""
|
||||
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",
|
||||
output_dimensions=768,
|
||||
)
|
||||
await emb.initialize()
|
||||
|
||||
init_call_args = mock_litellm.aembedding.call_args
|
||||
assert init_call_args.kwargs["dimensions"] == 768
|
||||
|
||||
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["dimensions"] == 768
|
||||
|
||||
async def test_output_dimensions_omitted_when_unset(self, mock_litellm):
|
||||
"""Test output dimensions are omitted when not configured."""
|
||||
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 "dimensions" 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 "dimensions" not in encode_call_args.kwargs
|
||||
|
||||
async def test_output_dimensions_and_api_base_passed_when_both_set(self, mock_litellm):
|
||||
"""Test both dimensions and api_base are forwarded when configured together."""
|
||||
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",
|
||||
api_base="https://custom.api.com",
|
||||
output_dimensions=768,
|
||||
)
|
||||
await emb.initialize()
|
||||
|
||||
init_call_args = mock_litellm.aembedding.call_args
|
||||
assert init_call_args.kwargs["api_base"] == "https://custom.api.com"
|
||||
assert init_call_args.kwargs["dimensions"] == 768
|
||||
|
||||
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["api_base"] == "https://custom.api.com"
|
||||
assert encode_call_args.kwargs["dimensions"] == 768
|
||||
|
||||
async def test_openai_invalid_output_dimensions_raises(self, mock_litellm):
|
||||
"""Invalid dimensions fail during initialize() (probe call), not per HTTP request.
|
||||
|
||||
MemoryEngine runs this at app lifespan startup; the process typically fails to become
|
||||
ready rather than returning a JSON error for a single API call. The RuntimeError
|
||||
message should still chain the underlying provider/LiteLLM detail for logs.
|
||||
"""
|
||||
mock_litellm.aembedding.side_effect = Exception("invalid dimensions for model")
|
||||
|
||||
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="openai/text-embedding-3-small",
|
||||
output_dimensions=9999,
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
RuntimeError,
|
||||
match="Failed to initialize LiteLLM SDK embeddings:.*invalid dimensions for model",
|
||||
):
|
||||
await emb.initialize()
|
||||
|
||||
|
||||
class TestLiteLLMSDKEmbeddingsFactory:
|
||||
"""Test the factory function for creating LiteLLM SDK embeddings."""
|
||||
|
|
@ -324,6 +417,21 @@ class TestLiteLLMSDKEmbeddingsFactory:
|
|||
assert isinstance(embeddings, LiteLLMSDKEmbeddings)
|
||||
assert embeddings.api_base == "https://custom.api.com"
|
||||
|
||||
def test_create_from_env_with_output_dimensions(self, monkeypatch):
|
||||
"""Test creating embeddings with configured output dimensions."""
|
||||
mock_config = MagicMock()
|
||||
mock_config.embeddings_provider = "litellm-sdk"
|
||||
mock_config.embeddings_litellm_sdk_api_key = "test_key"
|
||||
mock_config.embeddings_litellm_sdk_model = "gemini/gemini-embedding-2-preview"
|
||||
mock_config.embeddings_litellm_sdk_api_base = None
|
||||
mock_config.embeddings_litellm_sdk_output_dimensions = 768
|
||||
|
||||
with patch("hindsight_api.config.get_config", return_value=mock_config):
|
||||
embeddings = create_embeddings_from_env()
|
||||
|
||||
assert isinstance(embeddings, LiteLLMSDKEmbeddings)
|
||||
assert embeddings.output_dimensions == 768
|
||||
|
||||
|
||||
class TestLiteLLMSDKCohereEmbeddings:
|
||||
"""Integration tests calling real Cohere API (matches CI pattern)."""
|
||||
|
|
|
|||
|
|
@ -366,6 +366,7 @@ export HINDSIGHT_API_RETAIN_LLM_MAX_BACKOFF=120.0 # Cap at 2min instead of 1m
|
|||
| `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_KEY` | LiteLLM SDK API key for direct embedding provider access | - |
|
||||
| `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) | - |
|
||||
|
||||
```bash
|
||||
# Local (default) - uses SentenceTransformers
|
||||
|
|
@ -413,6 +414,8 @@ export HINDSIGHT_API_EMBEDDINGS_LITELLM_MODEL=text-embedding-3-small # or coher
|
|||
export HINDSIGHT_API_EMBEDDINGS_PROVIDER=litellm-sdk
|
||||
export HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_KEY=your-provider-api-key
|
||||
export HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MODEL=cohere/embed-english-v3.0
|
||||
# Optional: request a specific output dimension when the provider supports it
|
||||
# export HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS=768
|
||||
|
||||
# Supported LiteLLM SDK embedding providers:
|
||||
# - cohere/embed-english-v3.0 (1024 dimensions)
|
||||
|
|
@ -426,6 +429,8 @@ export HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MODEL=cohere/embed-english-v3.0
|
|||
|
||||
Hindsight automatically detects the embedding dimension from the model at startup and adjusts the database schema accordingly. The default model (`BAAI/bge-small-en-v1.5`) produces 384-dimensional vectors, while OpenAI models produce 1536 or 3072 dimensions.
|
||||
|
||||
For `litellm-sdk`, if you set `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS`, startup uses that output size when the underlying provider supports LiteLLM's `dimensions` parameter (otherwise behavior is unchanged). The same dimension-change rules below apply.
|
||||
|
||||
:::warning Dimension Changes
|
||||
Once memories are stored, you cannot change the embedding dimension without losing data. If you need to switch to a model with different dimensions:
|
||||
|
||||
|
|
|
|||
|
|
@ -366,6 +366,7 @@ export HINDSIGHT_API_RETAIN_LLM_MAX_BACKOFF=120.0 # Cap at 2min instead of 1m
|
|||
| `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_KEY` | LiteLLM SDK API key for direct embedding provider access | - |
|
||||
| `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) | - |
|
||||
|
||||
```bash
|
||||
# Local (default) - uses SentenceTransformers
|
||||
|
|
@ -413,6 +414,8 @@ export HINDSIGHT_API_EMBEDDINGS_LITELLM_MODEL=text-embedding-3-small # or coher
|
|||
export HINDSIGHT_API_EMBEDDINGS_PROVIDER=litellm-sdk
|
||||
export HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_API_KEY=your-provider-api-key
|
||||
export HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MODEL=cohere/embed-english-v3.0
|
||||
# Optional: request a specific output dimension when the provider supports it
|
||||
# export HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS=768
|
||||
|
||||
# Supported LiteLLM SDK embedding providers:
|
||||
# - cohere/embed-english-v3.0 (1024 dimensions)
|
||||
|
|
@ -426,6 +429,8 @@ export HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_MODEL=cohere/embed-english-v3.0
|
|||
|
||||
Hindsight automatically detects the embedding dimension from the model at startup and adjusts the database schema accordingly. The default model (`BAAI/bge-small-en-v1.5`) produces 384-dimensional vectors, while OpenAI models produce 1536 or 3072 dimensions.
|
||||
|
||||
For `litellm-sdk`, if you set `HINDSIGHT_API_EMBEDDINGS_LITELLM_SDK_OUTPUT_DIMENSIONS`, startup uses that output size when the underlying provider supports LiteLLM's `dimensions` parameter (otherwise behavior is unchanged). The same dimension-change rules below apply.
|
||||
|
||||
:::warning Dimension Changes
|
||||
Once memories are stored, you cannot change the embedding dimension without losing data. If you need to switch to a model with different dimensions:
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue