Skip to content

Commit

Permalink
feat: created contract classes
Browse files Browse the repository at this point in the history
  • Loading branch information
Alejandro-Morales committed Apr 5, 2023
1 parent 3b70f13 commit b0e2c52
Show file tree
Hide file tree
Showing 3 changed files with 116 additions and 86 deletions.
100 changes: 20 additions & 80 deletions src/uagents/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
from uagents.storage import KeyValueStore, get_or_create_private_keys
from uagents.network import (
get_ledger,
get_reg_contract,
get_almanac_contract,
get_service_contract,
wait_for_tx_to_complete,
)
Expand All @@ -36,7 +36,6 @@
REGISTRATION_FEE,
REGISTRATION_DENOM,
LEDGER_PREFIX,
BLOCK_INTERVAL,
parse_endpoint_config,
parse_mailbox_config,
get_logger,
Expand All @@ -59,7 +58,6 @@ async def _handle_error(ctx: Context, destination: str, msg: ErrorMessage):
await ctx.send(destination, msg)


# pylint: disable=R0904
class Agent(Sink):
def __init__(
self,
Expand Down Expand Up @@ -111,7 +109,7 @@ def __init__(
]

self._ledger = get_ledger()
self._reg_contract = get_reg_contract()
self._almanac_contract = get_almanac_contract()
self._service_contract = get_service_contract()
self._storage = KeyValueStore(self.address[0:16])
self._interval_handlers: List[Tuple[IntervalCallback, float]] = []
Expand Down Expand Up @@ -183,9 +181,10 @@ def sign_digest(self, digest: bytes) -> str:
return self._identity.sign_digest(digest)

def sign_registration(self) -> str:
assert self._reg_contract.address is not None
assert self._almanac_contract.address is not None
return self._identity.sign_registration(
str(self._reg_contract.address), self.get_registration_sequence()
str(self._almanac_contract.address),
self._almanac_contract.get_sequence(self.address),
)

def update_loop(self, loop):
Expand All @@ -208,43 +207,25 @@ async def _register(self, ctx: Context):

self._logger.info("Registering on contract...")

if not self.is_name_available(ctx.name) and not self.is_owner(ctx.name):
if not self._service_contract.is_name_available(
ctx.name
) and not self._service_contract.is_owner(ctx.name, str(self.wallet.address())):
self._logger.error(
f"Please select another name for your agent, {ctx.name} is owned by another address"
)
return

transaction = Transaction()

almanac_msg = {
"register": {
"record": {
"service": {
"protocols": list(
map(lambda x: x.digest, self.protocols.values())
),
"endpoints": self._endpoints,
}
},
"signature": signature,
"sequence": self.get_registration_sequence(),
"agent_address": self.address,
}
}

ownership_msg = {
"update_ownership": {
"domain": f"{ctx.name}.agent",
"owner": {"address": {"address": str(ctx.wallet.address())}},
"permissions": "admin",
}
}

register_msg = {
"register": {
"domain": f"{ctx.name}.agent",
"agent_address": self.address,
}
}
almanac_msg = self._almanac_contract.get_registration_msg(
self.protocols, self._endpoints, signature, self.address
)
ownership_msg = self._service_contract.get_ownership_msg(
ctx.name, str(self.wallet.address())
)
register_msg = self._service_contract.get_registration_msg(
ctx.name, self.address
)

transaction.add_message(
create_cosmwasm_execute_msg(
Expand Down Expand Up @@ -273,32 +254,7 @@ async def _register(self, ctx: Context):
self._logger.info("Registering on contract...complete")

def _schedule_registration(self):
query_msg = {"query_records": {"agent_address": self.address}}
response = self._reg_contract.query(query_msg)

if not response["record"]:
contract_state = self._reg_contract.query({"query_contract_state": {}})
expiry = contract_state.get("state").get("expiry_height")
return expiry * BLOCK_INTERVAL

expiry = response.get("record")[0].get("expiry")
height = response.get("height")

return (expiry - height) * BLOCK_INTERVAL

def _is_almanac_registered(self) -> bool:
query_msg = {"query_records": {"agent_address": self.address}}
response = self._reg_contract.query(query_msg)

if not response["record"]:
return False
return True

def get_registration_sequence(self) -> int:
query_msg = {"query_sequence": {"agent_address": self.address}}
sequence = self._reg_contract.query(query_msg)["sequence"]

return sequence
return self._almanac_contract.get_expiry(self.address)

def get_agent_address(self, name: str) -> str:
query_msg = {"domain_record": {"domain": f"{name}.agent"}}
Expand All @@ -310,22 +266,6 @@ def get_agent_address(self, name: str) -> str:
return 0
return 1

def is_name_available(self, name: str):
query_msg = {"domain_record": {"domain": f"{name}.agent"}}
return self._service_contract.query(query_msg)["is_available"]

def is_owner(self, name: str):
query_msg = {
"permissions": {
"domain": f"{name}.agent",
"owner": {"address": {"address": str(self.wallet.address())}},
}
}
permission = self._service_contract.query(query_msg)["permissions"]
if permission == "admin":
return True
return False

def on_interval(
self,
period: float,
Expand Down Expand Up @@ -432,7 +372,7 @@ def start_background_tasks(self):
# start the contract registration update loop
if self._endpoints is not None:
if (
not self._is_almanac_registered()
not self._almanac_contract.is_registered(self.address)
or self._schedule_registration() < 3600
):
self._loop.create_task(
Expand Down
98 changes: 94 additions & 4 deletions src/uagents/network.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import asyncio
from datetime import datetime, timedelta
from typing import Optional
from typing import Optional, Dict, List, Union

from cosmpy.aerial.contract import LedgerContract
from cosmpy.aerial.client import (
Expand All @@ -18,9 +18,98 @@
CONTRACT_ALMANAC,
CONTRACT_SERVICE,
AGENT_NETWORK,
BLOCK_INTERVAL,
)


class AlmanacContract(LedgerContract):
def is_registered(self, address: str) -> bool:
query_msg = {"query_records": {"agent_address": address}}
response = self.query(query_msg)

if not response["record"]:
return False
return True

def get_expiry(self, address: str):
query_msg = {"query_records": {"agent_address": address}}
response = self.query(query_msg)

if not response["record"]:
contract_state = self.query({"query_contract_state": {}})
expiry = contract_state.get("state").get("expiry_height")
return expiry * BLOCK_INTERVAL

expiry = response.get("record")[0].get("expiry")
height = response.get("height")

return (expiry - height) * BLOCK_INTERVAL

def get_registration_msg(
self,
protocols: Dict,
endpoints: Optional[Union[List[str], Dict[str, dict]]],
signature: str,
address: str,
) -> dict:
return {
"register": {
"record": {
"service": {
"protocols": list(map(lambda x: x.digest, protocols.values())),
"endpoints": endpoints,
}
},
"signature": signature,
"sequence": self.get_sequence(address),
"agent_address": address,
}
}

def get_sequence(self, address: str) -> int:
query_msg = {"query_sequence": {"agent_address": address}}
sequence = self.query(query_msg)["sequence"]

return sequence


class ServiceContract(LedgerContract):
def is_name_available(self, name: str):
query_msg = {"domain_record": {"domain": f"{name}.agent"}}
return self.query(query_msg)["is_available"]

def is_owner(self, name: str, wallet_address: str):
query_msg = {
"permissions": {
"domain": f"{name}.agent",
"owner": {"address": {"address": wallet_address}},
}
}
permission = self.query(query_msg)["permissions"]
if permission == "admin":
return True
return False

@staticmethod
def get_ownership_msg(name: str, wallet_address: str):
return {
"update_ownership": {
"domain": f"{name}.agent",
"owner": {"address": {"address": wallet_address}},
"permissions": "admin",
}
}

@staticmethod
def get_registration_msg(name: str, address: str):
return {
"register": {
"domain": f"{name}.agent",
"agent_address": address,
}
}


if AGENT_NETWORK == AgentNetwork.FETCHAI_TESTNET:
_ledger = LedgerClient(NetworkConfig.fetchai_stable_testnet())
_faucet_api = FaucetApi(NetworkConfig.fetchai_stable_testnet())
Expand All @@ -29,8 +118,9 @@
else:
raise NotImplementedError

_almanac_contract = LedgerContract(None, _ledger, CONTRACT_ALMANAC)
_service_contract = LedgerContract(None, _ledger, CONTRACT_SERVICE)

_almanac_contract = AlmanacContract(None, _ledger, CONTRACT_ALMANAC)
_service_contract = ServiceContract(None, _ledger, CONTRACT_SERVICE)


def get_ledger() -> LedgerClient:
Expand All @@ -41,7 +131,7 @@ def get_faucet() -> FaucetApi:
return _faucet_api


def get_reg_contract() -> LedgerContract:
def get_almanac_contract() -> LedgerContract:
return _almanac_contract


Expand Down
4 changes: 2 additions & 2 deletions src/uagents/resolver.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,11 @@
from typing import Dict, Optional
import random

from uagents.network import get_reg_contract
from uagents.network import get_almanac_contract


def _query_record(agent_address: str, service: str) -> dict:
contract = get_reg_contract()
contract = get_almanac_contract()
query_msg = {
"query_record": {"agent_address": agent_address, "record_type": service}
}
Expand Down

0 comments on commit b0e2c52

Please sign in to comment.