fix: handle unsupported dimensions and ensure mock mode safety
This commit is contained in:
parent
76bfb3ac3e
commit
8fd96c7275
1 changed files with 3 additions and 1 deletions
|
|
@ -28,6 +28,7 @@ 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")
|
||||||
|
|
||||||
|
|
||||||
|
|
@ -111,7 +112,8 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
if self.mock:
|
if self.mock:
|
||||||
response = {"data": [{"embedding": [0.0] * self.dimensions} for _ in text]}
|
dim = self.dimensions if self.dimensions is not None else 3072
|
||||||
|
response = {"data": [{"embedding": [0.0] * dim} for _ in text]}
|
||||||
return [data["embedding"] for data in response["data"]]
|
return [data["embedding"] for data in response["data"]]
|
||||||
else:
|
else:
|
||||||
async with embedding_rate_limiter_context_manager():
|
async with embedding_rate_limiter_context_manager():
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue