wei_gen is a python package that takes a natural language description of an experiment and generates WEI workflows, code, and instrument configs. wei_gen is packaged as a class, and can be imported into any python project or spun up as an API. Examples of usage are in /scripts. An API-ized version of wei_gen can be found in /api.
See the demo here
Interact with weigen with a simple next.js app
The api wrapper of weigen can be found in /src/api. The following endpoints are supported.
- /session/init - returns session history on success
- /session/<session_id>/workflow_step - returns session history on success
- /session/<session_id>/code_step - returns session history on success
- /session/<session_id>/config_step - returns session history on success
- /session/<session_id>/call_gen_env/ - returns session history on success
- /session/<session_id>/history - returns session history on success
- /session/<session_id>/update_generated/ - returns session history on success
- Preparation: Provide some sort of natural language description of the experiment along with values
- Validate experiment: Provide all
/abouts from all instruments in the lab as context (≈7k tokens). Ask underlying model to check if the experiment described in the input from step 0 can be carried out with the available resources. Use token logprobs for this, over 50% "yes" prob as the the experiment will continue. - Create experiment framework: Using
/abouts as context, generate a structure for the experiment (using the input from step 0). - Validate experiment framework: Using the z3 SMT Solver, check if the logical flow of the plan makes sense, retry generation with feedback if unsatisfiable. For more about this, check the Symbol AI section below.
- Make a short list of instruments to use: Using
/abouts as context, make a short list of instruments to use as a precursor to the following step. - Generate workflow: Module
/abouts are fed as context, list of instruments from previous step provided 8000 as suggestion, retrieve 2-3 examples of workflows from vector db (RAG) and insert into context. Use the experiment framework from step 2 as guidance. - Generate code: Using previous messages as context, create python code logic for the entire experiment. Using an AgentCoder inspired scheme.
- Generate misc: Using the previous information as context, generate any extra documents (init locations of pipets, etc.).
RAG is used in some steps of generation to do in context learning. WEIGen's RAG implementation is done with chromadb.
WEIGen uses the z3 SMT Solver from Microsoft Research to validate the flow of a generated experiment. This implementation goes as follows: 0. Using the natural language input and metadata from all instruments in the laboratory, determine if the experiment can be carried out
- Generate a markdown framework/plan for the experiment based on the input and metadata from all instruments in the laboratory.
- Using LLM, parse this plan into a json indexed by action with each entry having attributes such as the instrument in use.
- The json is parsed into a z3 solver instance, specifically checking if there is an instrument handling transportation between each step that has a change in the instrument in use
- Satisfiability of the resulting z3 statment is checked, if unsatisfiable, the restart at step 1 with some feedback.
Using z3 in this manner was inspired by Large Language Models Can Plan Your Travels Rigorously with Formal Verification Tools.
The underlying model for each model can be augmented by the config (see below), the standard during testing was GPT4.
A feature of wei_gen is the ability to load and store sessions. This enables wei_gen to not only be used in an iterative manner after an experiment is run (you can use wei_gen to help modify your workflow or code given errors that arose), but also allows wei_gen to work async. This enables wei_gen to be hosted as an API, as session data is persisted across various requests as long as a proper session UID is passed by the user.
An example a of a session history is below
{
"version": "0.0.1",
"session_id": "f8954ec2-b1b4-413e-b566-ed4fe641536d",
"timestamp": 0,
"validity": 0.6,
"original_user_description": "",
"original_user_values": "",
"framework_agent_ctx": None,
"workflow_agent_ctx": None,
"code_agent_ctx": None,
"validator_agent_ctx": None,
"config_agent_ctx": None,
"generated_framework": "",
"generated_code": "",
"generated_workflow": "",
"generated_config": "",
}The idea is that this entire json document can also be sent to a front end and be parsed into a clean and friendly UI (see here).
To run any tests, check out the scripts in /scripts that have various examples of testing the different functions accessible on the wei_gen session.
The weigen session class has the following functions that can be called on it
- def execute_experiment(self, user_description: str, user_values = None) -> None:
- def framework_step(self, user_description, user_values) -> None:
- def workflow_step(self) -> None:
- def code_step(self) -> None:
- def config_step(self) -> None:
- def call_gen_env(self, agent: str, user_input: str) -> str
- def modify_generated_data(self, agent: str, user_input: str) -> None
- def get_history(self) -> Dict[str, Any]
from wei_gen import WEIGen
config_path = "<path to config>"
weigen = WEIGen(config_path)
weigen_session = weigen.new_session()api_keys:
openai:
key: "sk-..."
org: "org-..."
settings:
framework_model: "gpt-4-turbo"
code_model: "gpt-4-turbo"
validator_model: "gpt-4-turbo"
workflow_model: "gpt-4-turbo"*in future it might be easier to just have a config class that can be constructed within the code.
V1
- Prompt templates
- All tool abouts are local, I think some are enhanced
- Critic ≈ validator
- Some plugging into pryankas work
V2
- Different flow of wei_gen
- Changed structure of agents
- Code/workflow gen upgraded
- Added history
- Fine tuning a code llama for workflow gen could yield better results than gpt4.