Using auto-discovery to build context options in the ui (#44)

* Using auto-discovery to build context options in the ui

* Fix for pypackage

beaker_kernel/contexts/dataset/default_payload.json

@@ -0,0 +1,4 @@
+{
+  "df_hosp": "truth-incident-hospitalization",
+  "df_cases": "truth-incident-case"
+}

beaker_kernel/contexts/decapodes/default_payload.json

@@ -0,0 +1,3 @@
+{
+  "halfar": "ice_dynamics-id"
+}

beaker_kernel/contexts/mira_model/default_payload.json

@@ -0,0 +1,3 @@
+{
+  "id": "sir-model-id"
+}

beaker_kernel/contexts/oceananigans/default_payload.json

@@ -0,0 +1 @@
+{}

beaker_kernel/contexts/pypackage/context.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Dict
+from typing import TYPE_CHECKING, Any, Dict, List
 
 from beaker_kernel.lib.context import BaseContext
 from beaker_kernel.lib.subkernels.python import PythonSubkernel
@@ -18,3 +18,7 @@ def __init__(self, beaker_kernel: "LLMKernel", subkernel: "BaseSubkernel", confi
         if not isinstance(subkernel, PythonSubkernel):
             raise ValueError("This context is only valid for Python.")
         super().__init__(beaker_kernel, subkernel, self.agent_cls, config)
+
+    @classmethod
+    def available_subkernels(cls) -> List["BaseSubkernel"]:
+        return ["python3"]

beaker_kernel/contexts/pypackage/default_payload.json

@@ -0,0 +1 @@
+{}

beaker_kernel/lib/context.py

@@ -2,7 +2,7 @@
 import json
 import logging
 import os.path
-from typing import TYPE_CHECKING, Any, Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 from beaker_kernel.lib.autodiscovery import autodiscover
 
@@ -60,10 +60,32 @@ def __init__(self, beaker_kernel: "LLMKernel", subkernel: "BaseSubkernel", agent
                 template = self.jinja_env.get_template(template_file)
                 self.templates[template_name] = template
 
-    async def setup(self, parent_header=None):
+    async def setup(self, config=None, parent_header=None):
+        if config:
+            self.config = config
         if callable(getattr(self.agent, 'setup', None)):
             await self.agent.setup(self.config, parent_header=parent_header)
 
+    @classmethod
+    def available_subkernels(cls) -> List["BaseSubkernel"]:
+        class_dir = inspect.getfile(cls)
+        proc_dir = os.path.join(os.path.dirname(class_dir), "procedures")
+        if os.path.exists(proc_dir):
+            subkernel_slugs = list(os.listdir(proc_dir))
+            return subkernel_slugs
+        else:
+            return []
+
+    @classmethod
+    def default_payload(cls) -> str:
+        class_dir = inspect.getfile(cls)
+        payload_file_path = os.path.join(os.path.dirname(class_dir), "default_payload.json")
+        if os.path.exists(payload_file_path):
+            with open(payload_file_path) as payload_file:
+                return payload_file.read().strip()
+        else:
+            return "{}"
+
     def send_response(self, *args, **kwargs):
         return self.beaker_kernel.send_response(*args, **kwargs)
 

service/main.py

@@ -1,5 +1,5 @@
 import os
-from typing import Any
+from typing import Dict
 
 from jupyter_server.base.handlers import JupyterHandler
 from jupyter_server.extension.handler import (
@@ -11,6 +11,10 @@
 from tornado import httputil
 from tornado.web import Application
 
+from beaker_kernel.lib.autodiscovery import autodiscover
+from beaker_kernel.lib.context import BaseContext
+from beaker_kernel.lib.subkernels.base import BaseSubkernel
+
 
 HERE = os.path.dirname(__file__)
 
@@ -29,50 +33,22 @@ class ContextHandler(ExtensionHandlerMixin, JupyterHandler):
     def get(self):
         """Get the main page for the application's interface."""
 
-        # TODO: Generate this from a single source of truth at startup time
+        contexts: Dict[str, BaseContext] = autodiscover("contexts")
+        subkernels: Dict[str, BaseSubkernel] = autodiscover("subkernels")
+
+        # Extract data from auto-discovered contexts and subkernels to provide options
         context_data = {
-            "dataset": {
+            context_slug: {
                 "languages": [
-                    ["python3", "python3"],
-                    ["julia", "julia-1.9"],
-                    ["rlang", "ir"]
+                    [subkernel_slug, getattr(subkernels.get(subkernel_slug), "KERNEL_NAME")]
+                    for subkernel_slug in context.available_subkernels()
+                    if subkernel_slug in subkernels
                 ],
-                "defaultPayload": '''
-{
-  "df_hosp": "truth-incident-hospitalization",
-  "df_cases": "truth-incident-case"
-}
-'''.strip(),
-            },
-            "mira_model": {
-                "languages": [
-                    ["python3", "python3"]
-                ],
-                "defaultPayload": '''
-{
-  "id": "sir-model-id"
-}
-'''.strip(),
-            },
-            "decapodes": {
-                "languages": [
-                    ["julia", "julia-1.9"]
-                ],
-                "defaultPayload": '''
-{
-  "halfar": "ice_dynamics-id"
-}
-'''.strip(),
-            },
-            "oceananigans": {
-                "languages": [
-                    ["julia", "julia-1.9"]
-                ],
-                "defaultPayload": '''
-{}
-'''.strip(),
-            },
+                "defaultPayload": context.default_payload()
+            }
+            for context_slug, context in contexts.items()
         }
+
         return self.write(context_data)