From 2e679b69e5197d7481a7d2989f30aac2be07d398 Mon Sep 17 00:00:00 2001 From: shijianglong Date: Wed, 7 Jan 2026 15:04:29 +0800 Subject: [PATCH] fix: address PR code review comments Signed-off-by: shijianglong --- .../vector/embeddings/LiteLLMEmbeddingEngine.py | 5 ++--- .../test_litellm_embedding_dimensions.py | 10 +++++++++- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py b/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py index e1bad3a8c..aa93906c8 100644 --- a/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py +++ b/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py @@ -30,7 +30,6 @@ from cognee.infrastructure.llm.tokenizer.TikToken import ( from cognee.shared.rate_limiting import embedding_rate_limiter_context_manager litellm.set_verbose = False -litellm.drop_params = True logger = get_logger("LiteLLMEmbeddingEngine") @@ -50,7 +49,7 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine): api_version: str provider: str model: str - dimensions: int + dimensions: Optional[int] mock: bool MAX_RETRIES = 5 @@ -293,4 +292,4 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine): ) logger.debug(f"Tokenizer loaded for model: {self.model}") - return tokenizer + return tokenizer \ No newline at end of file diff --git a/cognee/tests/unit/infrastructure/test_litellm_embedding_dimensions.py b/cognee/tests/unit/infrastructure/test_litellm_embedding_dimensions.py index af52d27d6..f6f5388e1 100644 --- a/cognee/tests/unit/infrastructure/test_litellm_embedding_dimensions.py +++ b/cognee/tests/unit/infrastructure/test_litellm_embedding_dimensions.py @@ -1,6 +1,8 @@ -import pytest import os from unittest.mock import patch + +import pytest + from cognee.infrastructure.databases.vector.embeddings.LiteLLMEmbeddingEngine import LiteLLMEmbeddingEngine @pytest.mark.asyncio @@ -46,3 +48,9 @@ async def test_litellm_embedding_invalid_dimensions(): with pytest.raises(ValueError, match="dimensions must be a positive integer"): LiteLLMEmbeddingEngine(dimensions=-100) + + with pytest.raises(ValueError, match="dimensions must be a positive integer"): + LiteLLMEmbeddingEngine(dimensions="1024") # type: ignore + + with pytest.raises(ValueError, match="dimensions must be a positive integer"): + LiteLLMEmbeddingEngine(dimensions=1024.5) # type: ignore \ No newline at end of file