Backend: support environment variables

backend/src/interfaces/agent.py

@@ -25,12 +25,12 @@ class UpdateAgentResponse(BaseModel):
 
 class PublicAgentData(BaseModel):
     id: str
+    subscription_id: str
     vm_hash: str | None
     last_update: int
 
 
 class Agent(PublicAgentData):
-    subscription_id: str
     encrypted_secret: str
     tags: list[str]
 

backend/src/main.py

@@ -1,3 +1,4 @@
+import json
 import time
 from http import HTTPStatus
 from uuid import uuid4
@@ -60,7 +61,7 @@ async def setup(body: SetupAgentBody) -> None:
 
     aleph_account = ETHAccount(config.ALEPH_SENDER_SK)
     async with AuthenticatedAlephHttpClient(
-        account=aleph_account, api_server=config.ALEPH_API_URL
+            account=aleph_account, api_server=config.ALEPH_API_URL
     ) as client:
         post_message, _ = await client.create_post(
             address=config.ALEPH_OWNER,
@@ -83,7 +84,10 @@ async def get_agent_public_info(agent_id: str) -> GetAgentResponse:
     agent = agents[0]
 
     return GetAgentResponse(
-        id=agent.id, vm_hash=agent.vm_hash, last_update=agent.last_update
+        id=agent.id,
+        vm_hash=agent.vm_hash,
+        last_update=agent.last_update,
+        subscription_id=agent.subscription_id,
     )
 
 
@@ -100,7 +104,7 @@ async def get_agent_secret(agent_id: str, signature: str) -> GetAgentSecretRespo
 
     async with aiohttp.ClientSession() as session:
         async with session.get(
-            url=f"{config.SUBSCRIPTION_BACKEND_URL}/subscriptions/{agent.subscription_id}"
+                url=f"{config.SUBSCRIPTION_BACKEND_URL}/subscriptions/{agent.subscription_id}"
         ) as response:
             data = await response.json()
             if response.status != HTTPStatus.OK:
@@ -138,11 +142,13 @@ def get_agent_secret_message(agent_id: str) -> GetAgentSecretMessage:
 
 @app.put("/agent/{agent_id}", description="Deploy an agent or update it")
 async def update(
-    agent_id: str,
-    secret: str = Form(),
-    code: UploadFile = File(...),
-    packages: UploadFile = File(...),
+        agent_id: str,
+        secret: str = Form(),
+        env_variables: str = Form(),  # actually dict[str, str] but the typing doesn't work well with forms
+        code: UploadFile = File(...),
+        packages: UploadFile = File(...),
 ) -> UpdateAgentResponse:
+    env_variables = json.loads(env_variables)
     agents = await fetch_agents([agent_id])
 
     if len(agents) != 1:
@@ -179,7 +185,7 @@ async def update(
     # Register the program
     aleph_account = ETHAccount(config.ALEPH_SENDER_SK)
     async with AuthenticatedAlephHttpClient(
-        account=aleph_account, api_server=config.ALEPH_API_URL
+            account=aleph_account, api_server=config.ALEPH_API_URL
     ) as client:
         vm_hash = agent.vm_hash
 
@@ -189,6 +195,7 @@ async def update(
                 program_ref=code_ref,
                 entrypoint="run",
                 runtime="63f07193e6ee9d207b7d1fcf8286f9aee34e6f12f101d2ec77c1229f92964696",
+                environment_variables=env_variables,
                 channel=config.ALEPH_CHANNEL,
                 encoding=Encoding.squashfs,
                 persistent=False,
@@ -204,6 +211,8 @@ async def update(
             )
             vm_hash = message.item_hash
 
+        # TODO: update env_vars also if vm already created
+
         # Updating the related POST message
         await client.create_post(
             address=config.ALEPH_OWNER,
@@ -233,10 +242,12 @@ async def delete(body: DeleteAgentBody):
 
     aleph_account = ETHAccount(config.ALEPH_SENDER_SK)
     async with AuthenticatedAlephHttpClient(
-        account=aleph_account, api_server=config.ALEPH_API_URL
+            account=aleph_account, api_server=config.ALEPH_API_URL
     ) as client:
+        # TODO: should we delete STORE messages of the code / deps too ?
         await client.forget(
             address=config.ALEPH_OWNER,
             hashes=[agent.vm_hash],
+            channel=config.ALEPH_CHANNEL,
             reason="LibertAI Agent subscription ended",
         )

backend/src/utils/agent.py

@@ -11,7 +11,7 @@ async def fetch_agents(ids: list[str] | None = None) -> list[FetchedAgent]:
         result = await client.get_posts(
             post_filter=PostFilter(
                 types=[config.ALEPH_AGENT_POST_TYPE],
-                addresses=[config.ALEPH_SENDER],
+                addresses=[config.ALEPH_OWNER],
                 tags=ids,
                 channels=[config.ALEPH_CHANNEL],
             )