From 82a32bfa19f040b2804544183299a683d92d8cf4 Mon Sep 17 00:00:00 2001 From: Brandon Rose Date: Wed, 24 Jan 2024 13:00:46 -0600 Subject: [PATCH 1/3] initial mira model configuration editing context --- .../contexts/mira_config_edit/__init__.py | 0 .../contexts/mira_config_edit/agent.py | 90 +++++++++++++++++++ .../contexts/mira_config_edit/context.py | 89 ++++++++++++++++++ .../mira_config_edit/default_payload.json | 3 + .../procedures/python3/get_initials.py | 1 + .../procedures/python3/get_params.py | 15 ++++ .../procedures/python3/load_model.py | 4 + .../procedures/python3/metadata.json | 1 + .../procedures/python3/model_preview.py | 13 +++ .../procedures/python3/setup.py | 4 + .../procedures/python3/update_params.py | 3 + 11 files changed, 223 insertions(+) create mode 100644 src/askem_beaker/contexts/mira_config_edit/__init__.py create mode 100644 src/askem_beaker/contexts/mira_config_edit/agent.py create mode 100644 src/askem_beaker/contexts/mira_config_edit/context.py create mode 100644 src/askem_beaker/contexts/mira_config_edit/default_payload.json create mode 100644 src/askem_beaker/contexts/mira_config_edit/procedures/python3/get_initials.py create mode 100644 src/askem_beaker/contexts/mira_config_edit/procedures/python3/get_params.py create mode 100644 src/askem_beaker/contexts/mira_config_edit/procedures/python3/load_model.py create mode 100644 src/askem_beaker/contexts/mira_config_edit/procedures/python3/metadata.json create mode 100644 src/askem_beaker/contexts/mira_config_edit/procedures/python3/model_preview.py create mode 100644 src/askem_beaker/contexts/mira_config_edit/procedures/python3/setup.py create mode 100644 src/askem_beaker/contexts/mira_config_edit/procedures/python3/update_params.py diff --git a/src/askem_beaker/contexts/mira_config_edit/__init__.py b/src/askem_beaker/contexts/mira_config_edit/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/askem_beaker/contexts/mira_config_edit/agent.py b/src/askem_beaker/contexts/mira_config_edit/agent.py new file mode 100644 index 0000000..eb18452 --- /dev/null +++ b/src/askem_beaker/contexts/mira_config_edit/agent.py @@ -0,0 +1,90 @@ +import json +import logging +import re + +import requests +from archytas.react import Undefined +from archytas.tool_utils import AgentRef, LoopControllerRef, tool + +from beaker_kernel.lib.agent import BaseAgent +from beaker_kernel.lib.context import BaseContext +from beaker_kernel.lib.jupyter_kernel_proxy import JupyterMessage + +logging.disable(logging.WARNING) # Disable warnings +logger = logging.Logger(__name__) + + +class MiraConfigEditAgent(BaseAgent): + """ + LLM agent used for working with the Mira Modeling framework ("mira_model" package) in Python 3. + This will be used to find pre-written functions which will be used to edit a model. + + A mira model is made up of multiple templates that are merged together like ... + + An example mira model will look like this when encoded in json: + ``` + { + "id": "foo", + "bar": "foobar", + ... + } + + Instead of manipulating the model directly, the agent will always return code that will be run externally in a jupyter notebook. + + """ + + def __init__(self, context: BaseContext = None, tools: list = None, **kwargs): + super().__init__(context, tools, **kwargs) + + @tool() + async def get_parameters_initials(self, _type: str, agent: AgentRef, loop: LoopControllerRef): + """ + This tool is used when a user wants to see the names and values of the model configuration's parameters. + 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. + + Args: + _type (str): either "parameters" or "initials" and determines whether to fetch values of the parameters or the initial conditions + """ + loop.set_state(loop.STOP_SUCCESS) + if _type == "parameters": + code = agent.context.get_code("get_params") + elif _type == "initials": + code = agent.context.get_code("get_initials") + return json.dumps( + { + "action": "code_cell", + "language": "python3", + "content": code.strip(), + } + ) + + @tool() + async def update_parameters(self, parameter_values: dict, agent: AgentRef, loop: LoopControllerRef): + """ + This tool is used when a user wants to update the model configuration parameter values. + It takes in a dictionary where the key is the parameter name and the value is the new value in the form: + + ``` + {'param1': 10, + 'param_n: 2, + ...} + ``` + + 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. + + Args: + parameter_values (dict): the dictionary of parameter names and the values to update them with + """ + loop.set_state(loop.STOP_SUCCESS) + code = agent.context.get_code("update_params", {"parameter_values": parameter_values}) + return json.dumps( + { + "action": "code_cell", + "language": "python3", + "content": code.strip(), + } + ) diff --git a/src/askem_beaker/contexts/mira_config_edit/context.py b/src/askem_beaker/contexts/mira_config_edit/context.py new file mode 100644 index 0000000..a8657b7 --- /dev/null +++ b/src/askem_beaker/contexts/mira_config_edit/context.py @@ -0,0 +1,89 @@ +import copy +import datetime +import json +import logging +import os +from typing import TYPE_CHECKING, Any, Dict, Optional + +import requests + +from beaker_kernel.lib.context import BaseContext +from beaker_kernel.lib.utils import intercept + +from .agent import MiraConfigEditAgent + +if TYPE_CHECKING: + from beaker_kernel.kernel import LLMKernel + from beaker_kernel.lib.subkernels.base import BaseSubkernel + +logger = logging.getLogger(__name__) + + +class MiraConfigEditContext(BaseContext): + + agent_cls = MiraConfigEditAgent + + model_config_id: Optional[str] + model_config_json: Optional[str] + model_config_dict: Optional[dict[str, Any]] + var_name: Optional[str] = "model_config" + + def __init__(self, beaker_kernel: "LLMKernel", subkernel: "BaseSubkernel", config: Dict[str, Any]) -> None: + self.reset() + logger.error("initializing...") + super().__init__(beaker_kernel, subkernel, self.agent_cls, config) + + def reset(self): + pass + + async def post_execute(self, message): + pass + + async def setup(self, config, parent_header): + logger.error(f"performing setup...") + self.config = config + item_id = config["id"] + item_type = config.get("type", "model_config") + logger.error(f"Processing {item_type} AMR {item_id} as a MIRA model") + await self.set_model_config( + item_id, item_type, parent_header=parent_header + ) + + async def set_model_config(self, item_id, agent=None, parent_header={}): + self.config_id = item_id + meta_url = f"{os.environ['DATA_SERVICE_URL']}/model_configurations/{self.config_id}" + logger.error(f"Meta url: {meta_url}") + self.configuration = requests.get(meta_url).json() + self.amr = self.configuration.get("configuration") + self.original_amr = copy.deepcopy(self.amr) + if self.amr: + await self.load_mira() + else: + raise Exception(f"Model config '{item_id}' not found.") + await self.send_mira_preview_message(parent_header=parent_header) + + async def load_mira(self): + command = "\n".join( + [ + self.get_code("setup"), + self.get_code("load_model", { + "var_name": self.var_name, + "model": self.amr, + }), + ] + ) + print(f"Running command:\n-------\n{command}\n---------") + await self.execute(command) + + async def send_mira_preview_message( + self, server=None, target_stream=None, data=None, parent_header={} + ): + try: + + preview = await self.evaluate(self.get_code("model_preview"), {"var_name": self.var_name}) + content = preview["return"] + self.beaker_kernel.send_response( + "iopub", "model_preview", content, parent_header=parent_header + ) + except Exception as e: + raise \ No newline at end of file diff --git a/src/askem_beaker/contexts/mira_config_edit/default_payload.json b/src/askem_beaker/contexts/mira_config_edit/default_payload.json new file mode 100644 index 0000000..c2b6232 --- /dev/null +++ b/src/askem_beaker/contexts/mira_config_edit/default_payload.json @@ -0,0 +1,3 @@ +{ + "id": "sir-model-id" +} diff --git a/src/askem_beaker/contexts/mira_config_edit/procedures/python3/get_initials.py b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/get_initials.py new file mode 100644 index 0000000..26553ab --- /dev/null +++ b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/get_initials.py @@ -0,0 +1 @@ +model_config.initials \ No newline at end of file diff --git a/src/askem_beaker/contexts/mira_config_edit/procedures/python3/get_params.py b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/get_params.py new file mode 100644 index 0000000..c0afd14 --- /dev/null +++ b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/get_params.py @@ -0,0 +1,15 @@ +def get_params(model_config): + params = "" + for kk, vv in model_config.parameters.items(): + if vv.display_name: + display_name = f"({vv.display_name})" + else: + display_name = "" + if vv.units: + units = f"({vv.units})" + else: + units = "" + params += f"{kk} {vv.display_name}: {vv.value} {units}\n" + return params + +print(get_params({{ var_name|default("model_config") }})) \ No newline at end of file diff --git a/src/askem_beaker/contexts/mira_config_edit/procedures/python3/load_model.py b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/load_model.py new file mode 100644 index 0000000..04cb30e --- /dev/null +++ b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/load_model.py @@ -0,0 +1,4 @@ +import copy, requests +amr_json = dict({{ model }}) +{{ var_name|default("model_config") }} = model_from_json(amr_json) +_model_orig = copy.deepcopy({{ var_name|default("model_config") }}) diff --git a/src/askem_beaker/contexts/mira_config_edit/procedures/python3/metadata.json b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/metadata.json new file mode 100644 index 0000000..9e26dfe --- /dev/null +++ b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/metadata.json @@ -0,0 +1 @@ +{} \ No newline at end of file diff --git a/src/askem_beaker/contexts/mira_config_edit/procedures/python3/model_preview.py b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/model_preview.py new file mode 100644 index 0000000..9f45a7a --- /dev/null +++ b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/model_preview.py @@ -0,0 +1,13 @@ +from IPython.core.interactiveshell import InteractiveShell; +from IPython.core import display_functions; +from mira.modeling.amr.petrinet import template_model_to_petrinet_json + +format_dict, md_dict = InteractiveShell.instance().display_formatter.format(GraphicalModel.for_jupyter(model_config)) +result = { + "application/json": template_model_to_petrinet_json(model_config) +} +for key, value in format_dict.items(): + if "image" in key: + result[key] = value + +result diff --git a/src/askem_beaker/contexts/mira_config_edit/procedures/python3/setup.py b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/setup.py new file mode 100644 index 0000000..c05a435 --- /dev/null +++ b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/setup.py @@ -0,0 +1,4 @@ +import requests; import pandas as pd; import numpy as np; import scipy; +import json; import mira; +import sympy; import itertools; from mira.metamodel import *; from mira.modeling import Model; +from mira.sources.amr import model_from_json; from mira.modeling.viz import GraphicalModel; diff --git a/src/askem_beaker/contexts/mira_config_edit/procedures/python3/update_params.py b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/update_params.py new file mode 100644 index 0000000..d7be80a --- /dev/null +++ b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/update_params.py @@ -0,0 +1,3 @@ +parameter_values = {{ parameter_values }}['parameter_values'] +for kk, vv in parameter_values.items(): + {{ var_name|default("model_config") }}.parameters[kk].value = vv \ No newline at end of file From ab05fd5d151499994fc5936a859eab96f1575c81 Mon Sep 17 00:00:00 2001 From: Brandon Rose Date: Wed, 24 Jan 2024 13:24:42 -0600 Subject: [PATCH 2/3] add save of model configuration --- .../contexts/mira_config_edit/context.py | 32 ++++++++++++++++++- .../procedures/python3/update_params.py | 2 +- 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/src/askem_beaker/contexts/mira_config_edit/context.py b/src/askem_beaker/contexts/mira_config_edit/context.py index a8657b7..fd3f695 100644 --- a/src/askem_beaker/contexts/mira_config_edit/context.py +++ b/src/askem_beaker/contexts/mira_config_edit/context.py @@ -54,6 +54,7 @@ async def set_model_config(self, item_id, agent=None, parent_header={}): meta_url = f"{os.environ['DATA_SERVICE_URL']}/model_configurations/{self.config_id}" logger.error(f"Meta url: {meta_url}") self.configuration = requests.get(meta_url).json() + logger.error(f"Succeeded in fetching model configuration, proceeding.") self.amr = self.configuration.get("configuration") self.original_amr = copy.deepcopy(self.amr) if self.amr: @@ -86,4 +87,33 @@ async def send_mira_preview_message( "iopub", "model_preview", content, parent_header=parent_header ) except Exception as e: - raise \ No newline at end of file + raise + + @intercept() + async def save_model_config_request(self, message): + ''' + Updates the model configuration in place. + ''' + content = message.content + + new_model: dict = ( + await self.evaluate( + f"template_model_to_petrinet_json({self.var_name})" + ) + )["return"] + + model_config = self.configuration + model_config["configuration"] = new_model + + create_req = requests.put( + f"{os.environ['DATA_SERVICE_URL']}/model_configurations/{self.config_id}", json=model_config + ) + + if create_req.status_code == 200: + logger.error(f"Successfuly updated model config {self.config_id}") + response_id = create_req.json()["id"] + + content = {"model_configuration_id": response_id} + self.beaker_kernel.send_response( + "iopub", "save_model_response", content, parent_header=message.header + ) \ No newline at end of file diff --git a/src/askem_beaker/contexts/mira_config_edit/procedures/python3/update_params.py b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/update_params.py index d7be80a..2d5c7bb 100644 --- a/src/askem_beaker/contexts/mira_config_edit/procedures/python3/update_params.py +++ b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/update_params.py @@ -1,3 +1,3 @@ -parameter_values = {{ parameter_values }}['parameter_values'] +parameter_values = {{ parameter_values['parameter_values'] }} for kk, vv in parameter_values.items(): {{ var_name|default("model_config") }}.parameters[kk].value = vv \ No newline at end of file From 26abea3a6bd99ef2de27dc9016039ce949e1710c Mon Sep 17 00:00:00 2001 From: Brandon Rose Date: Wed, 24 Jan 2024 16:22:05 -0600 Subject: [PATCH 3/3] update storage/retrieval for config edits to support HMI server --- env.example | 3 +++ src/askem_beaker/contexts/mira_config_edit/context.py | 11 ++++++++--- .../mira_config_edit/procedures/python3/get_params.py | 6 +++--- 3 files changed, 14 insertions(+), 6 deletions(-) diff --git a/env.example b/env.example index 92072c6..49119fc 100644 --- a/env.example +++ b/env.example @@ -1,5 +1,8 @@ OPENAI_API_KEY={OPENAI_API_KEY} DATA_SERVICE_URL=http://localhost:8001 +HMI_SERVER_URL=http://localhost:3000 +HMI_SERVER_USER=user +HMI_SERVER_PASSWORD=pass JUPYTER_SERVER=http://localhost:8888 JUPYTER_TOKEN=ebcec7fcf42f28baccfab1cbc07bfb3f ENABLE_USER_PROMPT=true diff --git a/src/askem_beaker/contexts/mira_config_edit/context.py b/src/askem_beaker/contexts/mira_config_edit/context.py index fd3f695..aaa5113 100644 --- a/src/askem_beaker/contexts/mira_config_edit/context.py +++ b/src/askem_beaker/contexts/mira_config_edit/context.py @@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any, Dict, Optional import requests +from requests.auth import HTTPBasicAuth from beaker_kernel.lib.context import BaseContext from beaker_kernel.lib.utils import intercept @@ -51,9 +52,12 @@ async def setup(self, config, parent_header): async def set_model_config(self, item_id, agent=None, parent_header={}): self.config_id = item_id - meta_url = f"{os.environ['DATA_SERVICE_URL']}/model_configurations/{self.config_id}" + meta_url = f"{os.environ['HMI_SERVER_URL']}/model-configurations/{self.config_id}" logger.error(f"Meta url: {meta_url}") - self.configuration = requests.get(meta_url).json() + self.configuration = requests.get(meta_url, + auth=(os.environ['HMI_SERVER_USER'], + os.environ['HMI_SERVER_PASSWORD']) + ).json() logger.error(f"Succeeded in fetching model configuration, proceeding.") self.amr = self.configuration.get("configuration") self.original_amr = copy.deepcopy(self.amr) @@ -106,7 +110,8 @@ async def save_model_config_request(self, message): model_config["configuration"] = new_model create_req = requests.put( - f"{os.environ['DATA_SERVICE_URL']}/model_configurations/{self.config_id}", json=model_config + f"{os.environ['HMI_SERVER_URL']}/model-configurations/{self.config_id}", json=model_config, + auth =(os.environ['HMI_SERVER_USER'], os.environ['HMI_SERVER_PASSWORD']) ) if create_req.status_code == 200: diff --git a/src/askem_beaker/contexts/mira_config_edit/procedures/python3/get_params.py b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/get_params.py index c0afd14..1beff78 100644 --- a/src/askem_beaker/contexts/mira_config_edit/procedures/python3/get_params.py +++ b/src/askem_beaker/contexts/mira_config_edit/procedures/python3/get_params.py @@ -2,14 +2,14 @@ def get_params(model_config): params = "" for kk, vv in model_config.parameters.items(): if vv.display_name: - display_name = f"({vv.display_name})" + display_name = f" ({vv.display_name})" else: display_name = "" if vv.units: - units = f"({vv.units})" + units = f" ({vv.units})" else: units = "" - params += f"{kk} {vv.display_name}: {vv.value} {units}\n" + params += f"{kk}{display_name}: {vv.value}{units}\n" return params print(get_params({{ var_name|default("model_config") }})) \ No newline at end of file