add rate limiting

This commit is contained in:
siddas27
2024-11-30 22:08:29 -06:00
parent 5532ea8ff7
commit e7e059d02a
2 changed files with 12 additions and 5 deletions

View File

@@ -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:

View File

@@ -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__":