fix: address PR code review comments
Signed-off-by: shijianglong <stonyme@vip.qq.com>
This commit is contained in:
parent
a6f57bc8ad
commit
2e679b69e5
2 changed files with 11 additions and 4 deletions
|
|
@ -30,7 +30,6 @@ from cognee.infrastructure.llm.tokenizer.TikToken import (
|
||||||
from cognee.shared.rate_limiting import embedding_rate_limiter_context_manager
|
from cognee.shared.rate_limiting import embedding_rate_limiter_context_manager
|
||||||
|
|
||||||
litellm.set_verbose = False
|
litellm.set_verbose = False
|
||||||
litellm.drop_params = True
|
|
||||||
logger = get_logger("LiteLLMEmbeddingEngine")
|
logger = get_logger("LiteLLMEmbeddingEngine")
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -50,7 +49,7 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
|
||||||
api_version: str
|
api_version: str
|
||||||
provider: str
|
provider: str
|
||||||
model: str
|
model: str
|
||||||
dimensions: int
|
dimensions: Optional[int]
|
||||||
mock: bool
|
mock: bool
|
||||||
|
|
||||||
MAX_RETRIES = 5
|
MAX_RETRIES = 5
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,8 @@
|
||||||
import pytest
|
|
||||||
import os
|
import os
|
||||||
from unittest.mock import patch
|
from unittest.mock import patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
from cognee.infrastructure.databases.vector.embeddings.LiteLLMEmbeddingEngine import LiteLLMEmbeddingEngine
|
from cognee.infrastructure.databases.vector.embeddings.LiteLLMEmbeddingEngine import LiteLLMEmbeddingEngine
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@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"):
|
with pytest.raises(ValueError, match="dimensions must be a positive integer"):
|
||||||
LiteLLMEmbeddingEngine(dimensions=-100)
|
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
|
||||||
Loading…
Add table
Reference in a new issue