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

hindsight-api-slim/hindsight_api/config.py

@@ -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),

hindsight-api-slim/hindsight_api/engine/embeddings.py

@@ -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(

hindsight-api-slim/tests/test_litellm_sdk_embeddings.py

@@ -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)."""

hindsight-docs/docs/developer/configuration.md

@@ -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:
 

hindsight-docs/versioned_docs/version-0.4/developer/configuration.md

@@ -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: