chore: add health check proto to post extractor agent (#48)

* chore: add health check proto to post extractor agent

* chore: add rate limit protocol

Author: jrriehl

1-uagents/utility/post-extractor-agent/agent.py

@@ -1,7 +1,9 @@
 import os
 import requests
 from datetime import datetime
+from enum import Enum
 from uagents import Agent, Context, Model, Protocol
+from uagents.experimental.quota import QuotaProtocol, RateLimit
 
 class RedditPostsRequest(Model):
     limit: int
@@ -18,6 +20,7 @@ class RedditPostsResponse(Model):
     posts: list[RedditPost]
 
 
+AGENT_NAME = "Post Extractor Agent"
 AGENT_SEED = os.getenv("AGENT_SEED", "your-post-agent-seed")
 REDDIT_ID = os.getenv("REDDIT_ID")
 REDDIT_SECRET = os.getenv("REDDIT_SECRET")
@@ -26,14 +29,18 @@ class RedditPostsResponse(Model):
 
 PORT = 8000
 agent = Agent(
-    name="Post Extractor Agent",
+    name=AGENT_NAME,
     seed=AGENT_SEED,
     port=PORT,
     endpoint=f"http://localhost:{PORT}/submit",
 )
 
-proto = Protocol(name="Post-Extractor", version="0.1.0")
-
+proto = QuotaProtocol(
+    storage_reference=agent.storage,
+    name="Post-Extractor",
+    version="0.1.0",
+    default_rate_limit=RateLimit(window_size_minutes=60, max_requests=6),
+)
 
 auth = requests.auth.HTTPBasicAuth(REDDIT_ID, REDDIT_SECRET)
 data = {
@@ -93,5 +100,51 @@ async def handle_request(ctx: Context, sender: str, msg: RedditPostsRequest):
 
 agent.include(proto)
 
+### Health check related code
+def agent_is_healthy() -> bool:
+    """
+    Implement the actual health check logic here.
+
+    For example, check if the agent can connect to a third party API,
+    check if the agent has enough resources, etc.
+    """
+    condition = True  # TODO: logic here
+    return bool(condition)
+
+
+class HealthCheck(Model):
+    pass
+
+
+class HealthStatus(str, Enum):
+    HEALTHY = "healthy"
+    UNHEALTHY = "unhealthy"
+
+
+class AgentHealth(Model):
+    agent_name: str
+    status: HealthStatus
+
+
+health_protocol = QuotaProtocol(
+    storage_reference=agent.storage, name="HealthProtocol", version="0.1.0"
+)
+
+
+@health_protocol.on_message(HealthCheck, replies={AgentHealth})
+async def handle_health_check(ctx: Context, sender: str, msg: HealthCheck):
+    status = HealthStatus.UNHEALTHY
+    try:
+        if agent_is_healthy():
+            status = HealthStatus.HEALTHY
+    except Exception as err:
+        ctx.logger.error(err)
+    finally:
+        await ctx.send(sender, AgentHealth(agent_name=AGENT_NAME, status=status))
+
+
+agent.include(health_protocol, publish_manifest=True)
+
+
 if __name__ == "__main__":
-    agent.run()
+    agent.run()