Skip to content

Commit

Permalink
creating Logger and RPMController
Browse files Browse the repository at this point in the history
  • Loading branch information
joaomdmoura committed Jan 14, 2024
1 parent 4fb7014 commit 1954be4
Show file tree
Hide file tree
Showing 2 changed files with 68 additions and 0 deletions.
11 changes: 11 additions & 0 deletions src/crewai/utilities/logger.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
class Logger:
def __init__(self, verbose_level=0):
verbose_level = (
2 if isinstance(verbose_level, bool) and verbose_level else verbose_level
)
self.verbose_level = verbose_level

def log(self, level, message):
level_map = {"debug": 1, "info": 2}
if self.verbose_level and level_map.get(level, 0) <= self.verbose_level:
print(f"\n[{level.upper()}]: {message}")
57 changes: 57 additions & 0 deletions src/crewai/utilities/rpm_controller.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
import threading
import time
from typing import Union

from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator

from crewai.utilities.logger import Logger


class RPMController(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
max_rpm: Union[int, None] = Field(default=None)
logger: Logger = Field(default=None)
_current_rpm: int = PrivateAttr(default=0)
_timer: threading.Timer = PrivateAttr(default=None)
_lock: threading.Lock = PrivateAttr(default=None)

@model_validator(mode="after")
def reset_counter(self):
if self.max_rpm:
self._lock = threading.Lock()
self._reset_request_count()
return self

def check_or_wait(self):
if not self.max_rpm:
return True

with self._lock:
if self._current_rpm < self.max_rpm:
self._current_rpm += 1
return True
else:
self.logger.log(
"info", "Max RPM reached, waiting for next minute to start."
)
self._wait_for_next_minute()
self._current_rpm = 1
return True

def stop_rpm_counter(self):
if self._timer:
self._timer.cancel()
self._timer = None

def _wait_for_next_minute(self):
time.sleep(60)
with self._lock:
self._current_rpm = 0

def _reset_request_count(self):
with self._lock:
self._current_rpm = 0
if self._timer:
self._timer.cancel()
self._timer = threading.Timer(60.0, self._reset_request_count)
self._timer.start()

0 comments on commit 1954be4

Please sign in to comment.