Skip to content

feat: Enable dimension passing to llitellm #1942

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dexters1 merged 5 commits into from
Jan 21, 2026
Merged

feat: Enable dimension passing to llitellm #1942

Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a47cfe7
added dimensions in a simple way
Vasilije1990 Dec 28, 2025
12d5a7d
Merge dev into COG-2988
Vasilije1990 Jan 12, 2026
3a6bb77
refactor: Handle embedding case in data
dexters1 Jan 20, 2026
9d73b49
refactor: add support for Ollama embedding size definition
dexters1 Jan 20, 2026
38f0ab8
Merge branch 'dev' into COG-2988
dexters1 Jan 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 12 additions & 7 deletions cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,13 +111,18 @@ async def embed_text(self, text: List[str]) -> List[List[float]]:
return [data["embedding"] for data in response["data"]]
else:
async with embedding_rate_limiter_context_manager():
response = await litellm.aembedding(
model=self.model,
input=text,
api_key=self.api_key,
api_base=self.endpoint,
api_version=self.api_version,
)
embedding_kwargs = {
"model": self.model,
"input": text,
"api_key": self.api_key,
"api_base": self.endpoint,
"api_version": self.api_version,
}
# Pass through target embedding dimensions when supported
if self.dimensions is not None:
embedding_kwargs["dimensions"] = self.dimensions

response = await litellm.aembedding(**embedding_kwargs)
Copy link
Contributor

@coderabbitai coderabbitai bot Dec 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

🧩 Analysis chain

🏁 Script executed:

# First, let's check the file content to understand the context
cat -n cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py | head -150

Repository: topoteretes/cognee

Length of output: 6543

🏁 Script executed:

# Check if litellm is available and get version info
pip show litellm 2>/dev/null | grep -E "Version|Name"

Repository: topoteretes/cognee

Length of output: 148

🏁 Script executed:

# Check the aembedding function signature if available
python3 << 'EOF'
try:
    import inspect
    import litellm
    sig = inspect.signature(litellm.aembedding)
    print("aembedding signature:")
    print(sig)
    print("\nParameters:")
    for param_name, param in sig.parameters.items():
        print(f"  {param_name}: {param.annotation if param.annotation != inspect.Parameter.empty else 'no annotation'}")
except Exception as e:
    print(f"Error: {e}")
EOF

Repository: topoteretes/cognee

Length of output: 148

🌐 Web query:

litellm aembedding dimensions parameter API documentation

💡 Result:

The LiteLLM embedding API (embedding / aembedding) accepts an optional integer parameter dimensions which requests the output embedding dimensionality. It’s documented as an optional field (used with OpenAI/Azure text-embedding-3+ models and some provider-specific models) and appears in the aembedding examples. See LiteLLM docs for details and supported-model notes. [1][2]

Sources
[1] LiteLLM aembedding / embedding docs (parameters list, includes dimensions).
[2] LiteLLM “supported embedding” / provider notes (dimensions usage and examples).

🌐 Web query:

OpenAI embedding models dimensions parameter support text-embedding-3

💡 Result:

Yes — the text-embedding-3 models support a dimensions parameter. Summary:

  • text-embedding-3-large: default 3072, can be shortened to 1024 or 256 (supports 3072, 1024, 256). [1][2]
  • text-embedding-3-small: default 1536, can be returned at 512 or 1536 (supports 1536, 512). [1][3]

References:

  • OpenAI model page (text-embedding-3-large / embeddings). [1]
  • OpenAI announcement describing the dimensions parameter and shortening examples. [2]
  • Weaviate docs listing available dimensions per model. [3]

Add docstring documentation for the dimensions parameter.

The implementation correctly constructs embedding kwargs and conditionally passes the dimensions parameter to litellm.aembedding. LiteLLM's API supports the dimensions parameter (confirmed for version 1.76.0+), and the default model (openai/text-embedding-3-large) accepts it with valid values of 3072, 1024, or 256.

However, the class docstring (lines 35-43) does not document the dimensions parameter or its behavior. Update the docstring to explain:

  • That dimensions controls output embedding dimensionality
  • That it's passed through to the underlying embedding API when supported by the model
  • The valid range or values for the default model
🤖 Prompt for AI Agents 
In cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py
around lines 35 to 43 and implementation at 114-125, the class docstring is
missing documentation for the dimensions parameter; update the class docstring
to state that dimensions is an optional int controlling output embedding
dimensionality, that when set it is passed through to the underlying LiteLLM
embedding API (litellm.aembedding) only if the model supports it, and document
the valid values for the default model (openai/text-embedding-3-large accepts
3072, 1024, or 256). Also note the default behavior when dimensions is None (use
model default) and mention LiteLLM version requirement (1.76.0+) for dimensions
support.


return [data["embedding"] for data in response.data]

Expand Down
Loading