Skip to content

[Bug]: LiteLLMEmbeddingEngine Does Not Pass EMBEDDING_DIMENSIONS to litellm, Making Custom Embedding Dimensions Ineffective #1961

@stonyme

Description

@stonyme

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

  1. Configure cognee's EMBEDDING_DIMENSIONS=512.
  2. Initialize LiteLLMEmbeddingEngine and call embed_text("test").
  3. 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

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions