From 7dea1d54d7f0f97012c4e4be4a29dd661fa80d19 Mon Sep 17 00:00:00 2001
From: Igor Ilic <igorilic03@gmail.com>
Date: Thu, 23 Jan 2025 18:18:45 +0100
Subject: [PATCH] refactor: Add specific max token values to embedding models

---
 .../databases/vector/embeddings/LiteLLMEmbeddingEngine.py      | 3 +--
 cognee/infrastructure/llm/tokenizer/Gemini/adapter.py          | 2 +-
 cognee/infrastructure/llm/tokenizer/HuggingFace/adapter.py     | 2 +-
 cognee/infrastructure/llm/tokenizer/TikToken/adapter.py        | 2 +-
 4 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py b/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py
index cb84337c..10992b22 100644
--- a/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py
+++ b/cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py
@@ -31,12 +31,11 @@ def __init__(
         api_key: str = None,
         endpoint: str = None,
         api_version: str = None,
-        max_tokens: int = float("inf"),
+        max_tokens: int = 512,
     ):
         self.api_key = api_key
         self.endpoint = endpoint
         self.api_version = api_version
-        # TODO: Add or remove provider info
         self.provider = provider
         self.model = model
         self.dimensions = dimensions
diff --git a/cognee/infrastructure/llm/tokenizer/Gemini/adapter.py b/cognee/infrastructure/llm/tokenizer/Gemini/adapter.py
index 697bc957..f3131ea0 100644
--- a/cognee/infrastructure/llm/tokenizer/Gemini/adapter.py
+++ b/cognee/infrastructure/llm/tokenizer/Gemini/adapter.py
@@ -7,7 +7,7 @@ class GeminiTokenizer(TokenizerInterface):
     def __init__(
         self,
         model: str,
-        max_tokens: int = float("inf"),
+        max_tokens: int = 3072,
     ):
         self.model = model
         self.max_tokens = max_tokens
diff --git a/cognee/infrastructure/llm/tokenizer/HuggingFace/adapter.py b/cognee/infrastructure/llm/tokenizer/HuggingFace/adapter.py
index 7b92fb76..a8eac29d 100644
--- a/cognee/infrastructure/llm/tokenizer/HuggingFace/adapter.py
+++ b/cognee/infrastructure/llm/tokenizer/HuggingFace/adapter.py
@@ -9,7 +9,7 @@ class HuggingFaceTokenizer(TokenizerInterface):
     def __init__(
         self,
         model: str,
-        max_tokens: int = float("inf"),
+        max_tokens: int = 512,
     ):
         self.model = model
         self.max_tokens = max_tokens
diff --git a/cognee/infrastructure/llm/tokenizer/TikToken/adapter.py b/cognee/infrastructure/llm/tokenizer/TikToken/adapter.py
index 6ba1e002..862a7929 100644
--- a/cognee/infrastructure/llm/tokenizer/TikToken/adapter.py
+++ b/cognee/infrastructure/llm/tokenizer/TikToken/adapter.py
@@ -13,7 +13,7 @@ class TikTokenTokenizer(TokenizerInterface):
     def __init__(
         self,
         model: str,
-        max_tokens: int = float("inf"),
+        max_tokens: int = 8191,
     ):
         self.model = model
         self.max_tokens = max_tokens