-
Notifications
You must be signed in to change notification settings - Fork 1.1k
Description
Bug Description
In the file cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py, the init method correctly receives and stores the dimensions parameter (corresponding to cognee's EMBEDDING_DIMENSIONS configuration, default 3072), which is used for get_vector_size() return value and mock mode vector generation ([0.0] * self.dimensions).
However, in the embed_text() method's actual embedding branch, the call to litellm.aembedding() only passes model, input, api_key, api_base, and api_version parameters, without passing the dimensions parameter. This causes LiteLLM to use the model's default embedding dimension (e.g., 3072 for OpenAI text-embedding-3-large), ignoring cognee's custom configuration.
Steps to Reproduce
- Configure cognee's EMBEDDING_DIMENSIONS=512.
- Initialize LiteLLMEmbeddingEngine and call embed_text("test").
- Check the returned vector length: expected 512, but actual is model default (e.g., 3072).
Expected Behavior
litellm.aembedding() should receive the dimensions=self.dimensions parameter (supported by models like OpenAI), ensuring the output dimension matches cognee's configuration.
Actual Behavior
Add dimensions=self.dimensions to the litellm.aembedding() call in the embed_text() method:
response = await litellm.aembedding(
model=self.model,
input=texts,
api_key=self.api_key,
api_base=self.api_base,
api_version=self.api_version,
dimensions=self.dimensions # Add this line
)Environment
- OS: [macOS 15.6.1]
- Python version: [3.12.11]
- Cognee version: [0.5.1]
- LLM Provider: [OpenAI]
- Database: [Neo4j,qdrant]
Logs/Error Messages
Additional Context
No response
Pre-submission Checklist
- I have searched existing issues to ensure this bug hasn't been reported already
- I have provided a clear and detailed description of the bug
- I have included steps to reproduce the issue
- I have included my environment details