mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-11 00:58:30 +00:00
add rate limiting
This commit is contained in:
@@ -1,7 +1,7 @@
|
|||||||
import datetime
|
import datetime
|
||||||
import os
|
import os
|
||||||
from functools import lru_cache
|
import time
|
||||||
from typing import Any, Optional, Type
|
from typing import Any, ClassVar, Optional, Type
|
||||||
|
|
||||||
import requests
|
import requests
|
||||||
from pydantic import BaseModel, Field
|
from pydantic import BaseModel, Field
|
||||||
@@ -47,6 +47,8 @@ class BraveSearchTool(BaseTool):
|
|||||||
country: Optional[str] = ""
|
country: Optional[str] = ""
|
||||||
n_results: int = 10
|
n_results: int = 10
|
||||||
save_file: bool = False
|
save_file: bool = False
|
||||||
|
_last_request_time: ClassVar[float] = 0
|
||||||
|
_min_request_interval: ClassVar[float] = 1.0 # seconds
|
||||||
|
|
||||||
def __init__(self, *args, **kwargs):
|
def __init__(self, *args, **kwargs):
|
||||||
super().__init__(*args, **kwargs)
|
super().__init__(*args, **kwargs)
|
||||||
@@ -55,11 +57,16 @@ class BraveSearchTool(BaseTool):
|
|||||||
"BRAVE_API_KEY environment variable is required for BraveSearchTool"
|
"BRAVE_API_KEY environment variable is required for BraveSearchTool"
|
||||||
)
|
)
|
||||||
|
|
||||||
@lru_cache(maxsize=100)
|
|
||||||
def _run(
|
def _run(
|
||||||
self,
|
self,
|
||||||
**kwargs: Any,
|
**kwargs: Any,
|
||||||
) -> Any:
|
) -> Any:
|
||||||
|
current_time = time.time()
|
||||||
|
if (current_time - self._last_request_time) < self._min_request_interval:
|
||||||
|
time.sleep(
|
||||||
|
self._min_request_interval - (current_time - self._last_request_time)
|
||||||
|
)
|
||||||
|
BraveSearchTool._last_request_time = time.time()
|
||||||
try:
|
try:
|
||||||
search_query = kwargs.get("search_query") or kwargs.get("query")
|
search_query = kwargs.get("search_query") or kwargs.get("query")
|
||||||
if not search_query:
|
if not search_query:
|
||||||
|
|||||||
@@ -40,8 +40,8 @@ def test_brave_tool():
|
|||||||
tool = BraveSearchTool(
|
tool = BraveSearchTool(
|
||||||
n_results=2,
|
n_results=2,
|
||||||
)
|
)
|
||||||
|
x = tool.run(search_query="ChatGPT")
|
||||||
print(tool.run(search_query="ChatGPT"))
|
print(x)
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
Reference in New Issue
Block a user