feat(core): add experimental rate limiter agent and protocol (#562)

Co-authored-by: James Riehl <[email protected]>

python/src/uagents/experimental/quota/__init__.py

@@ -0,0 +1,272 @@
+"""
+This Protocol class can be used to rate limit `on_message` message handlers.
+
+The rate limiter uses the agents storage to keep track of the number of requests
+made by another agent within a given time window. If the number of requests exceeds
+a specified limit, the rate limiter will block further requests until the time
+window resets.
+
+> Default: Not rate limited, but you can set a default during initialization.
+
+Additionally, the protocol can be used to set access control rules for handlers
+allowing or blocking specific agents from accessing the handler.
+The default access control rule can be set to allow or block all agents.
+
+Both rules can work together to provide a secure and rate-limited environment for
+message handlers.
+
+
+Usage examples:
+
+```python
+from uagents.experimental.quota import AccessControlList, QuotaProtocol, RateLimit
+
+# Initialize the QuotaProtocol instance
+quota_protocol = QuotaProtocol(
+ storage_reference=agent.storage,
+ name="quota_proto",
+ version=agent._version,
+ # default_rate_limit=RateLimit(window_size_minutes=1, max_requests=3), # Optional
+)
+
+# This message handler is not rate limited
+@quota_protocol.on_message(ExampleMessage1)
+async def handle(ctx: Context, sender: str, msg: ExampleMessage1):
+ ...
+
+# This message handler is rate limited with custom window size and request limit
+@quota_protocol.on_message(
+ ExampleMessage2,
+ rate_limit=RateLimit(window_size_minutes=1, max_requests=3),
+)
+async def handle(ctx: Context, sender: str, msg: ExampleMessage2):
+ ...
+
+# This message handler has access control rules set
+@quota_protocol.on_message(
+ ExampleMessage3,
+ acl=AccessControlList(default=False, allowed={"<agent_address>"}),
+)
+async def handle(ctx: Context, sender: str, msg: ExampleMessage3):
+ ...
+
+agent.include(quota_protocol)
+```
+
+Tip: The `AccessControlList` object can be used to set access control rules during
+runtime. This can be useful for dynamic access control rules based on the state of the
+agent or the network.
+```python
+acl = AccessControlList(default=True, allowed={""}, blocked={""})
+
+@proto.on_message(model=Message, access_control_list=acl)
+async def message_handler(ctx: Context, sender: str, msg: Message):
+ if REASON_TO_BLOCK:
+ acl.blocked.add(sender)
+ ctx.logger.info(f"Received message from {sender}: {msg.text}")
+```
+
+"""
+
+import functools
+import time
+from typing import Optional, Set, Type, Union
+
+from pydantic import BaseModel
+
+from uagents import Context, Model, Protocol
+from uagents.models import ErrorMessage
+from uagents.storage import StorageAPI
+from uagents.types import MessageCallback
+
+WINDOW_SIZE_MINUTES = 60
+MAX_REQUESTS = 6
+
+
+class Usage(BaseModel):
+ time_window_start: float
+ window_size_minutes: int
+ requests: int
+ max_requests: int
+
+
+class RateLimit(BaseModel):
+ window_size_minutes: int
+ max_requests: int
+
+
+class AccessControlList(BaseModel):
+ default: bool
+ allowed: set[str]
+ blocked: set[str]
+
+
+class QuotaProtocol(Protocol):
+ def __init__(
+ self,
+ storage_reference: StorageAPI,
+ name: Optional[str] = None,
+ version: Optional[str] = None,
+ default_rate_limit: Optional[RateLimit] = None,
+ ):
+ """
+ Initialize a QuotaProtocol instance.
+
+ Args:
+ storage_reference (StorageAPI): The storage reference to use for rate limiting.
+ name (Optional[str], optional): The name of the protocol. Defaults to None.
+ version (Optional[str], optional): The version of the protocol. Defaults to None.
+ acl (Optional[AccessControlList], optional): The access control list. Defaults to None.
+ """
+ super().__init__(name=name, version=version)
+ self.storage_ref = storage_reference
+ self.default_rate_limit = default_rate_limit
+
+ def on_message(
+ self,
+ model: Type[Model],
+ replies: Optional[Union[Type[Model], Set[Type[Model]]]] = None,
+ allow_unverified: Optional[bool] = False,
+ rate_limit: Optional[RateLimit] = None,
+ access_control_list: Optional[AccessControlList] = None,
+ ):
+ """
+ Overwritten decorator to register a message handler for the protocol
+ including rate limiting.
+
+ Args:
+ model (Type[Model]): The message model type.
+ replies (Optional[Union[Type[Model], Set[Type[Model]]]], optional): The associated
+ reply types. Defaults to None.
+ allow_unverified (Optional[bool], optional): Whether to allow unverified messages.
+ Defaults to False.
+ rate_limit (Optional[RateLimit], optional): The rate limit to apply. Defaults to None.
+ access_control_list (Optional[AccessControlList], optional): The access control list to
+ apply.
+
+ Returns:
+ Callable: The decorator to register the message handler.
+ """
+
+ def decorator_on_message(func: MessageCallback):
+ handler = self.wrap(func, rate_limit, access_control_list)
+ self._add_message_handler(model, handler, replies, allow_unverified)
+ return handler
+
+ return decorator_on_message
+
+ def wrap(
+ self,
+ func: MessageCallback,
+ rate_limit: Optional[RateLimit] = None,
+ acl: Optional[AccessControlList] = None,
+ ) -> MessageCallback:
+ """
+ Decorator to wrap a function with rate limiting.
+
+ Args:
+ func: The function to wrap with rate limiting
+ rate_limit: The rate limit to apply
+ acl: The access control list to apply
+
+ Returns:
+ Callable: The decorated
+ """
+ if acl is None:
+ acl = AccessControlList(default=True, allowed=set(), blocked=set())
+
+ rate_limit = rate_limit or self.default_rate_limit
+
+ @functools.wraps(func)
+ async def decorator(ctx: Context, sender: str, msg: Type[Model]):
+ if (acl.default and sender in acl.blocked) or (
+ not acl.default and sender not in acl.allowed
+ ):
+ return await ctx.send(
+ sender,
+ ErrorMessage(
+ error=("You are not allowed to access this endpoint.")
+ ),
+ )
+ if not rate_limit or self.add_request(
+ sender,
+ func.__name__,
+ rate_limit.window_size_minutes,
+ rate_limit.max_requests,
+ ):
+ result = await func(ctx, sender, msg)
+ else:
+ result = await ctx.send(
+ sender,
+ ErrorMessage(
+ error=(
+ f"Rate limit exceeded for {msg.schema()["title"]}. "
+ f"This endpoint allows for {rate_limit.max_requests} calls per "
+ f"{rate_limit.window_size_minutes} minutes. Try again later."
+ )
+ ),
+ )
+ return result
+
+ return decorator # type: ignore
+
+ def _clean_usage(self, usage: dict[str, dict]):
+ """
+ Remove all time windows that are older than the current time window.
+
+ Args:
+ usage: The usage dictionary to clean
+ """
+ now = int(time.time())
+ for key in list(usage.keys()):
+ if (now - usage[key]["time_window_start"]) > usage[key][
+ "window_size_minutes"
+ ] * 60:
+ del usage[key]
+
+ def add_request(
+ self,
+ agent_address: str,
+ function_name: str,
+ window_size_minutes: int,
+ max_requests: int,
+ ) -> bool:
+ """
+ Add a request to the rate limiter if the current time is still within the
+ time window since the beginning of the most recent time window. Otherwise,
+ reset the time window and add the request.
+
+ Args:
+ agent_address: The address of the agent making the request
+
+ Returns:
+ False if the maximum number of requests has been exceeded, True otherwise
+ """
+
+ now = int(time.time())
+
+ usage = self.storage_ref.get(agent_address) or {}
+
+ if function_name in usage:
+ quota = Usage(**usage[function_name])
+ if (now - quota.time_window_start) <= window_size_minutes * 60:
+ if quota.requests >= max_requests:
+ return False
+ quota.requests += 1
+ else:
+ quota.time_window_start = now
+ quota.requests = 1
+ usage[function_name] = quota.model_dump()
+ else:
+ usage[function_name] = Usage(
+ time_window_start=now,
+ window_size_minutes=window_size_minutes,
+ requests=1,
+ max_requests=max_requests,
+ ).model_dump()
+
+ self._clean_usage(usage)
+
+ self.storage_ref.set(agent_address, usage)
+
+ return True