Skip to content

feat(core): add is active almanac update #564

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 10 commits into from
Oct 22, 2024
20 changes: 16 additions & 4 deletions python/src/uagents/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,12 @@
from uagents.models import ErrorMessage, Model
from uagents.network import InsufficientFundsError, get_almanac_contract, get_ledger
from uagents.protocol import Protocol
from uagents.registration import AgentRegistrationPolicy, DefaultRegistrationPolicy
from uagents.registration import (
AgentRegistrationPolicy,
AgentStatusUpdate,
DefaultRegistrationPolicy,
update_agent_status,
)
from uagents.resolver import GlobalResolver, Resolver
from uagents.storage import KeyValueStore, get_or_create_private_keys
from uagents.types import (
Expand Down Expand Up @@ -339,10 +344,10 @@ def __init__(
else:
self._mailbox_client = None

almanac_api_url = f"{self._agentverse['http_prefix']}://{self._agentverse['base_url']}/v1/almanac"
self._almanac_api_url = f"{self._agentverse['http_prefix']}://{self._agentverse['base_url']}/v1/almanac"
self._resolver = resolve or GlobalResolver(
max_endpoints=max_resolver_endpoints,
almanac_api_url=almanac_api_url,
almanac_api_url=self._almanac_api_url,
)

self._ledger = get_ledger(test)
Expand Down Expand Up @@ -371,7 +376,7 @@ def __init__(
self._almanac_contract,
self._test,
logger=self._logger,
almanac_api=almanac_api_url,
almanac_api=self._almanac_api_url,
)
self._metadata = self._initialize_metadata(metadata)

Expand Down Expand Up @@ -1088,6 +1093,13 @@ async def _shutdown(self):
Perform shutdown actions.

"""
try:
status = AgentStatusUpdate(agent_address=self.address, is_active=False)
status.sign(self._identity)
await update_agent_status(status, self._almanac_api_url)
except Exception as ex:
self._logger.exception(f"Failed to update agent registration status: {ex}")

for handler in self._on_shutdown:
try:
ctx = self._build_context()
Expand Down
37 changes: 37 additions & 0 deletions python/src/uagents/registration.py
Original file line number Diff line number Diff line change
Expand Up @@ -293,3 +293,40 @@ async def register(
except Exception as e:
self._logger.error(f"Failed to register on Almanac contract: {e}")
raise


class AgentStatusUpdate(BaseModel):
agent_address: str
is_active: bool
signature: Optional[str] = None

def sign(self, identity: Identity):
digest = self._build_digest()
self.signature = identity.sign_digest(digest)

def verify(self) -> bool:
return self.signature is not None and Identity.verify_digest(
self.agent_address, self._build_digest(), self.signature
)

def _build_digest(self) -> bytes:
sha256 = hashlib.sha256()
sha256.update(
json.dumps(
self.model_dump(exclude={"signature"}),
sort_keys=True,
separators=(",", ":"),
).encode("utf-8")
)
return sha256.digest()


async def update_agent_status(status: AgentStatusUpdate, almanac_api: str):
async with aiohttp.ClientSession() as session: # noqa: SIM117
async with session.post(
f"{almanac_api}/agents/{status.agent_address}/status",
headers={"content-type": "application/json"},
data=status.model_dump_json(),
timeout=aiohttp.ClientTimeout(total=ALMANAC_API_TIMEOUT_SECONDS),
) as resp:
resp.raise_for_status()
Loading