fix: handle unsupported dimensions and ensure mock mode safety

This commit is contained in:
shijianglong 2026-01-05 16:00:36 +08:00
parent 76bfb3ac3e
commit 8fd96c7275

View file

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