diff --git a/python/src/uagents/agent.py b/python/src/uagents/agent.py index c54b9ec5..7fa988c6 100644 --- a/python/src/uagents/agent.py +++ b/python/src/uagents/agent.py @@ -804,7 +804,11 @@ async def register(self): assert self._registration_policy is not None, "Agent has no registration policy" await self._registration_policy.register( - self.address, list(self.protocols.keys()), self._endpoints, self._metadata + self.address, + self._identity, + list(self.protocols.keys()), + self._endpoints, + self._metadata, ) async def _schedule_registration(self): diff --git a/python/src/uagents/registration.py b/python/src/uagents/registration.py index 62826517..2ce93f18 100644 --- a/python/src/uagents/registration.py +++ b/python/src/uagents/registration.py @@ -138,6 +138,7 @@ class AgentRegistrationPolicy(ABC): async def register( self, agent_address: str, + identity: Identity, protocols: List[str], endpoints: List[AgentEndpoint], metadata: Optional[Dict[str, Any]] = None, @@ -159,7 +160,6 @@ def add_agent(self, agent_info: AgentInfo, identity: Identity): class AlmanacApiRegistrationPolicy(AgentRegistrationPolicy): def __init__( self, - identity: Identity, *, almanac_api: Optional[str] = None, max_retries: int = ALMANAC_API_MAX_RETRIES, @@ -167,12 +167,12 @@ def __init__( ): self._almanac_api = almanac_api or ALMANAC_API_URL self._max_retries = max_retries - self._identity = identity self._logger = logger or logging.getLogger(__name__) async def register( self, agent_address: str, + identity: Identity, protocols: List[str], endpoints: List[AgentEndpoint], metadata: Optional[Dict[str, Any]] = None, @@ -186,7 +186,7 @@ async def register( ) # sign the attestation - attestation.sign(self._identity) + attestation.sign(identity) success = await almanac_api_post( f"{self._almanac_api}/agents", attestation, retries=self._max_retries @@ -233,7 +233,6 @@ async def register(self): class LedgerBasedRegistrationPolicy(AgentRegistrationPolicy): def __init__( self, - identity: Identity, ledger: LedgerClient, wallet: LocalWallet, almanac_contract: AlmanacContract, @@ -241,7 +240,6 @@ def __init__( *, logger: Optional[logging.Logger] = None, ): - self._identity = identity self._wallet = wallet self._ledger = ledger self._testnet = testnet @@ -265,6 +263,7 @@ def check_contract_version(self): async def register( self, agent_address: str, + identity: Identity, protocols: List[str], endpoints: List[AgentEndpoint], metadata: Optional[Dict[str, Any]] = None, @@ -301,7 +300,7 @@ async def register( current_time = int(time.time()) - ALMANAC_REGISTRATION_WAIT - signature = self._sign_registration(current_time) + signature = self._sign_registration(identity, current_time) await self._almanac_contract.register( self._ledger, self._wallet, @@ -318,7 +317,7 @@ async def register( def _get_balance(self) -> int: return self._ledger.query_bank_balance(Address(self._wallet.address())) - def _sign_registration(self, timestamp: int) -> str: + def _sign_registration(self, identity: Identity, timestamp: int) -> str: """ Sign the registration data for Almanac contract. @@ -333,7 +332,7 @@ def _sign_registration(self, timestamp: int) -> str: """ assert self._almanac_contract.address is not None - return self._identity.sign_registration( + return identity.sign_registration( str(self._almanac_contract.address), timestamp, str(self._wallet.address()), @@ -400,7 +399,6 @@ async def register(self): class DefaultRegistrationPolicy(AgentRegistrationPolicy): def __init__( self, - identity: Identity, ledger: LedgerClient, wallet: LocalWallet, almanac_contract: Optional[AlmanacContract], @@ -411,18 +409,19 @@ def __init__( ): self._logger = logger or logging.getLogger(__name__) self._api_policy = AlmanacApiRegistrationPolicy( - identity, almanac_api=almanac_api, logger=logger + almanac_api=almanac_api, logger=logger ) if almanac_contract is None: self._ledger_policy = None else: self._ledger_policy = LedgerBasedRegistrationPolicy( - identity, ledger, wallet, almanac_contract, testnet, logger=logger + ledger, wallet, almanac_contract, testnet, logger=logger ) async def register( self, agent_address: str, + identity: Identity, protocols: List[str], endpoints: List[AgentEndpoint], metadata: Optional[Dict[str, Any]] = None, @@ -430,7 +429,7 @@ async def register( # prefer the API registration policy as it is faster try: await self._api_policy.register( - agent_address, protocols, endpoints, metadata + agent_address, identity, protocols, endpoints, metadata ) except Exception as e: self._logger.warning( @@ -443,7 +442,7 @@ async def register( # schedule the ledger registration try: await self._ledger_policy.register( - agent_address, protocols, endpoints, metadata + agent_address, identity, protocols, endpoints, metadata ) except InsufficientFundsError: self._logger.warning(