Fix: Refactor the ssh management.

backend/src/backend/blockchain.py

@@ -30,6 +30,10 @@
 w3 = Web3(Web3.HTTPProvider("https://mainnet.base.org"))
 
 
+def web3_from_wei(number: int, unit: str) -> int | Decimal:
+    return Web3.from_wei(number, unit)
+
+
 def convert_aleph_to_eth(required_tokens: Decimal) -> Decimal:
     aleph_pool_contract = w3.eth.contract(
         address=UNISWAP_ALEPH_POOL_ADDRESS, abi=POOL_ABI

backend/src/backend/main.py

@@ -5,11 +5,10 @@
 from starlette.middleware.cors import CORSMiddleware
 
 from aleph.sdk.client.authenticated_http import AuthenticatedAlephHttpClient
-from aleph.sdk.client.abstract import MessageFilter
-from aleph_message.models import ItemHash, Chain, MessageType
+from aleph_message.models import Chain
 
 from .agent import get_agent
-from .blockchain import CustomETHAccount
+from .blockchain import CustomETHAccount, web3_from_wei
 from .config import config
 from .models import AgentDeployment, AgentDeploymentStatus, AgentRequest
 from .orchestrator import DeploymentOrchestrator
@@ -131,12 +130,15 @@ async def deploy_agent(
     agent_proof_key = generate_predictable_key(agent_account_key)
     aleph_account = CustomETHAccount(agent_proof_key, chain=Chain.BASE)
     wallet_address = aleph_account.get_address()
-    eth_balance = aleph_account.get_eth_balance()
+    eth_balance = web3_from_wei(int(aleph_account.get_eth_balance()), "ether")
+
+    print(wallet_address)
+
     if eth_balance < MINIMUM_REQUIRED_AMOUNT:
         return {
             "error": True,
             "message": f"Insufficient balance, {'{0:.18f}'.format(MINIMUM_REQUIRED_AMOUNT)} ETH required "
-                       f"on wallet {wallet_address} instead {eth_balance}"
+                       f"on wallet {wallet_address} instead {'{0:.18f}'.format(eth_balance)}"
         }
 
     # Create and start the autonomous agent deployment
@@ -161,14 +163,14 @@ async def get_agent_info(
 ):
     """Get an agent by an agent ID"""
     deployment = orchestrator.get(agent_id)
-    if not deployment:
-        agent = await get_agent(agent_id)
-        if not agent:
-            return {
-                "error": True,
-                "message": f"Agent with id {agent_id} not found",
-            }
-    else:
-        agent = deployment.deployment
+    if deployment:
+        return deployment.deployment
+
+    agent = await get_agent(agent_id)
+    if not agent:
+        return {
+            "error": True,
+            "message": f"Agent with id {agent_id} not found",
+        }
 
     return agent

backend/src/backend/orchestrator.py

@@ -191,7 +191,13 @@ async def deploy_code(self):
                     raise
 
     async def _ssh_deployment(self):
-        await agent_ssh_deployment(agent_orchestration=self)
+        await agent_ssh_deployment(
+            deployment=self.deployment,
+            aleph_account=self.aleph_account,
+            ssh_private_key=self.ssh_private_key,
+            creator_wallet=self.creator_wallet,
+            env_variables=self.env_variables,
+        )
 
         self.deployment.status = AgentDeploymentStatus.ALIVE
         self.deployment.last_update = int(time.time())

backend/src/backend/ssh.py

@@ -1,22 +1,32 @@
 import io
+from typing import Dict
+
 import paramiko
 
+from aleph.sdk.chains.ethereum import ETHAccount
+
 from backend.agent import generate_fixed_env_variables, generate_env_file_content
 from backend.aleph import get_code_hash, get_code_file
 from backend.config import config
-from backend.orchestrator import AgentOrchestration
+from backend.models import FetchedAgentDeployment
 
 
-async def agent_ssh_deployment(agent_orchestration: AgentOrchestration):
+async def agent_ssh_deployment(
+        deployment: FetchedAgentDeployment,
+        aleph_account: ETHAccount,
+        ssh_private_key: str,
+        creator_wallet: str,
+        env_variables: Dict[str, str]
+):
     # Create a Paramiko SSH client
     ssh_client = paramiko.SSHClient()
     ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
 
     # Load private key from string
-    rsa_key = paramiko.RSAKey(file_obj=io.StringIO(agent_orchestration.ssh_private_key))
+    rsa_key = paramiko.RSAKey(file_obj=io.StringIO(ssh_private_key))
 
     # Get code file
-    agent_hash = agent_orchestration.deployment.agent_hash
+    agent_hash = deployment.agent_hash
     code_hash = await get_code_hash(agent_hash)
     if not code_hash:
         raise ValueError(f"Code hash not found for Agent hash {agent_hash}")
@@ -26,7 +36,7 @@ async def agent_ssh_deployment(agent_orchestration: AgentOrchestration):
 
     try:
         # Connect to the server
-        ssh_client.connect(hostname=agent_orchestration.deployment.instance_ip, username="root", pkey=rsa_key)
+        ssh_client.connect(hostname=deployment.instance_ip, username="root", pkey=rsa_key)
 
         # Send the zip with the code
         sftp = ssh_client.open_sftp()
@@ -36,14 +46,14 @@ async def agent_ssh_deployment(agent_orchestration: AgentOrchestration):
 
         # Send the env variable file
         sftp = ssh_client.open_sftp()
-        wallet_private_key = agent_orchestration.aleph_account.export_private_key()
+        wallet_private_key = aleph_account.export_private_key()
 
         fixed_env_variables = generate_fixed_env_variables(
             private_key=wallet_private_key,
-            creator_address=agent_orchestration.creator_wallet,
-            owner_address=agent_orchestration.deployment.owner,
+            creator_address=creator_wallet,
+            owner_address=deployment.owner,
         )
-        content = generate_env_file_content(fixed_env_variables, agent_orchestration.env_variables)
+        content = generate_env_file_content(fixed_env_variables, env_variables)
         remote_path = "/tmp/.env"
         sftp.putfo(io.BytesIO(content), remote_path)
         sftp.close()