gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

docs: add docstrings and fix type annotations for LiteLLMEmbeddingEngine

Browse source
This commit is contained in:
shijianglong 2026-01-05 16:19:49 +08:00
parent 8fd96c7275
commit 338c9e4366
1 changed files with 5 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py
View file

@ -64,6 +64,10 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
max_completion_tokens: int = 512, max_completion_tokens: int = 512,
batch_size: int = 100, batch_size: int = 100,
): ):
"""
Initializes the embedding engine with the specified model and provider.
Supports dimension configuration for compatible models.
"""
self.api_key = api_key self.api_key = api_key
self.endpoint = endpoint self.endpoint = endpoint
self.api_version = api_version self.api_version = api_version
@ -78,7 +82,6 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
if isinstance(enable_mocking, bool): if isinstance(enable_mocking, bool):
enable_mocking = str(enable_mocking).lower() enable_mocking = str(enable_mocking).lower()
self.mock = enable_mocking in ("true", "1", "yes") self.mock = enable_mocking in ("true", "1", "yes")
if dimensions is not None: if dimensions is not None:
if not isinstance(dimensions, int) or dimensions <= 0: if not isinstance(dimensions, int) or dimensions <= 0:
raise ValueError("dimensions must be a positive integer") raise ValueError("dimensions must be a positive integer")
@ -184,7 +187,7 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
- int: The size (dimensionality) of the embedding vectors. - int: The size (dimensionality) of the embedding vectors.
""" """
return self.dimensions return self.dimensions if self.dimensions is not None else 3072
def get_batch_size(self) -> int: def get_batch_size(self) -> int:
""" """