Add Oceananigans Mode (#39)

* Use git submodule

* Init Oceananigans

* Remove submodule

* Create save data procedure

* Register Oceananigans context

* Add to UI

* Finish fixes for saving

* Enable LLM code generation

* Tweak LLM integration

* Add dependencies in Dockerfile

* Fix to work with refactor AND `askem-julia`

* Use intercept decorator

.gitignore

@@ -63,3 +63,4 @@ create.sql
 .eslintcache
 **/.venv/
 .openai.toml
+shell.nix

Dockerfile

@@ -20,7 +20,8 @@ WORKDIR /home/jupyter
 # Install r-lang and kernel
 RUN apt update && \
     apt install -y r-base r-cran-irkernel \
-        graphviz libgraphviz-dev && \
+        graphviz libgraphviz-dev \
+        libevent-core-2.1-7 libevent-pthreads-2.1-7 && \
     apt clean -y && \
     apt autoclean -y
 

beaker_kernel/contexts/oceananigans/agent.py

@@ -0,0 +1,81 @@
+import json
+import logging
+import re
+
+from archytas.tool_utils import AgentRef, LoopControllerRef, tool, toolset
+
+from beaker_kernel.lib.agent import BaseAgent
+from beaker_kernel.lib.context import BaseContext
+
+logging.disable(logging.WARNING)  # Disable warnings
+logger = logging.Logger(__name__)
+
+
+@toolset()
+class OceananigansToolset:
+    """
+    Toolset used for working with the Julia language and the package Oceananigans
+    """
+
+    @tool()
+    async def generate_code(
+        self, query: str, agent: AgentRef, loop: LoopControllerRef
+    ) -> None:
+        """
+        Generated Julia code to be run in an interactive Jupyter notebook for the purpose of creating models, visualizations, and simulations in Oceananigans.
+
+        Input is a full grammatically correct question about or request for an action to be performed with the Julia language or Oceananigans.
+
+        Args:
+            query (str): A fully grammatically correct queistion about the current model.
+        """
+        prompt = f"""
+You are a programmer writing code to help with writing simulations in Julia and Oceananigans.jl, a fast, friendly, flexible software package for finite volume simulations of the nonhydrostatic and hydrostatic Boussinesq equations on CPUs and GPUs.
+
+As an Oceananigans example, you can run a two-dimensional, horizontally-periodic simulation of turbulence using 128² finite volume cells for 4 non-dimensional time units:
+```
+using Oceananigans
+grid = RectilinearGrid(CPU(), size=(128, 128), x=(0, 2π), y=(0, 2π), topology=(Periodic, Periodic, Flat))
+model = NonhydrostaticModel(; grid, advection=WENO())
+ϵ(x, y) = 2rand() - 1
+set!(model, u=ϵ, v=ϵ)
+simulation = Simulation(model; Δt=0.01, stop_time=4)
+run!(simulation)
+```
+
+If you have other knowledge of Oceananigans.jl, please use that as well.
+
+When doing visualization tasks, please use `WGLMakie`.
+
+Note that `using Oceananigans, WGLMakie` has already been run so do not include it in your output.
+        
+Here are the currently active Oceananigans-related variables in state, please don't redefine these since they are in state:
+```
+{await agent.context.get_available_vars()}
+```
+
+Please write code that satisfies the user's request below.
+
+Please generate the code as if you were programming inside a Jupyter Notebook and the code is to be executed inside a cell.
+You MUST wrap the code with a line containing three backticks (```) before and after the generated code.
+No addtional text is needed in the response, just the code block.
+"""
+
+        llm_response = await agent.oneshot(prompt=prompt, query=query)
+        loop.set_state(loop.STOP_SUCCESS)
+        preamble, code, coda = re.split("```\w*", llm_response)
+        result = json.dumps(
+            {
+                "action": "code_cell",
+                "language": "julia-1.9",
+                "content": code.strip(),
+            }
+        )
+        return result
+
+
+class OceananigansAgent(BaseAgent):
+
+    def __init__(self, context: BaseContext = None, tools: list = None, **kwargs):
+        tools = [OceananigansToolset]
+        super().__init__(context, tools, **kwargs)

beaker_kernel/contexts/oceananigans/context.py

@@ -0,0 +1,83 @@
+import json
+import logging
+import os
+from typing import TYPE_CHECKING, Any, Dict
+
+import requests
+
+from beaker_kernel.lib.context import BaseContext
+from beaker_kernel.lib.utils import intercept
+
+from .agent import OceananigansAgent
+
+if TYPE_CHECKING:
+    from beaker_kernel.kernel import LLMKernel
+    from beaker_kernel.lib.subkernels.base import BaseSubkernel
+
+
+logger = logging.getLogger(__name__)
+
+
+class OceananigansContext(BaseContext):
+
+    slug = "oceananigans"
+    agent_cls = OceananigansAgent
+
+    def __init__(self, beaker_kernel: "LLMKernel", subkernel: "BaseSubkernel", config: Dict[str, Any]) -> None:
+        self.target = "oceananigan"
+        self.reset()
+        super().__init__(beaker_kernel, subkernel, self.agent_cls, config)
+
+
+    async def setup(self, config, parent_header):
+        await self.execute(self.get_code("setup"))
+        print("Oceananigans creation environment set up")
+
+
+    async def post_execute(self, message):
+        pass
+
+    def reset(self):
+        pass
+
+    async def auto_context(self):
+        return f"""You are a scientific modeler whose goal is to help the user with Julia and Oceananigans.jl
+
+Here are the currently active Oceananigans-related variables in state:
+```
+{await self.get_available_vars()}
+```
+Please answer any user queries to the best of your ability, but do not guess if you are not sure of an answer.
+If you are asked to write code, please use the generate_code tool. Please use existing variables if you can and try
+not to redefine existing variables or redo a task the user has already done (unless it needs correcting).
+"""
+
+    async def get_available_vars(self, parent_header={}):
+        code = self.get_code("var_info")
+        var_info_response = await self.beaker_kernel.evaluate(
+            code,
+            parent_header=parent_header,
+        )
+        return  var_info_response.get('return')
+
+
+
+    @intercept(msg_type="save_data_request") 
+    async def save_data(self, message):
+        content = message.content
+
+        result = await self.evaluate(
+            self.get_code("save_data", 
+                {
+                    "dataservice_url": os.environ["DATA_SERVICE_URL"], 
+                    "name": content.get("name"),
+                    "description": content.get("description", ""),
+                    "filenames": content.get("filenames"),
+                }
+            ),
+        )
+
+        self.beaker_kernel.send_response(
+            "iopub", "save_data_response", result["return"], parent_header=message.header
+        )
+

beaker_kernel/contexts/oceananigans/procedures/julia/save_data.jl

@@ -0,0 +1,31 @@
+import HTTP, JSON3, DisplayAs
+
+_DATA_SERVICE_URL = "{{ dataservice_url }}" 
+_DATASET_NAME = "{{ name }}"
+_DATASET_DESCRIPTION = "{{ description }}"
+_FILENAMES = split("{{ filenames }}", ",")
+
+
+_dataset = Dict(
+    "name" => _DATASET_NAME,
+    "description" => _DATASET_DESCRIPTION,
+    "file_names" => _FILENAMES
+)
+
+_create_req = HTTP.post("$_DATA_SERVICE_URL/datasets", body=JSON3.write(_dataset))
+_new_dataset_id = JSON3.read(String(_create_req.body))["id"]
+
+for filename in _FILENAMES 
+    _new_dataset_url = "$_DATA_SERVICE_URL/datasets/$_new_dataset_id"
+    _data_url_req = HTTP.get("$(_new_dataset_url)/upload-url?filename=$filename")
+    _data_url = get!(JSON3.read(String(_data_url_req.body)), "url", nothing)
+    _filesize = stat(filename).size
+    _upload_response = HTTP.put(_data_url, ["content-length" => _filesize], open(filename, "r"))
+    if _upload_response.status != 200
+        error("Error uploading dataframe: $(String(_upload_response.body))")
+    end
+
+end
+
+_result = Dict("dataset_id" => _new_dataset_id)
+JSON3.write(_result) |> DisplayAs.unlimited

beaker_kernel/contexts/oceananigans/procedures/julia/setup.jl

@@ -0,0 +1 @@
+using Oceananigans, WGLMakie

beaker_kernel/contexts/oceananigans/procedures/julia/var_info.jl

@@ -0,0 +1,24 @@
+using Oceananigans
+_result = Any[]
+_var_syms = names(Main)
+
+_oceananigans_types = [
+    Oceananigans.AbstractGrid,
+    Oceananigans.Simulation,
+    Oceananigans.AbstractModel,
+    Oceananigans.TurbulenceClosures.AbstractTurbulenceClosure,
+    Oceananigans.Grids.AbstractTopology,
+]
+
+
+for _var_sym in _var_syms
+    _var = eval(_var_sym)
+    if any(_type -> isa(_var, _type), _oceananigans_types)
+        push!(_result, Dict(
+            "var_name" => _var_sym,
+            "value" => string(_var),
+        ))
+    end
+end
+
+JSON3.write(_result) |> DisplayAs.unlimited

service/main.py

@@ -64,7 +64,14 @@ def get(self):
 }
 '''.strip(),
             },
-
+            "oceananigans": {
+                "languages": [
+                    ["julia", "julia-1.9"]
+                ],
+                "defaultPayload": '''
+{}
+'''.strip(),
+            },
         }
         return self.write(context_data)