diff --git a/hindsight-api-slim/hindsight_api/config.py b/hindsight-api-slim/hindsight_api/config.py index ddc7fcd9..c6365901 100644 --- a/hindsight-api-slim/hindsight_api/config.py +++ b/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), diff --git a/hindsight-api-slim/hindsight_api/engine/embeddings.py b/hindsight-api-slim/hindsight_api/engine/embeddings.py index 37d446c6..cb9615aa 100644 --- a/hindsight-api-slim/hindsight_api/engine/embeddings.py +++ b/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( diff --git a/hindsight-api-slim/tests/test_litellm_sdk_embeddings.py b/hindsight-api-slim/tests/test_litellm_sdk_embeddings.py index 256662bc..2754517b 100644 --- a/hindsight-api-slim/tests/test_litellm_sdk_embeddings.py +++ b/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).""" diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index dbdc7753..84cadd92 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/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: diff --git a/hindsight-docs/versioned_docs/version-0.4/developer/configuration.md b/hindsight-docs/versioned_docs/version-0.4/developer/configuration.md index 92464fd5..ba6b02ed 100644 --- a/hindsight-docs/versioned_docs/version-0.4/developer/configuration.md +++ b/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: