docs/bioimage-chatbot-extension-tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## BioImage.IO Chatbot Extensions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Documentation: https://github.com/bioimage-io/bioimageio-chatbot/blob/main/docs/development.md"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using the chatbot in an ImJoy plugin\n",
    "\n",
    "Below you will find an example of how to use the chatbot in an ImJoy plugin. Using `api.createWindow`, we can display the chatbot inline as an imjoy plugin and interact with its api."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "trusted": true
   },
   "outputs": [
    {
     "data": {
      "application/javascript": "window.connectPlugin && window.connectPlugin(\"b78046eb-f9bd-432c-a1ea-5776c21ba3be\")",
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div id=\"fe5b0500-59ef-459d-8b2a-f01e2cd520c6\"></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<_GatheringFuture pending>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from imjoy import api\n",
    "\n",
    "async def setup():\n",
    "    chatbot = await api.createWindow(\n",
    "        src=\"https://bioimage.io/chat\",\n",
    "        name=\"BioImage.IO Chatbot\",\n",
    "    )\n",
    "    \n",
    "    await api.showMessage(str(await chatbot.getAllExtensions()))\n",
    "\n",
    "api.export({\"setup\": setup})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a chatbot extension\n",
    "\n",
    "A chatbot extension object is a dictionary with the following keys:\n",
    " - `id`: a unique identifier for the extension\n",
    " - `name`: the name of the extension\n",
    " - `description`: a short description of the extension\n",
    " - `type`: it must be `bioimageio-chatbot-extension`\n",
    " - `tools`: a dictionary with functions of tools, it represents the set of functions your extension offers, each accepting configuration parameters as input. These functions should carry out specific tasks and return their results in a dictionary.\n",
    " - `get_schema`: a function returns the schema for the tools, it returns a JSON schema for each tool function, specifying the structure and types of the expected parameters. This schema is crucial for instructing the chatbot to generate the correct input paramters and validate the inputs and ensuring they adhere to the expected format. Importantly, the chatbot uses the title and description for each field to understand what expected for the tool will generating a function call to run the tool (also see the detailed instructions below).\n",
    "\n",
    "\n",
    " The following is an example of creating a chatbot extension:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "from skimage import io, filters, data\n",
    "\n",
    "\n",
    "class ApplyFilterInput(BaseModel):\n",
    "    \"\"\"Apply a gaussian filter to an example image\"\"\"\n",
    "    sigma: float = Field(..., description=\"Standard deviation for Gaussian kernel\")\n",
    "\n",
    "async def apply_filter(kwargs):\n",
    "    config = ApplyFilterInput(**kwargs)\n",
    "\n",
    "    # Load the image\n",
    "    image = data.cell()\n",
    "\n",
    "    # Apply the filter\n",
    "    filtered_image = filters.gaussian(image, sigma=config.sigma)\n",
    "\n",
    "    # Display the image\n",
    "    viewer = await api.showDialog(src=\"https://kitware.github.io/itk-vtk-viewer/app/\")\n",
    "    await viewer.setImage(filtered_image)\n",
    "    \n",
    "    return \"Filter applied successfully\"\n",
    "\n",
    "def get_schema():\n",
    "    return {\n",
    "        \"apply_filter\": ApplyFilterInput.schema(),\n",
    "    }\n",
    "\n",
    "# Define a chatbot extension\n",
    "image_processing_extension = {\n",
    "    \"_rintf\": True,\n",
    "    \"id\": \"image-processing\",\n",
    "    \"type\": \"bioimageio-chatbot-extension\",\n",
    "    \"name\": \"Image Processing\",\n",
    "    \"description\": \"Apply gaussian filter to an example image\",\n",
    "    \"get_schema\": get_schema,\n",
    "    \"tools\": {\n",
    "        \"apply_filter\": apply_filter,\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Option 1: Register a chatbot extension as an ImJoy plugin\n",
    "\n",
    "The following will show the chatbot window, then you can type for example `run the gaussian filter example` to invoke the tool you defined in the extension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from imjoy_rpc import api\n",
    "\n",
    "async def setup():\n",
    "    chatbot = await api.createWindow(src=\"https://bioimage.io/chat\", name=\"BioImage.IO Chatbot\")\n",
    "    await chatbot.registerExtension(image_processing_extension)\n",
    "\n",
    "api.export({\"setup\": setup})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Option 2: Serving extension remotely using hypha\n",
    "\n",
    "Besides running the extension in the browser using ImJoy, you can also run the extension remotely using [hypha](https://ha.amun.ai). This allows the extension to run in a native Python environment which have easier access to hardware devices (e.g. the actual microscope) and more computational resources.\n",
    "\n",
    "Similar to the ImJoy plugin, you need to register the extension as a hypha service."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import micropip\n",
    "await micropip.install(\"hypha-rpc\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hypha_rpc import connect_to_server, login\n",
    "\n",
    "server_url = \"https://chat.bioimage.io\"\n",
    "token = await login({\"server_url\": server_url})\n",
    "server = await connect_to_server({\"server_url\": server_url, \"token\": token})\n",
    "# Below, we set the visibility to public\n",
    "image_processing_extension['config'] = {\"visibility\": \"public\"}\n",
    "svc = await server.register_service(image_processing_extension)\n",
    "print(f\"Extension service registered with id: {svc.id}, you can visit the service at: https://bioimage.io/chat?server={server_url}&extension={svc.id}&assistant=Skyler\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use the chatbot, you can now click the link above, or use the following code to run the chatbot in the browser."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Display the chatbot as embedded iframe\n",
    "from IPython.display import display, IFrame\n",
    "display(IFrame(src=\"https://bioimage.io/chat\", width=\"100%\", height=\"600px\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Additional Examples\n",
    "### Example 1: Controlling a microscope using the chatbot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "    \n",
    "class MoveStageInput(BaseModel):\n",
    "    \"\"\"Move the microscope stage\"\"\"\n",
    "    x: float = Field(..., description=\"x offset\")\n",
    "    y: float = Field(..., description=\"y offset\")\n",
    "\n",
    "class SnapImageInput(BaseModel):\n",
    "    \"\"\"Move the microscope stage\"\"\"\n",
    "    exposure: float = Field(..., description=\"exposure time\")\n",
    "\n",
    "async def move_stage(kwargs):\n",
    "    config = MoveStageInput(**kwargs)\n",
    "    print(config.x, config.y)\n",
    "\n",
    "    return \"success\"\n",
    "\n",
    "async def snap_image(kwargs):\n",
    "    config = SnapImageInput(**kwargs)\n",
    "    print(config.exposure)\n",
    "    await api.showDialog(src=\"https://bioimage.io\")\n",
    "    return \"success\"\n",
    "\n",
    "def get_schema():\n",
    "        return {\n",
    "            \"move_stage\": MoveStageInput.schema(),\n",
    "            \"snap_image\": SnapImageInput.schema()\n",
    "        }\n",
    "\n",
    "# Define an chatbot extension\n",
    "microscope_control_extension = {\n",
    "    \"_rintf\": True,\n",
    "    \"id\": \"microscope-control\",\n",
    "    \"type\": \"bioimageio-chatbot-extension\",\n",
    "    \"name\": \"Microscope Control\",\n",
    "    \"description\": \"Contorl the microscope....\",\n",
    "    \"get_schema\": get_schema,\n",
    "    \"tools\": {\n",
    "        \"move_stage\": move_stage,\n",
    "        \"snap_image\": snap_image,\n",
    "    }\n",
    "}\n",
    "\n",
    "from hypha_rpc import connect_to_server, login\n",
    "\n",
    "server_url = \"https://chat.bioimage.io\"\n",
    "token = await login({\"server_url\": server_url})\n",
    "server = await connect_to_server({\"server_url\": server_url, \"token\": token})\n",
    "# Below, we set the visibility to public\n",
    "microscope_control_extension['config'] = {\"visibility\": \"public\"}\n",
    "svc = await server.register_service(microscope_control_extension)\n",
    "print(f\"Extension service registered with id: {svc.id}, you can visit the service at: https://bioimage.io/chat?server={server_url}&extension={svc.id}&assistant=Skyler\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 2: Query the models in the model zoo via code generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "trusted": true
   },
   "outputs": [
    {
     "data": {
      "application/javascript": "window.connectPlugin && window.connectPlugin(\"b78046eb-f9bd-432c-a1ea-5776c21ba3be\")",
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div id=\"450f29b8-375c-4316-9f0e-b1ff2972ba5f\"></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<_GatheringFuture pending>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from imjoy_rpc import api\n",
    "import sys\n",
    "import io\n",
    "from imjoy import api\n",
    "from js import fetch\n",
    "from pydantic import BaseModel, Field\n",
    "from typing import Optional\n",
    "from typing import List, Optional, Dict, Any\n",
    "\n",
    "class ResourceType(str):\n",
    "    MODEL = \"model\"\n",
    "    DATASET = \"dataset\"\n",
    "    APPLICATION = \"application\"\n",
    "\n",
    "def normalize_text(text: str) -> str:\n",
    "    return text.replace('_', ' ').lower()\n",
    "\n",
    "def matches_keywords(text: str, keywords: List[str]) -> bool:\n",
    "    normalized_text = normalize_text(text)\n",
    "    return any(keyword in normalized_text for keyword in keywords)\n",
    "\n",
    "def search_item(item: Dict[str, Any], keywords: List[str]) -> bool:\n",
    "    search_fields = [item.get('id', ''), item.get('nickname', ''), item.get('name', ''),\n",
    "        item.get('nickname_icon', ''), item.get('license', ''), item.get('description', '')\n",
    "    ] + [tag for tag in item.get('tags', [])]\n",
    "    search_fields += [author['name'] for author in item.get('authors', [])]\n",
    "    return any(matches_keywords(field, keywords) for field in search_fields)\n",
    "\n",
    "def search(keywords, type, top_k, resource_items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:\n",
    "    keywords = [normalize_text(keyword) for keyword in keywords]\n",
    "    filtered_items = []\n",
    "    for item in resource_items:\n",
    "        if type and item.get('type') != type:\n",
    "            continue\n",
    "        if search_item(item, keywords):\n",
    "            filtered_items.append(item)\n",
    "        if len(filtered_items) == top_k:\n",
    "            break\n",
    "    return filtered_items\n",
    "\n",
    "async def load_model_info():\n",
    "    response = await fetch(\"https://bioimage-io.github.io/collection-bioimage-io/collection.json\")\n",
    "    model_info = await response.json()\n",
    "    model_info = model_info.to_py()\n",
    "    resource_items = model_info['collection']\n",
    "    return resource_items\n",
    "\n",
    "def execute_code(script, context=None):\n",
    "    if context is None:\n",
    "        context = {}\n",
    "\n",
    "    # Redirect stdout and stderr to capture their output\n",
    "    original_stdout = sys.stdout\n",
    "    original_stderr = sys.stderr\n",
    "    sys.stdout = io.StringIO()\n",
    "    sys.stderr = io.StringIO()\n",
    "\n",
    "    try:\n",
    "        # Create a copy of the context to avoid modifying the original\n",
    "        local_vars = context.copy()\n",
    "\n",
    "        # Execute the provided Python script with access to context variables\n",
    "        exec(script, local_vars)\n",
    "\n",
    "        # Capture the output from stdout and stderr\n",
    "        stdout_output = sys.stdout.getvalue()\n",
    "        stderr_output = sys.stderr.getvalue()\n",
    "\n",
    "        return {\n",
    "            \"stdout\": stdout_output,\n",
    "            \"stderr\": stderr_output,\n",
    "            # \"context\": local_vars  # Include context variables in the result\n",
    "        }\n",
    "    except Exception as e:\n",
    "        return {\n",
    "            \"stdout\": \"\",\n",
    "            \"stderr\": str(e),\n",
    "            # \"context\": context  # Include context variables in the result even if an error occurs\n",
    "        }\n",
    "    finally:\n",
    "        # Restore the original stdout and stderr\n",
    "        sys.stdout = original_stdout\n",
    "        sys.stderr = original_stderr\n",
    "\n",
    "async def register_chatbot_extension(register):\n",
    "    resource_items = await load_model_info()\n",
    "    types = set()\n",
    "    tags = set()\n",
    "    for resource in resource_items:\n",
    "        types.add(resource['type'])\n",
    "        tags.update(resource['tags'])\n",
    "    types = list(types)\n",
    "    tags = list(tags)[:5]\n",
    "    resource_item_stats = f\"\"\"- keys: {list(resource_items[0].keys())}\\n- resource types: {types}\\n- Exampletags: {tags}\\n\"\"\" #Here is an example: {resource_items[0]}\n",
    "\n",
    "    class ModelZooInfoScript(BaseModel):\n",
    "        script: str = Field(..., description=\"\"\"Executable python script (Python runtime: Pyodide) for querying information\"\"\")\n",
    "    \n",
    "    ModelZooInfoScript.__doc__ = (\n",
    "        \"Search the BioImage Model Zoo for statistical information by executing Python3 scripts on the resource items.\"\n",
    "        \"For exampling counting models, applications, and datasets filtered by tags in the BioImage Model Zoo (bioimage.io). \"\n",
    "        \"The generated scripts will be executed browser pyodide environment, the script can access data through the 'resources' local variable, containing zoo resources as dictionaries. \"\n",
    "        \"Handle any missing fields in zoo items, and ensure outputs are directed to stdout. \"\n",
    "        \"Filter resources by the 'type' key without making remote server requests. 'resources' variable details:\\\\n\"\n",
    "    ) + resource_item_stats\n",
    "\n",
    "\n",
    "    class ModelZooSearchInput(BaseModel):\n",
    "        \"\"\"Search the BioImage Model Zoo (bioimage.io) resource items such as models, applications, datasets, etc. in the model zoo and return detailed information. The link format to the models etc. is: https://bioimage.io/#/?id=[ResourceID]\"\"\"\n",
    "        keywords: List[str] = Field(..., description=\"List of keywords to search for in the model zoo.\")\n",
    "        top_k: int = Field(3, description=\"The maximum number of search results to return. Default is 3. Please be aware each item may contain a large amount of data.\")\n",
    "        type: Optional[ResourceType] = Field(None, description=\"The type of resource to search for. Options include 'model', 'dataset', 'application'.\")\n",
    "\n",
    "\n",
    "    def get_schema():\n",
    "        return {\n",
    "            \"run_script\": ModelZooInfoScript.schema(),\n",
    "            \"search\": ModelZooSearchInput.schema()\n",
    "        }\n",
    "\n",
    "    async def execute_script(kwargs):\n",
    "        info_script = ModelZooInfoScript.parse_obj(kwargs)\n",
    "        result = execute_code(info_script.script, {\"resources\": resource_items})\n",
    "        return result\n",
    "\n",
    "    async def execute_search(kwargs):\n",
    "        config = ModelZooSearchInput.parse_obj(kwargs)\n",
    "        result = search(config.keywords, config.type, config.top_k, resource_items)\n",
    "        return result\n",
    "\n",
    "    await register({\n",
    "        \"_rintf\": True,\n",
    "        \"id\": \"bioimage_model_zoo\",\n",
    "        \"type\": \"bioimageio-chatbot-extension\",\n",
    "        \"name\": \"BioImage Model Zoo\",\n",
    "        \"description\": \"Getting information about models, applications, datasets, etc. in the BioImage Model Zoo. It takes a list of keywords or a python script to query the resources in the BioImage Model Zoo.\",\n",
    "        \"get_schema\": get_schema,\n",
    "        \"tools\": {\n",
    "            \"run_script\": execute_script,\n",
    "            \"search\": execute_search,\n",
    "        }\n",
    "    })\n",
    "\n",
    "\n",
    "async def setup():\n",
    "    chatbot = await api.createWindow(src=\"https://bioimage.io/chat\")\n",
    "    await register_chatbot_extension(chatbot.registerExtension)\n",
    "\n",
    "api.export({\"setup\": setup})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (Pyodide)",
   "language": "python",
   "name": "python"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "python",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}