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"]] return [data["embedding"] for data in response["data"]]
else: else:
async with embedding_rate_limiter_context_manager(): async with embedding_rate_limiter_context_manager():
response = await litellm.aembedding( embedding_kwargs = {
model=self.model, "model": self.model,
input=text, "input": text,
api_key=self.api_key, "api_key": self.api_key,
api_base=self.endpoint, "api_base": self.endpoint,
api_version=self.api_version, "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] return [data["embedding"] for data in response.data]