fix: Resolve issue with BAML rate limit handling

cognee/infrastructure/databases/vector/embeddings/FastembedEmbeddingEngine.py

@@ -17,6 +17,7 @@
 from cognee.infrastructure.llm.tokenizer.TikToken import (
     TikTokenTokenizer,
 )
+from cognee.shared.rate_limiting import embedding_rate_limiter_context_manager
 
 litellm.set_verbose = False
 logger = get_logger("FastembedEmbeddingEngine")
@@ -68,7 +69,7 @@ def __init__(
 
     @retry(
         stop=stop_after_delay(128),
-        wait=wait_exponential_jitter(2, 128),
+        wait=wait_exponential_jitter(8, 128),
         retry=retry_if_not_exception_type(litellm.exceptions.NotFoundError),
         before_sleep=before_sleep_log(logger, logging.DEBUG),
         reraise=True,
@@ -96,11 +97,12 @@ async def embed_text(self, text: List[str]) -> List[List[float]]:
             if self.mock:
                 return [[0.0] * self.dimensions for _ in text]
             else:
-                embeddings = self.embedding_model.embed(
-                    text,
-                    batch_size=len(text),
-                    parallel=None,
-                )
+                async with embedding_rate_limiter_context_manager():
+                    embeddings = self.embedding_model.embed(
+                        text,
+                        batch_size=len(text),
+                        parallel=None,
+                    )
 
                 return list(embeddings)
 

cognee/infrastructure/databases/vector/embeddings/LiteLLMEmbeddingEngine.py

@@ -25,6 +25,7 @@
 from cognee.infrastructure.llm.tokenizer.TikToken import (
     TikTokenTokenizer,
 )
+from cognee.shared.rate_limiting import embedding_rate_limiter_context_manager
 
 litellm.set_verbose = False
 logger = get_logger("LiteLLMEmbeddingEngine")
@@ -109,13 +110,14 @@ async def embed_text(self, text: List[str]) -> List[List[float]]:
                 response = {"data": [{"embedding": [0.0] * self.dimensions} for _ in text]}
                 return [data["embedding"] for data in response["data"]]
             else:
-                response = await litellm.aembedding(
-                    model=self.model,
-                    input=text,
-                    api_key=self.api_key,
-                    api_base=self.endpoint,
-                    api_version=self.api_version,
-                )
+                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,
+                    )
 
                 return [data["embedding"] for data in response.data]
 

cognee/infrastructure/databases/vector/embeddings/OllamaEmbeddingEngine.py

@@ -18,10 +18,7 @@
 from cognee.infrastructure.llm.tokenizer.HuggingFace import (
     HuggingFaceTokenizer,
 )
-from cognee.infrastructure.databases.vector.embeddings.embedding_rate_limiter import (
-    embedding_rate_limit_async,
-    embedding_sleep_and_retry_async,
-)
+from cognee.shared.rate_limiting import embedding_rate_limiter_context_manager
 from cognee.shared.utils import create_secure_ssl_context
 
 logger = get_logger("OllamaEmbeddingEngine")
@@ -101,7 +98,7 @@ async def embed_text(self, text: List[str]) -> List[List[float]]:
 
     @retry(
         stop=stop_after_delay(128),
-        wait=wait_exponential_jitter(2, 128),
+        wait=wait_exponential_jitter(8, 128),
         retry=retry_if_not_exception_type(litellm.exceptions.NotFoundError),
         before_sleep=before_sleep_log(logger, logging.DEBUG),
         reraise=True,
@@ -120,14 +117,15 @@ async def _get_embedding(self, prompt: str) -> List[float]:
         ssl_context = create_secure_ssl_context()
         connector = aiohttp.TCPConnector(ssl=ssl_context)
         async with aiohttp.ClientSession(connector=connector) as session:
-            async with session.post(
-                self.endpoint, json=payload, headers=headers, timeout=60.0
-            ) as response:
-                data = await response.json()
-                if "embeddings" in data:
-                    return data["embeddings"][0]
-                else:
-                    return data["data"][0]["embedding"]
+            async with embedding_rate_limiter_context_manager():
+                async with session.post(
+                    self.endpoint, json=payload, headers=headers, timeout=60.0
+                ) as response:
+                    data = await response.json()
+                    if "embeddings" in data:
+                        return data["embeddings"][0]
+                    else:
+                        return data["data"][0]["embedding"]
 
     def get_vector_size(self) -> int:
         """