feat: bureau overwrites agent endpoints (#94)

Co-authored-by: Ed FitzGerald <[email protected]>

src/uagents/agent.py

@@ -1,6 +1,6 @@
 import asyncio
 import functools
-from typing import Dict, List, Optional, Set, Union, Type, Tuple
+from typing import Dict, List, Optional, Set, Union, Type, Tuple, Any
 
 from cosmpy.aerial.wallet import LocalWallet, PrivateKey
 from cosmpy.crypto.address import Address
@@ -54,7 +54,7 @@ def __init__(
         name: Optional[str] = None,
         port: Optional[int] = None,
         seed: Optional[str] = None,
-        endpoint: Optional[Union[List[str], Dict[str, dict]]] = None,
+        endpoint: Optional[Union[str, List[str], Dict[str, dict]]] = None,
         mailbox: Optional[Union[str, Dict[str, str]]] = None,
         resolve: Optional[Resolver] = None,
         version: Optional[str] = None,
@@ -184,6 +184,9 @@ def sign_registration(self) -> str:
             str(self._reg_contract.address), self.get_registration_sequence()
         )
 
+    def update_endpoints(self, endpoints: List[Dict[str, Any]]):
+        self._endpoints = endpoints
+
     def update_loop(self, loop):
         self._loop = loop
 
@@ -336,6 +339,10 @@ 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):
@@ -355,10 +362,6 @@ def start_background_tasks(self):
             self._loop.create_task(
                 _run_interval(self._register, self._ctx, self._schedule_registration())
             )
-        else:
-            self._logger.warning(
-                "I have no endpoint and won't be able to receive external messages"
-            )
 
     def run(self):
         self.setup()
@@ -423,9 +426,14 @@ async def _process_message_queue(self):
 
 
 class Bureau:
-    def __init__(self, port: Optional[int] = None):
+    def __init__(
+        self,
+        port: Optional[int] = None,
+        endpoint: Optional[Union[str, List[str], Dict[str, dict]]] = None,
+    ):
         self._loop = asyncio.get_event_loop_policy().get_event_loop()
         self._agents = []
+        self._endpoints = parse_endpoint_config(endpoint)
         self._port = port or 8000
         self._queries: Dict[str, asyncio.Future] = {}
         self._logger = get_logger("bureau")
@@ -437,6 +445,8 @@ def add(self, agent: Agent):
         agent.update_queries(self._queries)
         if agent.mailbox is not None:
             self._use_mailbox = True
+        else:
+            agent.update_endpoints(self._endpoints)
         self._agents.append(agent)
 
     def run(self):

src/uagents/config.py

@@ -29,7 +29,7 @@ class AgentNetwork(Enum):
 
 
 def parse_endpoint_config(
-    endpoint: Optional[Union[List[str], Dict[str, dict]]]
+    endpoint: Optional[Union[str, List[str], Dict[str, dict]]]
 ) -> List[Dict[str, Any]]:
     if isinstance(endpoint, dict):
         endpoints = [
@@ -38,6 +38,8 @@ def parse_endpoint_config(
         ]
     elif isinstance(endpoint, list):
         endpoints = [{"url": val, "weight": 1} for val in endpoint]
+    elif isinstance(endpoint, str):
+        endpoints = [{"url": endpoint, "weight": 1}]
     else:
         endpoints = None
     return endpoints