refactor(BA-3028): Move view objects up

Author: hhoikoo

src/ai/backend/agent/kernel.py

@@ -461,49 +461,6 @@ class AgentKernelRegistryKey:
     kernel_id: KernelId
 
 
-class KernelRegistry(MutableMapping[AgentKernelRegistryKey, AbstractKernel]):
-    _registry: MutableMapping[AgentKernelRegistryKey, AbstractKernel]
-    _global_registry: MutableMapping[KernelId, AbstractKernel]
-
-    def __init__(self) -> None:
-        super().__init__()
-
-        self._registry = {}
-        self._global_registry = {}
-
-    def agent_mapping(self, agent_id: AgentId) -> "KernelRegistryAgentMapping":
-        return KernelRegistryAgentMapping(self, agent_id)
-
-    def global_view(self) -> "KernelRegistryGlobalView":
-        return KernelRegistryGlobalView(self)
-
-    @overload
-    def __getitem__(self, key: KernelId) -> AbstractKernel: ...
-
-    @overload
-    def __getitem__(self, key: AgentKernelRegistryKey) -> AbstractKernel: ...
-
-    def __getitem__(self, key: KernelId | AgentKernelRegistryKey) -> AbstractKernel:
-        if isinstance(key, AgentKernelRegistryKey):
-            return self._registry[key]
-        else:
-            return self._global_registry[key]
-
-    def __setitem__(self, key: AgentKernelRegistryKey, value: AbstractKernel) -> None:
-        self._registry[key] = value
-        self._global_registry[key.kernel_id] = value
-
-    def __delitem__(self, key: AgentKernelRegistryKey) -> None:
-        del self._registry[key]
-        del self._global_registry[key.kernel_id]
-
-    def __iter__(self) -> Iterator[AgentKernelRegistryKey]:
-        return iter(self._registry)
-
-    def __len__(self) -> int:
-        return len(self._registry)
-
-
 class KernelRegistryAgentMapping(MutableMapping[KernelId, AbstractKernel]):
     _registry: KernelRegistry
     _agent_id: AgentId
@@ -551,6 +508,49 @@ def __len__(self) -> int:
         return len(self._registry)
 
 
+class KernelRegistry(MutableMapping[AgentKernelRegistryKey, AbstractKernel]):
+    _registry: MutableMapping[AgentKernelRegistryKey, AbstractKernel]
+    _global_registry: MutableMapping[KernelId, AbstractKernel]
+
+    def __init__(self) -> None:
+        super().__init__()
+
+        self._registry = {}
+        self._global_registry = {}
+
+    def agent_mapping(self, agent_id: AgentId) -> KernelRegistryAgentMapping:
+        return KernelRegistryAgentMapping(self, agent_id)
+
+    def global_view(self) -> KernelRegistryGlobalView:
+        return KernelRegistryGlobalView(self)
+
+    @overload
+    def __getitem__(self, key: KernelId) -> AbstractKernel: ...
+
+    @overload
+    def __getitem__(self, key: AgentKernelRegistryKey) -> AbstractKernel: ...
+
+    def __getitem__(self, key: KernelId | AgentKernelRegistryKey) -> AbstractKernel:
+        if isinstance(key, AgentKernelRegistryKey):
+            return self._registry[key]
+        else:
+            return self._global_registry[key]
+
+    def __setitem__(self, key: AgentKernelRegistryKey, value: AbstractKernel) -> None:
+        self._registry[key] = value
+        self._global_registry[key.kernel_id] = value
+
+    def __delitem__(self, key: AgentKernelRegistryKey) -> None:
+        del self._registry[key]
+        del self._global_registry[key.kernel_id]
+
+    def __iter__(self) -> Iterator[AgentKernelRegistryKey]:
+        return iter(self._registry)
+
+    def __len__(self) -> int:
+        return len(self._registry)
+
+
 _zctx = None