feat: Adding rate limiting

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

@@ -10,6 +10,10 @@
 from cognee.infrastructure.llm.tokenizer.HuggingFace import HuggingFaceTokenizer
 from cognee.infrastructure.llm.tokenizer.Mistral import MistralTokenizer
 from cognee.infrastructure.llm.tokenizer.TikToken import TikTokenTokenizer
+from cognee.infrastructure.llm.rate_limiter import (
+    embedding_rate_limit_async,
+    embedding_sleep_and_retry_async,
+)
 
 litellm.set_verbose = False
 logger = get_logger("LiteLLMEmbeddingEngine")
@@ -51,17 +55,12 @@ def __init__(
             enable_mocking = str(enable_mocking).lower()
         self.mock = enable_mocking in ("true", "1", "yes")
 
+    @embedding_sleep_and_retry_async()
+    @embedding_rate_limit_async
     async def embed_text(self, text: List[str]) -> List[List[float]]:
-        async def exponential_backoff(attempt):
-            wait_time = min(10 * (2**attempt), 60)  # Max 60 seconds
-            await asyncio.sleep(wait_time)
-
         try:
             if self.mock:
                 response = {"data": [{"embedding": [0.0] * self.dimensions} for _ in text]}
-
-                self.retry_count = 0
-
                 return [data["embedding"] for data in response["data"]]
             else:
                 response = await litellm.aembedding(
@@ -72,8 +71,6 @@ async def exponential_backoff(attempt):
                     api_version=self.api_version,
                 )
 
-                self.retry_count = 0  # Reset retry count on successful call
-
                 return [data["embedding"] for data in response.data]
 
         except litellm.exceptions.ContextWindowExceededError as error:
@@ -95,15 +92,6 @@ async def exponential_backoff(attempt):
             logger.error("Context window exceeded for embedding text: %s", str(error))
             raise error
 
-        except litellm.exceptions.RateLimitError:
-            if self.retry_count >= self.MAX_RETRIES:
-                raise Exception("Rate limit exceeded and no more retries left.")
-
-            await exponential_backoff(self.retry_count)
-            self.retry_count += 1
-
-            return await self.embed_text(text)
-
         except (
             litellm.exceptions.BadRequestError,
             litellm.exceptions.NotFoundError,

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

@@ -9,6 +9,10 @@
 from cognee.infrastructure.databases.vector.embeddings.EmbeddingEngine import EmbeddingEngine
 from cognee.infrastructure.databases.exceptions.EmbeddingException import EmbeddingException
 from cognee.infrastructure.llm.tokenizer.HuggingFace import HuggingFaceTokenizer
+from cognee.infrastructure.llm.rate_limiter import (
+    embedding_rate_limit_async,
+    embedding_sleep_and_retry_async,
+)
 
 logger = get_logger("OllamaEmbeddingEngine")
 
@@ -43,6 +47,7 @@ def __init__(
             enable_mocking = str(enable_mocking).lower()
         self.mock = enable_mocking in ("true", "1", "yes")
 
+    @embedding_rate_limit_async
     async def embed_text(self, text: List[str]) -> List[List[float]]:
         """
         Given a list of text prompts, returns a list of embedding vectors.
@@ -53,6 +58,7 @@ async def embed_text(self, text: List[str]) -> List[List[float]]:
         embeddings = await asyncio.gather(*[self._get_embedding(prompt) for prompt in text])
         return embeddings
 
+    @embedding_sleep_and_retry_async()
     async def _get_embedding(self, prompt: str) -> List[float]:
         """
         Internal method to call the Ollama embeddings endpoint for a single prompt.
@@ -66,26 +72,12 @@ async def _get_embedding(self, prompt: str) -> List[float]:
         if api_key:
             headers["Authorization"] = f"Bearer {api_key}"
 
-        retries = 0
-        while retries < self.MAX_RETRIES:
-            try:
-                async with aiohttp.ClientSession() as session:
-                    async with session.post(
-                        self.endpoint, json=payload, headers=headers, timeout=60.0
-                    ) as response:
-                        data = await response.json()
-                        return data["embedding"]
-            except aiohttp.http_exceptions.HttpBadRequest as e:
-                logger.error(f"HTTP error on attempt {retries + 1}: {e}")
-                retries += 1
-                await asyncio.sleep(min(2**retries, 60))
-            except Exception as e:
-                logger.error(f"Error on attempt {retries + 1}: {e}")
-                retries += 1
-                await asyncio.sleep(min(2**retries, 60))
-        raise EmbeddingException(
-            f"Failed to embed text using model {self.model} after {self.MAX_RETRIES} retries"
-        )
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                self.endpoint, json=payload, headers=headers, timeout=60.0
+            ) as response:
+                data = await response.json()
+                return data["embedding"]
 
     def get_vector_size(self) -> int:
         return self.dimensions

cognee/infrastructure/llm/anthropic/adapter.py

@@ -5,6 +5,7 @@
 from cognee.exceptions import InvalidValueError
 from cognee.infrastructure.llm.llm_interface import LLMInterface
 from cognee.infrastructure.llm.prompts import read_query_prompt
+from cognee.infrastructure.llm.rate_limiter import rate_limit_async, sleep_and_retry_async
 
 
 class AnthropicAdapter(LLMInterface):
@@ -23,6 +24,8 @@ def __init__(self, max_tokens: int, model: str = None):
         self.model = model
         self.max_tokens = max_tokens
 
+    @sleep_and_retry_async()
+    @rate_limit_async
     async def acreate_structured_output(
         self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
     ) -> BaseModel:

cognee/infrastructure/llm/config.py

@@ -16,6 +16,12 @@ class LLMConfig(BaseSettings):
     llm_max_tokens: int = 16384
     transcription_model: str = "whisper-1"
     graph_prompt_path: str = "generate_graph_prompt.txt"
+    llm_rate_limit_enabled: bool = False
+    llm_rate_limit_requests: int = 60
+    llm_rate_limit_interval: int = 60  # in seconds (default is 60 requests per minute)
+    embedding_rate_limit_enabled: bool = False
+    embedding_rate_limit_requests: int = 60
+    embedding_rate_limit_interval: int = 60  # in seconds (default is 60 requests per minute)
 
     model_config = SettingsConfigDict(env_file=".env", extra="allow")
 
@@ -85,6 +91,12 @@ def to_dict(self) -> dict:
             "max_tokens": self.llm_max_tokens,
             "transcription_model": self.transcription_model,
             "graph_prompt_path": self.graph_prompt_path,
+            "rate_limit_enabled": self.llm_rate_limit_enabled,
+            "rate_limit_requests": self.llm_rate_limit_requests,
+            "rate_limit_interval": self.llm_rate_limit_interval,
+            "embedding_rate_limit_enabled": self.embedding_rate_limit_enabled,
+            "embedding_rate_limit_requests": self.embedding_rate_limit_requests,
+            "embedding_rate_limit_interval": self.embedding_rate_limit_interval,
         }
 
 

cognee/infrastructure/llm/gemini/adapter.py

@@ -7,6 +7,11 @@
 from cognee.exceptions import InvalidValueError
 from cognee.infrastructure.llm.llm_interface import LLMInterface
 from cognee.infrastructure.llm.prompts import read_query_prompt
+from cognee.infrastructure.llm.rate_limiter import (
+    rate_limit_async,
+    rate_limit_sync,
+    sleep_and_retry_async,
+)
 from cognee.base_config import get_base_config
 
 logger = get_logger()
@@ -37,6 +42,8 @@ def __init__(
         self.max_tokens = max_tokens
 
     @observe(as_type="generation")
+    @sleep_and_retry_async()
+    @rate_limit_async
     async def acreate_structured_output(
         self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
     ) -> BaseModel:

cognee/infrastructure/llm/generic_llm_api/adapter.py

@@ -6,6 +6,7 @@
 import instructor
 from cognee.infrastructure.llm.llm_interface import LLMInterface
 from cognee.infrastructure.llm.config import get_llm_config
+from cognee.infrastructure.llm.rate_limiter import rate_limit_async, sleep_and_retry_async
 import litellm
 
 
@@ -27,6 +28,8 @@ def __init__(self, endpoint, api_key: str, model: str, name: str, max_tokens: in
             litellm.acompletion, mode=instructor.Mode.JSON, api_key=api_key
         )
 
+    @sleep_and_retry_async()
+    @rate_limit_async
     async def acreate_structured_output(
         self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
     ) -> BaseModel:

cognee/infrastructure/llm/ollama/adapter.py

@@ -3,6 +3,11 @@
 import instructor
 from cognee.infrastructure.llm.llm_interface import LLMInterface
 from cognee.infrastructure.llm.config import get_llm_config
+from cognee.infrastructure.llm.rate_limiter import (
+    rate_limit_async,
+    rate_limit_sync,
+    sleep_and_retry_async,
+)
 from openai import OpenAI
 import base64
 import os
@@ -22,6 +27,8 @@ def __init__(self, endpoint: str, api_key: str, model: str, name: str, max_token
             OpenAI(base_url=self.endpoint, api_key=self.api_key), mode=instructor.Mode.JSON
         )
 
+    @sleep_and_retry_async()
+    @rate_limit_async
     async def acreate_structured_output(
         self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
     ) -> BaseModel:
@@ -45,6 +52,7 @@ async def acreate_structured_output(
 
         return response
 
+    @rate_limit_sync
     def create_transcript(self, input_file: str) -> str:
         """Generate an audio transcript from a user query."""
 
@@ -64,6 +72,7 @@ def create_transcript(self, input_file: str) -> str:
 
         return transcription.text
 
+    @rate_limit_sync
     def transcribe_image(self, input_file: str) -> str:
         """Transcribe content from an image using base64 encoding."""
 
@@ -79,7 +88,7 @@ def transcribe_image(self, input_file: str) -> str:
                 {
                     "role": "user",
                     "content": [
-                        {"type": "text", "text": "What’s in this image?"},
+                        {"type": "text", "text": "What's in this image?"},
                         {
                             "type": "image_url",
                             "image_url": {"url": f"data:image/jpeg;base64,{encoded_image}"},

cognee/infrastructure/llm/openai/adapter.py

@@ -10,6 +10,12 @@
 from cognee.exceptions import InvalidValueError
 from cognee.infrastructure.llm.llm_interface import LLMInterface
 from cognee.infrastructure.llm.prompts import read_query_prompt
+from cognee.infrastructure.llm.rate_limiter import (
+    rate_limit_async,
+    rate_limit_sync,
+    sleep_and_retry_async,
+    sleep_and_retry_sync,
+)
 from cognee.base_config import get_base_config
 
 monitoring = get_base_config().monitoring_tool
@@ -49,6 +55,8 @@ def __init__(
         self.streaming = streaming
 
     @observe(as_type="generation")
+    @sleep_and_retry_async()
+    @rate_limit_async
     async def acreate_structured_output(
         self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
     ) -> BaseModel:
@@ -75,6 +83,8 @@ async def acreate_structured_output(
         )
 
     @observe
+    @sleep_and_retry_sync()
+    @rate_limit_sync
     def create_structured_output(
         self, text_input: str, system_prompt: str, response_model: Type[BaseModel]
     ) -> BaseModel:
@@ -100,6 +110,7 @@ def create_structured_output(
             max_retries=self.MAX_RETRIES,
         )
 
+    @rate_limit_sync
     def create_transcript(self, input):
         """Generate a audio transcript from a user query."""
 
@@ -120,6 +131,7 @@ def create_transcript(self, input):
 
         return transcription
 
+    @rate_limit_sync
     def transcribe_image(self, input) -> BaseModel:
         with open(input, "rb") as image_file:
             encoded_image = base64.b64encode(image_file.read()).decode("utf-8")
@@ -132,7 +144,7 @@ def transcribe_image(self, input) -> BaseModel:
                     "content": [
                         {
                             "type": "text",
-                            "text": "What’s in this image?",
+                            "text": "What's in this image?",
                         },
                         {
                             "type": "image_url",