examples/workflow.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Building Workflows and AI Agents through Multi-tool Collaboration\n",
    "\n",
    "OpenBioMed provides a user-friendly interface for connnecting different models and tools to build your own workflow and AI agent. Before starting, we recommend reading our tutorials for [tools](./explore_ai4s_tools.ipynb) and [data structure](./manipulate_molecules.ipynb) to help you clarify potential misunderstandings and build high-quality workflows.\n",
    "\n",
    "The workflow is typically configured with a yaml file composed of two parts: `tools` and `edges`. \n",
    "\n",
    "- The `tools` part is composed of a list of tools, where the user should cnofigure its name and inputs (if needed). Refer to `open_biomed/core/tool_registry.py` for our available tools and use the `print_usage()` function for configuring the required inputs of the tool. NOTE: If parts of the input takes the output of a prior tool, it will not be configured in this yaml file.\n",
    "- The `edges` part is compose of a list of connections between tools. The output of the `start` node is passed to the inputs of the `end` node. Typically, the workflow can automatically determine the part of the input that the output should flow into (e.g. if the output of a tool is a `Molecule` object, it will be passed as the `molecule` input for a downstream tool). We also provide a way to explicitly specify the input field by adding a `name_mapping` to the edge.\n",
    "\n",
    "### Demo Workflow\n",
    "Here we explain the [demo workflow](../configs/workflow/demo.yaml):\n",
    "```\n",
    "tools:\n",
    "  - name: molecule_name_request\n",
    "    inputs:\n",
    "      accession: dimethoxy-sulfanylidene-(3,5,6-trichloropyridin-2-yl)oxy-lambda5-phosphane\n",
    "  - name: molecule_question_answering\n",
    "    inputs:\n",
    "      text: How can we group this molecule according to its molecular structure and functional groups?\n",
    "edges:\n",
    "  - start: 0\n",
    "    end: 1\n",
    "```\n",
    "The workflow first obtains a molecule named 'dimethoxy-sulfanylidene-(3,5,6-trichloropyridin-2-yl)oxy-lambda5-phosphane' from PubChem (Tool No.0) and then asks a question of the molecule (Tool No.1). Tool No.0 takes an accession as input (specified in config) and generates a `Molecule` object. Tool No.1 takes a molecule (tool output) and a textual question (specified in config) as inputs and generates a textual answer. The output of `molecule_name_request` is automatically passed to the `molecule` input of `molecule_question_answering` (as the edge starts from Tool No.0 to Tool No.1).\n",
    "\n",
    "Once the configuration file is specified, you can put it under `cnofigs/workflow` and simply execute the workflow by running: \n",
    "```\n",
    "python open_biomed/core/workflow.py --workflow [YOUR_WORKFLOW_NAME]\n",
    "```\n",
    "or the following codes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/AIRvePFS/dair/luoyz-data/projects/OpenBioMed/OpenBioMed_arch\n"
     ]
    }
   ],
   "source": [
    "# Change working directory\n",
    "import os\n",
    "import sys\n",
    "parent = os.path.dirname(os.path.abspath(''))\n",
    "print(parent)\n",
    "sys.path.append(parent)\n",
    "os.chdir(parent)\n",
    "\n",
    "import logging\n",
    "logging.basicConfig(level=logging.ERROR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/AIRvePFS/dair/conda_envs/biomed/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "/AIRvePFS/dair/conda_envs/biomed/lib/python3.9/site-packages/transformers/deepspeed.py:23: FutureWarning: transformers.deepspeed module is deprecated and will be removed in a future version. Please import deepspeed modules directly from transformers.integrations\n",
      "  warnings.warn(\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Inference Steps: 100%|██████████| 1/1 [00:01<00:00,  1.70s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "([<open_biomed.data.text.Text object at 0x7fe55ea3f070>], ['N-[[(3R,9R,10S)-12-[(2S)-1-hydroxypropan-2-yl]-16-(methanesulfonamido)-3,10-dimethyl-13-oxo-2,8-dioxa-12-azabicyclo[12. 4. 0]octadeca-1(14),15,17-trien-9-yl]methyl]-N-methyl-2-phenylacetamide'])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Task pending name='Task-5' coro=<Workflow.run() running at /AIRvePFS/dair/luoyz-data/projects/OpenBioMed/OpenBioMed_arch/open_biomed/core/workflow.py:60>>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Now we have a workflow with the following tools: \n",
      "Tool No.1: PubChem query.\n",
      "Inputs: {\"accession\": molecule name}\n",
      "Outputs: A molecule from PubChem.\n",
      "\n",
      "\n",
      "Tool No.2: Molecule question answering.\n",
      "Inputs: {\"molecule\": a small molecule you are interested in. \"text\": a question about the molecule.}\n",
      "Outputs: An answer to the question.\n",
      "\n",
      "\n",
      "Repeating workflow execution No.1:\n",
      "Next we execute Tool No.1.\n",
      " The inputs of this tool are: {\"accession\": \"dimethoxy-sulfanylidene-(3,5,6-trichloropyridin-2-yl)oxy-lambda5-phosphane\"}\n",
      "./tmp/pubchem_query_pubchem_dimethoxy-sulfanylidene-(3,5,6-trichloropyridin-2-yl)oxy-lambda5-phosphane.sdf\n",
      "The outputs of this tool are: {\"molecule\": \"COP(=S)(OC)Oc1nc(Cl)c(Cl)cc1Cl\"}\n",
      "\n",
      "\n",
      "Tool No.1 passes its \"molecule\" output to Tool No.2\n",
      "Next we execute Tool No.2.\n",
      " The inputs of this tool are: {\"text\": \"How can we group this molecule according to its molecular structure and functional groups?\", \"molecule\": \"COP(=S)(OC)Oc1nc(Cl)c(Cl)cc1Cl\"}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Inference Steps: 100%|██████████| 1/1 [00:00<00:00,  3.47it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The outputs of this tool are: {\"text\": \"organic phosphonate; organothiophosphate insecticide\"}\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import asyncio\n",
    "from open_biomed.core.workflow import Workflow\n",
    "from open_biomed.utils.config import Config\n",
    "\n",
    "workflow = Workflow(config=Config(\"./configs/workflow/demo.yaml\"))\n",
    "# Workflow exectution may contain asynchronous calls\n",
    "asyncio.create_task(workflow.run(num_repeats=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Drug Design Workflow\n",
    "Here's a more [complicated workflow](../configs/workflow/stable_drug_design.yaml). It first visualizes a protein and a pocket within the protein, and generates a molecule that fits the pocket. Then, the workflow optimizes the molecule to reduce its liver toxcity. Finally, it visualizes the binding pose, reports the binding score, and estimates several properties of the molecule before and after optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Task pending name='Task-7' coro=<Workflow.run() running at /AIRvePFS/dair/luoyz-data/projects/OpenBioMed/OpenBioMed_arch/open_biomed/core/workflow.py:60>>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./tmp/protein_to_visualize.pdb [414, 415, 416, 417, 418, 292, 293, 294, 295, 424, 426, 302, 303, 304, 305, 427, 428, 429, 301, 313, 314, 315, 316, 317, 318, 329, 332, 333, 335, 336, 344, 345, 346, 347, 348, 349, 425, 359, 360, 361, 362, 363, 364, 365, 366, 367, 368, 371, 372, 375, 358]\n",
      " PyMOL not running, entering library mode (experimental)\n",
      " Warning: 'protein' may become a reserved selection keyword in the future\n",
      " Movie: frame    1 of   20, 5.47 sec. (0:01:49 - 0:01:49 to go).\n",
      " Movie: frame    2 of   20, 4.68 sec. (0:01:28 - 0:01:36 to go).\n",
      " Movie: frame    3 of   20, 4.37 sec. (0:01:18 - 0:01:27 to go).\n",
      " Movie: frame    4 of   20, 4.40 sec. (0:01:14 - 0:01:20 to go).\n",
      " Movie: frame    5 of   20, 4.25 sec. (0:01:07 - 0:01:14 to go).\n",
      " Movie: frame    6 of   20, 4.36 sec. (0:01:05 - 0:01:08 to go).\n",
      " Movie: frame    7 of   20, 4.51 sec. (0:01:03 - 0:01:04 to go).\n",
      " Movie: frame    8 of   20, 5.19 sec. (0:01:07 - 0:01:00 to go).\n",
      " Movie: frame    9 of   20, 5.18 sec. (0:01:02 - 0:00:56 to go).\n",
      " Movie: frame   10 of   20, 5.29 sec. (0:00:58 - 0:00:52 to go).\n",
      " Movie: frame   11 of   20, 5.29 sec. (0:00:52 - 0:00:48 to go).\n",
      " Movie: frame   12 of   20, 5.31 sec. (0:00:47 - 0:00:43 to go).\n",
      " Movie: frame   13 of   20, 5.24 sec. (0:00:41 - 0:00:39 to go).\n",
      " Movie: frame   14 of   20, 5.34 sec. (0:00:37 - 0:00:34 to go).\n",
      " Movie: frame   15 of   20, 5.27 sec. (0:00:31 - 0:00:29 to go).\n",
      " Movie: frame   16 of   20, 5.40 sec. (0:00:27 - 0:00:24 to go).\n",
      " Movie: frame   17 of   20, 5.34 sec. (0:00:21 - 0:00:19 to go).\n",
      " Movie: frame   18 of   20, 5.66 sec. (0:00:16 - 0:00:15 to go).\n",
      " Movie: frame   19 of   20, 5.82 sec. (0:00:11 - 0:00:10 to go).\n",
      " Movie: frame   20 of   20, 5.50 sec. (0:00:05 - 0:00:05 to go).\n",
      "['/AIRvePFS/dair/luoyz-data/projects/OpenBioMed/OpenBioMed_arch/tmp/pocket_4xli_B_pocket_2025_03_05_07_53_14_619.png']\n",
      "/AIRvePFS/dair/luoyz-data/projects/OpenBioMed/OpenBioMed_arch/tmp/visualize_pockets_in_protein_pocket_4xli_B_pocket_2025_03_05_07_53_14_619.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Sampling: 100%|██████████| 100/100 [00:04<00:00, 22.94it/s]\n",
      "Post processing molecules...: 100%|██████████| 1/1 [00:00<00:00, 100.19it/s]\n",
      "Inference Steps: 100%|██████████| 1/1 [00:04<00:00,  4.38s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./tmp/structure-based_drug_design_mol_1741161304846_0.sdf\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: At low exhaustiveness, it may be impossible to utilize all CPUs.\n",
      "Inference Steps: 100%|██████████| 1/1 [00:00<00:00,  1.91it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "./tmp/text-based_molecule_editing_mol_1741161308373_0.sdf\n",
      " PyMOL not running, entering library mode (experimental)\n",
      " Warning: 'protein' may become a reserved selection keyword in the future\n",
      " Movie: frame    1 of   20, 1.59 sec. (0:00:31 - 0:00:31 to go).\n",
      " Movie: frame    2 of   20, 1.61 sec. (0:00:30 - 0:00:30 to go).\n",
      " Movie: frame    3 of   20, 1.60 sec. (0:00:28 - 0:00:28 to go).\n",
      " Movie: frame    4 of   20, 1.55 sec. (0:00:26 - 0:00:26 to go).\n",
      " Movie: frame    5 of   20, 1.58 sec. (0:00:25 - 0:00:25 to go).\n",
      " Movie: frame    6 of   20, 1.66 sec. (0:00:24 - 0:00:23 to go).\n",
      " Movie: frame    7 of   20, 1.68 sec. (0:00:23 - 0:00:22 to go).\n",
      " Movie: frame    8 of   20, 1.71 sec. (0:00:22 - 0:00:21 to go).\n",
      " Movie: frame    9 of   20, 1.70 sec. (0:00:20 - 0:00:19 to go).\n",
      " Movie: frame   10 of   20, 1.77 sec. (0:00:19 - 0:00:18 to go).\n",
      " Movie: frame   11 of   20, 1.73 sec. (0:00:17 - 0:00:16 to go).\n",
      " Movie: frame   12 of   20, 1.73 sec. (0:00:15 - 0:00:14 to go).\n",
      " Movie: frame   13 of   20, 1.67 sec. (0:00:13 - 0:00:13 to go).\n",
      " Movie: frame   14 of   20, 1.60 sec. (0:00:11 - 0:00:11 to go).\n",
      " Movie: frame   15 of   20, 1.57 sec. (0:00:09 - 0:00:09 to go).\n",
      " Movie: frame   16 of   20, 1.62 sec. (0:00:08 - 0:00:08 to go).\n",
      " Movie: frame   17 of   20, 1.66 sec. (0:00:06 - 0:00:06 to go).\n",
      " Movie: frame   18 of   20, 1.68 sec. (0:00:05 - 0:00:04 to go).\n",
      " Movie: frame   19 of   20, 1.64 sec. (0:00:03 - 0:00:03 to go).\n",
      " Movie: frame   20 of   20, 1.63 sec. (0:00:01 - 0:00:01 to go).\n",
      "['/AIRvePFS/dair/luoyz-data/projects/OpenBioMed/OpenBioMed_arch/tmp/complex_mol_2025_03_05_07_55_08_397_4xli_B.png']\n",
      "/AIRvePFS/dair/luoyz-data/projects/OpenBioMed/OpenBioMed_arch/tmp/visualize_a_ligand-receptor_complex_complex_mol_2025_03_05_07_55_08_397_4xli_B.png\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Inference Steps: 100%|██████████| 1/1 [00:00<00:00,  1.48it/s]\n",
      "Inference Steps: 100%|██████████| 1/1 [00:00<00:00, 20.13it/s]\n",
      "Inference Steps: 100%|██████████| 1/1 [00:00<00:00, 20.20it/s]\n",
      "Inference Steps: 100%|██████████| 1/1 [00:00<00:00, 20.21it/s]\n",
      "Inference Steps: 100%|██████████| 1/1 [00:00<00:00, 220.08it/s]\n",
      "Sampling: 100%|██████████| 100/100 [00:04<00:00, 22.92it/s]\n",
      "Post processing molecules...: 100%|██████████| 1/1 [00:00<00:00, 98.91it/s]\n",
      "Sampling: 100%|██████████| 100/100 [00:04<00:00, 23.31it/s]\n",
      "Post processing molecules...: 100%|██████████| 1/1 [00:00<00:00, 98.99it/s]\n",
      "Sampling: 100%|██████████| 100/100 [00:04<00:00, 23.18it/s]\n",
      "Post processing molecules...: 100%|██████████| 1/1 [00:00<00:00, 99.78it/s]\n",
      "Sampling: 100%|██████████| 100/100 [00:04<00:00, 23.15it/s]\n",
      "Post processing molecules...: 100%|██████████| 1/1 [00:00<00:00, 100.12it/s]\n",
      "Sampling: 100%|██████████| 100/100 [00:04<00:00, 23.24it/s]\n",
      "Post processing molecules...: 100%|██████████| 1/1 [00:00<00:00, 99.68it/s]\n",
      "Sampling: 100%|██████████| 100/100 [00:04<00:00, 23.33it/s]\n",
      "Post processing molecules...: 100%|██████████| 1/1 [00:00<00:00, 99.37it/s]\n",
      "Sampling: 100%|██████████| 100/100 [00:04<00:00, 23.30it/s]\n",
      "Post processing molecules...: 100%|██████████| 1/1 [00:00<00:00, 99.43it/s]\n",
      "Sampling: 100%|██████████| 100/100 [00:04<00:00, 23.23it/s]\n",
      "Post processing molecules...: 100%|██████████| 1/1 [00:00<00:00, 100.04it/s]\n",
      "Sampling: 100%|██████████| 100/100 [00:04<00:00, 23.43it/s]\n",
      "Post processing molecules...: 100%|██████████| 1/1 [00:00<00:00, 100.14it/s]\n",
      "Sampling: 100%|██████████| 100/100 [00:04<00:00, 23.20it/s]\n",
      "Post processing molecules...: 100%|██████████| 1/1 [00:00<00:00, 99.52it/s]\n",
      "Inference Steps: 100%|██████████| 1/1 [00:43<00:00, 43.23s/it]\n",
      "Inference Steps: 100%|██████████| 1/1 [00:00<00:00,  1.11it/s]\n",
      "Inference Steps: 100%|██████████| 1/1 [00:00<00:00, 20.37it/s]\n",
      "Inference Steps: 100%|██████████| 1/1 [00:00<00:00, 20.30it/s]\n",
      "Inference Steps: 100%|██████████| 1/1 [00:00<00:00, 20.08it/s]\n",
      "Inference Steps: 100%|██████████| 1/1 [00:00<00:00, 245.37it/s]\n"
     ]
    }
   ],
   "source": [
    "import asyncio\n",
    "from open_biomed.core.workflow import Workflow\n",
    "from open_biomed.utils.config import Config\n",
    "\n",
    "workflow = Workflow(config=Config(\"./configs/workflow/stable_drug_design.yaml\"))\n",
    "# Workflow exectution may contain asynchronous calls\n",
    "# The exectution of this workflow may take a bit longer\n",
    "asyncio.create_task(workflow.run(num_repeats=1, context=open(\"./logs/workflow_outputs1.txt\", \"w\")))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Trial-and-Errors with Workflow\n",
    "\n",
    "Notably, some of the tools may generate multiple outputs, or generate different outputs in different runs. Due to our architectural design, we currently choose a random sample from the multi-outputs as the inputs for downstream tools. Fortunately, the `workflow.run()` takes a `num_repeats` argument to execute the workflow for multiple times, allowing a trail-and-error process. \n",
    "\n",
    "### Developing LLM Agent based on Workflow\n",
    "Details of the workflow execution are passed into the `context` argument of the function, which could be used by an LLM agent to summarize and obtain insights within the process. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from open_biomed.core.llm_request import ReportGeneratorSBDD\n",
    "\n",
    "# Initialize the agent\n",
    "# A simple report structure is defined in ReportGeneratorGeneral for general tasks. To obtain insights for your specific workflow, we suggest to create the corresponding structure. ReportGeneratorSBDD is an example for drug design task.\n",
    "agent = ReportGeneratorSBDD()\n",
    "\n",
    "# Pass the details of the workflow execution as the context\n",
    "context = agent._gen_context(\"./logs/workflow_outputs1.txt\")\n",
    "\n",
    "# Summarize and obtain insights\n",
    "resp = agent.llm.generate(query=\"\", context=context)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print (\"[Report]\\n\", resp['final_resp'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print (\"\\n[Thoughts]\\n\", resp['reasoning'])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.7 ('biomed')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2b5492c31ef84abdc69aadb95e4c210f44c226a5800d1d766b22f7a50017392c"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}