refactor(core): separate prefix from address in agent info

fetchai · Jan 17, 2025 · a758d57 · a758d57
1 parent 197296f
commit a758d57
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 12 deletions.
diff --git a/python/src/uagents/agent.py b/python/src/uagents/agent.py
@@ -422,7 +422,8 @@ async def _handle_error_message(ctx: Context, sender: str, msg: ErrorMessage):
             @self.on_rest_get("/agent_info", AgentInfo)  # type: ignore
             async def _handle_get_info(_ctx: Context):
                 return AgentInfo(
-                    identifier=self.identifier,
+                    address=self.address,
+                    prefix=TESTNET_PREFIX if self._test else MAINNET_PREFIX,
                     endpoints=self._endpoints,
                     protocols=list(self.protocols.keys()),
                 )
@@ -672,7 +673,8 @@ def info(self) -> AgentInfo:
             AgentInfo: The agent's address, endpoints, protocols, and metadata.
         """
         return AgentInfo(
-            identifier=self.identifier,
+            address=self.address,
+            prefix=TESTNET_PREFIX if self._test else MAINNET_PREFIX,
             endpoints=self._endpoints,
             protocols=list(self.protocols.keys()),
             metadata=self.metadata,

diff --git a/python/src/uagents/mailbox.py b/python/src/uagents/mailbox.py
@@ -12,7 +12,7 @@
 from uagents.dispatch import dispatcher
 from uagents.envelope import Envelope
 from uagents.models import Model
-from uagents.types import AgentEndpoint
+from uagents.types import AddressPrefix, AgentEndpoint
 from uagents.utils import get_logger
 
 logger = get_logger("mailbox")
@@ -46,9 +46,6 @@ class ChallengeProofResponse(Model):
     expiry: str
 
 
-AddressPrefix = Literal["agent", "test-agent"]
-
-
 class RegistrationRequest(BaseModel):
     address: str
     prefix: Optional[AddressPrefix] = "test-agent"

diff --git a/python/src/uagents/network.py b/python/src/uagents/network.py
@@ -446,7 +446,7 @@ async def register_batch(
                 endpoints=record.endpoints,
                 signature=record.signature,
                 sequence=record.timestamp,
-                address=record.identifier,
+                address=record.address,
             )
 
             denom = self._client.network_config.fee_denomination

diff --git a/python/src/uagents/registration.py b/python/src/uagents/registration.py
@@ -209,7 +209,7 @@ def __init__(
 
     def add_agent(self, agent_info: AgentInfo, identity: Identity):
         attestation = AgentRegistrationAttestation(
-            agent_identifier=agent_info.identifier,
+            agent_identifier=f"{agent_info.prefix}://{agent_info.address}",
             protocols=list(agent_info.protocols),
             endpoints=agent_info.endpoints,
             metadata=coerce_metadata_to_str(extract_geo_metadata(agent_info.metadata)),
@@ -362,22 +362,23 @@ def __init__(
 
     def add_agent(self, agent_info: AgentInfo, identity: Identity):
         agent_record = AlmanacContractRecord(
-            identifier=agent_info.identifier,
+            address=agent_info.address,
+            prefix=agent_info.prefix,
             protocols=agent_info.protocols,
             endpoints=agent_info.endpoints,
             contract_address=str(self._almanac_contract.address),
             sender_address=str(self._wallet.address()),
         )
         self._records.append(agent_record)
-        self._identities[agent_info.identifier] = identity
+        self._identities[agent_info.address] = identity
 
     def _get_balance(self) -> int:
         return self._ledger.query_bank_balance(Address(self._wallet.address()))
 
     async def register(self):
         self._logger.info("Registering agents on Almanac contract...")
         for record in self._records:
-            _, _, agent_address = parse_identifier(record.identifier)
+            _, _, agent_address = parse_identifier(record.address)
             record.sign(self._identities[agent_address])
 
         if self._get_balance() < REGISTRATION_FEE * len(self._records):

diff --git a/python/src/uagents/types.py b/python/src/uagents/types.py
@@ -38,13 +38,17 @@
 RestHandlerMap = Dict[Tuple[RestMethod, str], RestHandler]
 
 
+AddressPrefix = Literal["agent", "test-agent"]
+
+
 class AgentEndpoint(BaseModel):
     url: str
     weight: int
 
 
 class AgentInfo(BaseModel):
-    identifier: str
+    address: str
+    prefix: AddressPrefix
     endpoints: List[AgentEndpoint]
     protocols: List[str]
     metadata: Optional[Dict[str, Any]] = None