feat: add wallet messaging (#53)

Co-authored-by: UrAvgDeveloper <[email protected]>

.github/workflows/ci-tests.yml

@@ -107,8 +107,8 @@ jobs:
           key: venv-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ hashFiles('**/poetry.lock') }}
       - name: Install dependencies
         if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'
-        run: poetry install --no-interaction --no-root
+        run: poetry install -E all --no-interaction --no-root
 
-      - run: poetry install --no-interaction
+      - run: poetry install -E all --no-interaction
       - run: poetry run black --check .
       - run: poetry run pylint $(git ls-files '*.py')

python/.pylintrc

@@ -11,6 +11,6 @@ disable=missing-module-docstring,
         too-many-locals,
         too-many-return-statements,
         logging-fstring-interpolation,
+        import-outside-toplevel,
         too-many-lines,
         broad-exception-caught
-

python/examples/15-wallet-messaging/main.py

@@ -0,0 +1,32 @@
+from uagents import Agent, Bureau, Context
+from uagents.wallet_messaging import WalletMessage
+
+
+ALICE_SEED = "alice dorado recovery phrase"
+BOB_SEED = "bob dorado recovery phrase"
+
+alice = Agent(name="alice", seed=ALICE_SEED, enable_wallet_messaging=True)
+bob = Agent(name="bob", seed=BOB_SEED, enable_wallet_messaging=True)
+
+
+@alice.on_wallet_message()
+async def reply(ctx: Context, msg: WalletMessage):
+    ctx.logger.info(f"Got wallet message: {msg.text}")
+    await ctx.send_wallet_message(msg.sender, "hey, thanks for the message")
+
+
+@bob.on_interval(period=5)
+async def send_message(ctx: Context):
+    ctx.logger.info("Sending message...")
+    await ctx.send_wallet_message(alice.address, "hello")
+
+
+@bob.on_wallet_message()
+async def wallet_reply(ctx: Context, msg: WalletMessage):
+    ctx.logger.info(f"Got wallet message: {msg.text}")
+
+
+bureau = Bureau()
+bureau.add(alice)
+bureau.add(bob)
+bureau.run()

python/pyproject.toml

@@ -18,6 +18,12 @@ aiohttp = "^3.8.3"
 cosmpy = "^0.9.1"
 websockets = "^10.4"
 
+# extras
+fetchai-babble = {version = "^0.4.0", optional = true}
+tortoise-orm = {version = "^0.19.2", optional = true}
+geopy = {version = "^2.3.0", optional = true}
+pyngrok = {version = "^5.2.3", optional = true}
+
 [tool.poetry.dev-dependencies]
 aioresponses = "^0.7.4"
 black = "^23.1.0"
@@ -26,14 +32,13 @@ pylint = "^2.15.3"
 mkdocs = "^1.4.2"
 mkdocs-material = "^9.1.13"
 
-[tool.poetry.group.orm.dependencies]
-tortoise-orm = "^0.19.2"
-
-[tool.poetry.group.geo.dependencies]
-geopy = "^2.3.0"
 
-[tool.poetry.group.remote-agents.dependencies]
-pyngrok = "^5.2.3"
+[tool.poetry.extras]
+all = ["tortoise-orm", "geopy", "fetchai-babble", "pyngrok"]
+wallet = ["fetchai-babble"]
+orm = ["tortoise-orm"]
+geo = ["geopy"]
+remote-agents = ["pyngrok"]
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

python/src/uagents/agent.py

@@ -156,6 +156,8 @@ def __init__(
         agentverse: Optional[Union[str, Dict[str, str]]] = None,
         mailbox: Optional[Union[str, Dict[str, str]]] = None,
         resolve: Optional[Resolver] = None,
+        enable_wallet_messaging: Optional[Union[bool, Dict[str, str]]] = False,
+        wallet_key_derivation_index: Optional[int] = 0,
         max_resolver_endpoints: Optional[int] = None,
         version: Optional[str] = None,
         test: Optional[bool] = True,
@@ -172,6 +174,10 @@ def __init__(
             agentverse (Optional[Union[str, Dict[str, str]]]): The agentverse configuration.
             mailbox (Optional[Union[str, Dict[str, str]]]): The mailbox configuration.
             resolve (Optional[Resolver]): The resolver to use for agent communication.
+            enable_wallet_messaging (Optional[Union[bool, Dict[str, str]]]): Whether to enable
+            wallet messaging. If '{"chain_id": CHAIN_ID}' is provided, this sets the chain ID for
+            the messaging server.
+            wallet_key_derivation_index (Optional[int]): The index used for deriving the wallet key.
             max_resolver_endpoints (Optional[int]): The maximum number of endpoints to resolve.
             version (Optional[str]): The version of the agent.
         """
@@ -189,8 +195,8 @@ def __init__(
         else:
             self._loop = asyncio.get_event_loop_policy().get_event_loop()
 
-        self._initialize_wallet_and_identity(seed, name)
-
+        # initialize wallet and identity
+        self._initialize_wallet_and_identity(seed, name, wallet_key_derivation_index)
         self._logger = get_logger(self.name)
 
         # configure endpoints and mailbox
@@ -237,6 +243,8 @@ def __init__(
         self._test = test
         self._version = version or "0.1.0"
 
+        self.initialize_wallet_messaging(enable_wallet_messaging)
+
         # initialize the internal agent protocol
         self._protocol = Protocol(name=self._name, version=self._version)
 
@@ -255,6 +263,7 @@ def __init__(
             self._queries,
             replies=self._replies,
             interval_messages=self._interval_messages,
+            wallet_messaging_client=self._wallet_messaging_client,
             protocols=self.protocols,
             logger=self._logger,
         )
@@ -272,7 +281,7 @@ def __init__(
         async def _handle_error_message(ctx: Context, sender: str, msg: ErrorMessage):
             ctx.logger.exception(f"Received error message from {sender}: {msg.error}")
 
-    def _initialize_wallet_and_identity(self, seed, name):
+    def _initialize_wallet_and_identity(self, seed, name, wallet_key_derivation_index):
         """
         Initialize the wallet and identity for the agent.
 
@@ -282,6 +291,7 @@ def _initialize_wallet_and_identity(self, seed, name):
         Args:
             seed (str or None): The seed for generating keys.
             name (str or None): The name of the agent.
+            wallet_key_derivation_index (int): The index for deriving the wallet key.
         """
         if seed is None:
             if name is None:
@@ -294,12 +304,51 @@ def _initialize_wallet_and_identity(self, seed, name):
         else:
             self._identity = Identity.from_seed(seed, 0)
             self._wallet = LocalWallet(
-                PrivateKey(derive_key_from_seed(seed, LEDGER_PREFIX, 0)),
+                PrivateKey(
+                    derive_key_from_seed(
+                        seed, LEDGER_PREFIX, wallet_key_derivation_index
+                    )
+                ),
                 prefix=LEDGER_PREFIX,
             )
         if name is None:
             self._name = self.address[0:16]
 
+    def initialize_wallet_messaging(
+        self, enable_wallet_messaging: Union[bool, Dict[str, str]]
+    ):
+        """
+        Initialize wallet messaging for the agent.
+
+        Args:
+            enable_wallet_messaging (Union[bool, Dict[str, str]]): Wallet messaging configuration.
+        """
+        if enable_wallet_messaging:
+            wallet_chain_id = self._ledger.network_config.chain_id
+            if (
+                isinstance(enable_wallet_messaging, dict)
+                and "chain_id" in enable_wallet_messaging
+            ):
+                wallet_chain_id = enable_wallet_messaging["chain_id"]
+
+            try:
+                from uagents.wallet_messaging import WalletMessagingClient
+
+                self._wallet_messaging_client = WalletMessagingClient(
+                    self._identity,
+                    self._wallet,
+                    wallet_chain_id,
+                    self._logger,
+                )
+            except ModuleNotFoundError:
+                self._logger.exception(
+                    "Unable to include wallet messaging. "
+                    "Please install the 'wallet' extra to enable wallet messaging."
+                )
+                self._wallet_messaging_client = None
+        else:
+            self._wallet_messaging_client = None
+
     @property
     def name(self) -> str:
         """
@@ -679,6 +728,16 @@ def _add_event_handler(
         elif event_type == "shutdown":
             self._on_shutdown.append(func)
 
+    def on_wallet_message(
+        self,
+    ):
+        if self._wallet_messaging_client is None:
+            self._logger.warning(
+                "Discarding 'on_wallet_message' handler because wallet messaging is disabled"
+            )
+            return lambda func: func
+        return self._wallet_messaging_client.on_message()
+
     def include(self, protocol: Protocol, publish_manifest: Optional[bool] = False):
         """
         Include a protocol into the agent's capabilities.
@@ -802,10 +861,6 @@ def setup(self):
         # register the internal agent protocol
         self.include(self._protocol)
         self._loop.run_until_complete(self._startup())
-        if self._endpoints is None:
-            self._logger.warning(
-                "I have no endpoint and won't be able to receive external messages"
-            )
         self.start_background_tasks()
 
     def start_background_tasks(self):
@@ -824,6 +879,16 @@ def start_background_tasks(self):
         self._background_tasks.add(task)
         task.add_done_callback(self._background_tasks.discard)
 
+        # start the wallet messaging client if enabled
+        if self._wallet_messaging_client is not None:
+            for task in [
+                self._wallet_messaging_client.poll_server(),
+                self._wallet_messaging_client.process_message_queue(self._ctx),
+            ]:
+                task = self._loop.create_task(task)
+                self._background_tasks.add(task)
+                task.add_done_callback(self._background_tasks.discard)
+
     def run(self):
         """
         Run the agent.

python/src/uagents/config.py

@@ -37,6 +37,8 @@
 ALMANAC_API_URL = AGENTVERSE_URL + "/v1/almanac/"
 MAILBOX_POLL_INTERVAL_SECONDS = 1.0
 
+WALLET_MESSAGING_POLL_INTERVAL_SECONDS = 2.0
+
 RESPONSE_TIME_HINT_SECONDS = 5
 DEFAULT_ENVELOPE_TIMEOUT_SECONDS = 30
 DEFAULT_MAX_ENDPOINTS = 10

python/src/uagents/context.py

@@ -44,6 +44,7 @@
 IntervalCallback = Callable[["Context"], Awaitable[None]]
 MessageCallback = Callable[["Context", str, Any], Awaitable[None]]
 EventCallback = Callable[["Context"], Awaitable[None]]
+WalletMessageCallback = Callable[["Context", Any], Awaitable[None]]
 
 
 class DeliveryStatus(str, Enum):
@@ -147,6 +148,7 @@ def __init__(
         replies: Optional[Dict[str, Dict[str, Type[Model]]]] = None,
         interval_messages: Optional[Set[str]] = None,
         message_received: Optional[MsgDigest] = None,
+        wallet_messaging_client: Optional[Any] = None,
         protocols: Optional[Dict[str, Protocol]] = None,
         logger: Optional[logging.Logger] = None,
     ):
@@ -168,6 +170,7 @@ def __init__(
             for each type of incoming message.
             interval_messages (Optional[Set[str]]): The optional set of interval messages.
             message_received (Optional[MsgDigest]): The optional message digest received.
+            wallet_messaging_client (Optional[Any]): The optional wallet messaging client.
             protocols (Optional[Dict[str, Protocol]]): The optional dictionary of protocols.
             logger (Optional[logging.Logger]): The optional logger instance.
         """
@@ -184,6 +187,7 @@ def __init__(
         self._replies = replies
         self._interval_messages = interval_messages
         self._message_received = message_received
+        self._wallet_messaging_client = wallet_messaging_client
         self._protocols = protocols or {}
         self._logger = logger
 
@@ -554,3 +558,14 @@ async def send_raw_exchange_envelope(
             destination=destination,
             endpoint="",
         )
+
+    async def send_wallet_message(
+        self,
+        destination: str,
+        text: str,
+        msg_type: int = 1,
+    ):
+        if self._wallet_messaging_client is not None:
+            await self._wallet_messaging_client.send(destination, text, msg_type)
+        else:
+            self.logger.warning("Cannot send wallet message: no client available")

python/src/uagents/crypto/__init__.py

@@ -2,8 +2,12 @@
 import struct
 from secrets import token_bytes
 from typing import Tuple, Union
+import json
+import base64
+
 import bech32
 import ecdsa
+from ecdsa.util import sigencode_string_canonize
 
 from uagents.config import USER_PREFIX
 
@@ -78,6 +82,7 @@ def __init__(self, signing_key: ecdsa.SigningKey):
         # build the address
         pub_key_bytes = self._sk.get_verifying_key().to_string(encoding="compressed")
         self._address = _encode_bech32("agent", pub_key_bytes)
+        self._pub_key = pub_key_bytes.hex()
 
     @staticmethod
     def from_seed(seed: str, index: int) -> "Identity":
@@ -126,12 +131,20 @@ def address(self) -> str:
         """
         return self._address
 
+    @property
+    def pub_key(self) -> str:
+        return self._pub_key
+
     def sign(self, data: bytes) -> str:
         """
         Sign the provided data.
         """
         return _encode_bech32("sig", self._sk.sign(data))
 
+    def sign_b64(self, data: bytes) -> str:
+        raw_signature = bytes(self._sk.sign(data, sigencode=sigencode_string_canonize))
+        return base64.b64encode(raw_signature).decode()
+
     def sign_digest(self, digest: bytes) -> str:
         """
         Sign the provided digest.
@@ -148,6 +161,36 @@ def sign_registration(self, contract_address: str, sequence: int) -> str:
         hasher.update(encode_length_prefixed(sequence))
         return self.sign_digest(hasher.digest())
 
+    def sign_arbitrary(self, data: bytes) -> Tuple[str, str]:
+        # create the sign doc
+        sign_doc = {
+            "chain_id": "",
+            "account_number": "0",
+            "sequence": "0",
+            "fee": {
+                "gas": "0",
+                "amount": [],
+            },
+            "msgs": [
+                {
+                    "type": "sign/MsgSignData",
+                    "value": {
+                        "signer": self.address,
+                        "data": base64.b64encode(data).decode(),
+                    },
+                },
+            ],
+            "memo": "",
+        }
+
+        raw_sign_doc = json.dumps(
+            sign_doc, sort_keys=True, separators=(",", ":")
+        ).encode()
+        signature = self.sign_b64(raw_sign_doc)
+        enc_sign_doc = base64.b64encode(raw_sign_doc).decode()
+
+        return enc_sign_doc, signature
+
     @staticmethod
     def verify_digest(address: str, digest: bytes, signature: str) -> bool:
         """

python/src/uagents/network.py

@@ -15,8 +15,8 @@
 from cosmpy.aerial.exceptions import NotFoundError, QueryTimeoutError
 from cosmpy.aerial.contract.cosmwasm import create_cosmwasm_execute_msg
 from cosmpy.aerial.faucet import FaucetApi
-from cosmpy.aerial.tx_helpers import TxResponse
 from cosmpy.aerial.tx import Transaction
+from cosmpy.aerial.tx_helpers import TxResponse
 from cosmpy.aerial.wallet import LocalWallet
 
 from uagents.config import (