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

added dimensions in a simple way

Browse source
This commit is contained in:
vasilije 2025-12-28 21:37:22 +01:00
parent 310e9e97ae
commit a47cfe7272
1 changed files with 12 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -111,13 +111,18 @@ class LiteLLMEmbeddingEngine(EmbeddingEngine):
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)
return [data["embedding"] for data in response.data]