Remove oop

api/src/backend/entities.py

@@ -314,13 +314,13 @@ class AgentStatus(Enum):
     scored = "scored"                                  # All evaluations complete
     replaced = "replaced"                              # Replaced by newer version
     pruned = "pruned"                                  # Pruned due to low score compared to top agent
-    
+
     # Legacy statuses for backward compatibility during transition
     awaiting_screening = "awaiting_screening_1"        # Map to stage 1
     screening = "screening_1"                          # Map to stage 1
     failed_screening = "failed_screening_1"            # Map to stage 1 fail
     evaluation = "evaluating"                          # Map to evaluating (legacy alias)
-    
+
     @classmethod
     def from_string(cls, status: str) -> 'AgentStatus':
         """Map database status string to agent state enum"""
@@ -344,7 +344,6 @@ def from_string(cls, status: str) -> 'AgentStatus':
         }
         return mapping.get(status, cls.awaiting_screening_1)
 
-
 class EvaluationStatus(Enum):
     waiting = "waiting"
     running = "running"

api/src/backend/queries/agents.py

@@ -3,7 +3,7 @@
 import asyncpg
 
 from api.src.backend.db_manager import db_operation, db_transaction
-from api.src.backend.entities import MinerAgent
+from api.src.backend.entities import AgentStatus, MinerAgent
 from api.src.utils.models import TopAgentHotkey
 from loggers.logging_utils import get_logger
 
@@ -139,9 +139,51 @@ async def set_approved_agents_to_awaiting_screening(conn: asyncpg.Connection) ->
     return [MinerAgent(**dict(result)) for result in results]
 
 @db_operation
-async def get_all_approved_version_ids(conn: asyncpg.Connection) -> List[str]:
-    """
-    Get all approved version IDs
-    """
-    data = await conn.fetch("SELECT version_id FROM approved_version_ids WHERE approved_at <= NOW()")
-    return [str(row["version_id"]) for row in data]
+async def set_agent_status(conn: asyncpg.Connection, version_id: str, status: str):
+    try:
+        AgentStatus(status) # Check whether the status we are trying to set to is valid 
+    except ValueError:
+        logger.error(f"Tried to set agent to invalid status {status!r}")
+        raise ValueError("Invalid status")
+
+    await conn.execute(
+        "UPDATE miner_agents SET status = $1 WHERE version_id = $2", 
+        status,
+        version_id
+    )
+
+@db_operation
+async def upload_miner_agent(
+    conn: asyncpg.Connection,
+    version_id: str,
+    miner_hotkey: str,
+    agent_name: str,
+    version_num: int,
+    ip_address: str
+):
+    await conn.execute(
+        """
+            INSERT INTO miner_agents (version_id, miner_hotkey, agent_name, version_num, created_at, status, ip_address)
+            VALUES ($1, $2, $3, $4, NOW(), 'awaiting_screening_1', $5)
+        """,
+        version_id,
+        miner_hotkey,
+        agent_name,
+        version_num,
+        ip_address
+    )
+
+@db_operation
+async def agent_startup_recovery(conn: asyncpg.Connection):
+    # Reset agent statuses for multi-stage screening
+    await conn.execute("UPDATE miner_agents SET status = 'awaiting_screening_1' WHERE status = 'screening_1'")
+    await conn.execute("UPDATE miner_agents SET status = 'awaiting_screening_2' WHERE status = 'screening_2'")
+    await conn.execute("UPDATE miner_agents SET status = 'waiting' WHERE status = 'evaluating'")
+
+    # Legacy status recovery for backward compatibility
+    await conn.execute("UPDATE miner_agents SET status = 'awaiting_screening_1' WHERE status = 'screening'")
+    await conn.execute("UPDATE miner_agents SET status = 'waiting' WHERE status = 'evaluation'")  # Legacy alias
+
+@db_operation
+async def set_agent_status_by_version_id(conn: asyncpg.Connection, version_id: str, status: str):
+    await conn.execute("UPDATE miner_agents SET status = $1 WHERE version_id = $2", status, version_id)

api/src/backend/queries/agents.pyi

@@ -12,4 +12,8 @@ async def get_agents_by_hotkey(miner_hotkey: str) -> List[MinerAgent]: ...
 async def ban_agents(miner_hotkeys: List[str], reason: str) -> None: ...
 async def approve_agent_version(version_id: str) -> None: ...
 async def set_approved_agents_to_awaiting_screening() -> List[MinerAgent]: ...
-async def get_all_approved_version_ids() -> List[str]: ...
+async def get_all_approved_version_ids() -> List[str]: ...
+async def set_agent_status(version_id: str, status: str): ...
+async def upload_miner_agent(version_id: str, miner_hotkey: str, agent_name: str, version_num: int, ip_address: str): ...
+async def agent_startup_recovery() -> None: ...
+async def set_agent_status_by_version_id(version_id: str, status: str): ...

api/src/backend/queries/evaluation_runs.py

@@ -324,8 +324,7 @@ async def reset_validator_evaluations(conn: asyncpg.Connection, version_id: str)
         WHERE evaluation_id = ANY($1::uuid[])
     """, evaluation_ids_to_cancel)
 
-
-
-
-
-
+@db_operation
+async def cancel_evaluation_runs(conn: asyncpg.Connection, evaluation_id: str):
+    """Cancel existing eval runs - e.g. for errored runs or disconnections etc"""
+    await conn.execute("UPDATE evaluation_runs SET status = 'cancelled' WHERE evaluation_id = $1", evaluation_id)

api/src/backend/queries/evaluation_runs.pyi

@@ -15,4 +15,6 @@ async def update_evaluation_run_logs(run_id: str, logs: str): ...
 async def get_evaluation_run_logs(run_id: str) -> str: ...
 
 async def fully_reset_evaluations(version_id: str): ...
-async def reset_validator_evaluations(version_id: str): ...
+async def reset_validator_evaluations(version_id: str): ...
+
+async def cancel_evaluation_runs(evaluation_id: str): ...