demo_app_with_super_agent.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!pip install squidpy scanpy --upgrade"
      ],
      "metadata": {
        "id": "54uw_8McPd8d"
      },
      "id": "54uw_8McPd8d",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install langchain langgraph"
      ],
      "metadata": {
        "id": "Ql0AhMu6UZDi"
      },
      "id": "Ql0AhMu6UZDi",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install langchain_google_vertexai"
      ],
      "metadata": {
        "id": "2rWZzgKFUUNe"
      },
      "id": "2rWZzgKFUUNe",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install synapseclient"
      ],
      "metadata": {
        "id": "Hbn04DkXDYqi"
      },
      "id": "Hbn04DkXDYqi",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install commot"
      ],
      "metadata": {
        "id": "SMNOPCDx1W98"
      },
      "id": "SMNOPCDx1W98",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")"
      ],
      "metadata": {
        "id": "eEtFcg4SfD8F"
      },
      "id": "eEtFcg4SfD8F",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import re\n",
        "import io\n",
        "import os\n",
        "import sys\n",
        "import synapseclient\n",
        "from synapseclient import Synapse\n",
        "import gcsfs\n",
        "\n",
        "import vertexai\n",
        "import requests\n",
        "from vertexai.preview import reasoning_engines\n",
        "from langchain_google_vertexai import HarmBlockThreshold, HarmCategory\n",
        "\n",
        "import scanpy as sc\n",
        "import pandas as pd\n",
        "from google.cloud import bigquery\n",
        "import pandas_gbq\n",
        "\n",
        "\n",
        "from IPython.display import HTML, display\n",
        "\n",
        "def set_css():\n",
        "  display(HTML('''\n",
        "  <style>\n",
        "    pre {\n",
        "        white-space: pre-wrap;\n",
        "    }\n",
        "  </style>\n",
        "  '''))\n",
        "get_ipython().events.register('pre_run_cell', set_css)\n",
        "\n",
        "syn = synapseclient.login(authToken=\"\", silent=True)\n",
        "\n",
        "\n",
        "os.environ['SYNAPSE_AUTH_TOKEN'] = \"\""
      ],
      "metadata": {
        "id": "ex1_NzVnZXSp"
      },
      "id": "ex1_NzVnZXSp",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Path to your GCS file\n",
        "file_path = 'gs://htan_st_datasets/HTAN_ST_metadata_ad.tsv'\n",
        "\n",
        "# Create a GCS filesystem object\n",
        "fs = gcsfs.GCSFileSystem()\n",
        "\n",
        "# Load the TSV file directly from GCS\n",
        "with fs.open(file_path) as f:\n",
        "    metadata_df = pd.read_csv(f, sep='\\t')\n",
        "\n",
        "# Display the first few rows of the dataframe\n",
        "metadata_df.head()"
      ],
      "metadata": {
        "id": "YxPLHy1carkI",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 206
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1731078327815,
          "user_tz": 300,
          "elapsed": 1015,
          "user": {
            "displayName": "",
            "userId": ""
          }
        },
        "outputId": "90f1b693-5040-446b-8801-c2eae8a6a5fc"
      },
      "id": "YxPLHy1carkI",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "                    Attribute:  \\\n",
              "0                    Filename:   \n",
              "1                      Run ID:   \n",
              "2                 File Format:   \n",
              "3  HTAN Parent Biospecimen ID:   \n",
              "4           HTAN Data File ID:   \n",
              "\n",
              "                                         Description  \n",
              "0                                     Name of a file  \n",
              "1  A unique identifier for this individual run (t...  \n",
              "2  Format of a file (e.g. txt, csv, fastq, bam, e...  \n",
              "3  HTAN Biospecimen Identifier (eg HTANx_yyy_zzz)...  \n",
              "4  Self-identifier for this data file - HTAN ID o...  "
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-b92bd6d4-5d2b-4f25-bb52-7534c173fbf2\" class=\"colab-df-container\">\n",
              "    <div>\n",
              "<style scoped>\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "</style>\n",
              "<table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              "    <tr style=\"text-align: right;\">\n",
              "      <th></th>\n",
              "      <th>Attribute:</th>\n",
              "      <th>Description</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>Filename:</td>\n",
              "      <td>Name of a file</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>Run ID:</td>\n",
              "      <td>A unique identifier for this individual run (t...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>File Format:</td>\n",
              "      <td>Format of a file (e.g. txt, csv, fastq, bam, e...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>HTAN Parent Biospecimen ID:</td>\n",
              "      <td>HTAN Biospecimen Identifier (eg HTANx_yyy_zzz)...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>HTAN Data File ID:</td>\n",
              "      <td>Self-identifier for this data file - HTAN ID o...</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-b92bd6d4-5d2b-4f25-bb52-7534c173fbf2')\"\n",
              "            title=\"Convert this dataframe to an interactive table.\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "    <script>\n",
              "      const buttonEl =\n",
              "        document.querySelector('#df-b92bd6d4-5d2b-4f25-bb52-7534c173fbf2 button.colab-df-convert');\n",
              "      buttonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "      async function convertToInteractive(key) {\n",
              "        const element = document.querySelector('#df-b92bd6d4-5d2b-4f25-bb52-7534c173fbf2');\n",
              "        const dataTable =\n",
              "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                    [key], {});\n",
              "        if (!dataTable) return;\n",
              "\n",
              "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "          + ' to learn more about interactive tables.';\n",
              "        element.innerHTML = '';\n",
              "        dataTable['output_type'] = 'display_data';\n",
              "        await google.colab.output.renderOutput(dataTable, element);\n",
              "        const docLink = document.createElement('div');\n",
              "        docLink.innerHTML = docLinkHtml;\n",
              "        element.appendChild(docLink);\n",
              "      }\n",
              "    </script>\n",
              "  </div>\n",
              "\n",
              "\n",
              "<div id=\"df-978bcb56-7593-4c60-8870-dcb5751342d1\">\n",
              "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-978bcb56-7593-4c60-8870-dcb5751342d1')\"\n",
              "            title=\"Suggest charts\"\n",
              "            style=\"display:none;\">\n",
              "\n",
              "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "     width=\"24px\">\n",
              "    <g>\n",
              "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
              "    </g>\n",
              "</svg>\n",
              "  </button>\n",
              "\n",
              "<style>\n",
              "  .colab-df-quickchart {\n",
              "      --bg-color: #E8F0FE;\n",
              "      --fill-color: #1967D2;\n",
              "      --hover-bg-color: #E2EBFA;\n",
              "      --hover-fill-color: #174EA6;\n",
              "      --disabled-fill-color: #AAA;\n",
              "      --disabled-bg-color: #DDD;\n",
              "  }\n",
              "\n",
              "  [theme=dark] .colab-df-quickchart {\n",
              "      --bg-color: #3B4455;\n",
              "      --fill-color: #D2E3FC;\n",
              "      --hover-bg-color: #434B5C;\n",
              "      --hover-fill-color: #FFFFFF;\n",
              "      --disabled-bg-color: #3B4455;\n",
              "      --disabled-fill-color: #666;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart {\n",
              "    background-color: var(--bg-color);\n",
              "    border: none;\n",
              "    border-radius: 50%;\n",
              "    cursor: pointer;\n",
              "    display: none;\n",
              "    fill: var(--fill-color);\n",
              "    height: 32px;\n",
              "    padding: 0;\n",
              "    width: 32px;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart:hover {\n",
              "    background-color: var(--hover-bg-color);\n",
              "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "    fill: var(--button-hover-fill-color);\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart-complete:disabled,\n",
              "  .colab-df-quickchart-complete:disabled:hover {\n",
              "    background-color: var(--disabled-bg-color);\n",
              "    fill: var(--disabled-fill-color);\n",
              "    box-shadow: none;\n",
              "  }\n",
              "\n",
              "  .colab-df-spinner {\n",
              "    border: 2px solid var(--fill-color);\n",
              "    border-color: transparent;\n",
              "    border-bottom-color: var(--fill-color);\n",
              "    animation:\n",
              "      spin 1s steps(1) infinite;\n",
              "  }\n",
              "\n",
              "  @keyframes spin {\n",
              "    0% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "      border-left-color: var(--fill-color);\n",
              "    }\n",
              "    20% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    30% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    40% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    60% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    80% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "    90% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "  }\n",
              "</style>\n",
              "\n",
              "  <script>\n",
              "    async function quickchart(key) {\n",
              "      const quickchartButtonEl =\n",
              "        document.querySelector('#' + key + ' button');\n",
              "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
              "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
              "      try {\n",
              "        const charts = await google.colab.kernel.invokeFunction(\n",
              "            'suggestCharts', [key], {});\n",
              "      } catch (error) {\n",
              "        console.error('Error during call to suggestCharts:', error);\n",
              "      }\n",
              "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
              "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
              "    }\n",
              "    (() => {\n",
              "      let quickchartButtonEl =\n",
              "        document.querySelector('#df-978bcb56-7593-4c60-8870-dcb5751342d1 button');\n",
              "      quickchartButtonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "    })();\n",
              "  </script>\n",
              "</div>\n",
              "    </div>\n",
              "  </div>\n"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# ChatVertexAI Multi-agent LLM Using LangChain"
      ],
      "metadata": {
        "id": "WSPknVyiasPo"
      },
      "id": "WSPknVyiasPo"
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "osigfVmK1VGW"
      },
      "id": "osigfVmK1VGW",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "from contextlib import redirect_stdout\n"
      ],
      "metadata": {
        "id": "2OpIWxZCd9_A"
      },
      "id": "2OpIWxZCd9_A",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "## %%capture captured_output\n",
        "\n",
        "import re\n",
        "import os\n",
        "import sys\n",
        "import base64\n",
        "import vertexai\n",
        "from vertexai.generative_models import GenerativeModel, Part, SafetySetting\n",
        "\n",
        "from langchain_google_vertexai import ChatVertexAI\n",
        "from langchain.tools import tool\n",
        "from langchain.agents import create_tool_calling_agent\n",
        "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
        "from langgraph.prebuilt.tool_executor import ToolExecutor\n",
        "from typing import TypedDict, Annotated,Union\n",
        "from langchain_core.agents import AgentAction, AgentFinish\n",
        "from langchain.agents.output_parsers.tools import ToolAgentAction\n",
        "from langchain_core.messages import BaseMessage\n",
        "import operator\n",
        "from langgraph.graph import END, StateGraph\n",
        "from langgraph.checkpoint.memory import MemorySaver\n",
        "\n",
        "vertexai.init(\n",
        "    project=\"isb-cgc-external-004\",\n",
        "    location=\"us-central1\",\n",
        "    staging_bucket=\"gs://htan_st_datasets/staging_dir/\",\n",
        ")\n",
        "\n",
        "model = \"gemini-1.5-flash-002\"\n",
        "\n",
        "safety_settings = {\n",
        "    HarmCategory.HARM_CATEGORY_UNSPECIFIED: HarmBlockThreshold.BLOCK_NONE,\n",
        "    HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,\n",
        "    HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_ONLY_HIGH,\n",
        "    HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,\n",
        "    HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,\n",
        "}\n",
        "\n",
        "model_kwargs = {\n",
        "    # temperature (float): The sampling temperature controls the degree of\n",
        "    # randomness in token selection.\n",
        "    \"temperature\": 0.28,\n",
        "    # max_output_tokens (int): The token limit determines the maximum amount of\n",
        "    # text output from one prompt.\n",
        "    \"max_output_tokens\": 8000,\n",
        "    # top_p (float): Tokens are selected from most probable to least until\n",
        "    # the sum of their probabilities equals the top-p value.\n",
        "    \"top_p\": 0.95,\n",
        "    # top_k (int): The next token is selected from among the top-k most\n",
        "    # probable tokens.\n",
        "    \"top_k\": 40,\n",
        "    # safety_settings (Dict[HarmCategory, HarmBlockThreshold]): The safety\n",
        "    # settings to use for generating content.\n",
        "    \"safety_settings\": safety_settings,\n",
        "}\n",
        "\n",
        "safety_settings = [\n",
        "    SafetySetting(\n",
        "        category=SafetySetting.HarmCategory.HARM_CATEGORY_HATE_SPEECH,\n",
        "        threshold=SafetySetting.HarmBlockThreshold.OFF\n",
        "    ),\n",
        "    SafetySetting(\n",
        "        category=SafetySetting.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,\n",
        "        threshold=SafetySetting.HarmBlockThreshold.OFF\n",
        "    ),\n",
        "    SafetySetting(\n",
        "        category=SafetySetting.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,\n",
        "        threshold=SafetySetting.HarmBlockThreshold.OFF\n",
        "    ),\n",
        "    SafetySetting(\n",
        "        category=SafetySetting.HarmCategory.HARM_CATEGORY_HARASSMENT,\n",
        "        threshold=SafetySetting.HarmBlockThreshold.OFF\n",
        "    ),\n",
        "]\n",
        "\n",
        "@tool\n",
        "def agent_1(question: str) -> str:\n",
        "    \"\"\"\n",
        "   Download files from HTAN or Synapse. You can directly download files by generating and executing Python code to retrieve datasets using APIs.\n",
        "\n",
        "    Parameters:\n",
        "    question (str): A user-defined question regarding dataset retrieval from HTAN or Synapse.\n",
        "\n",
        "    Returns:\n",
        "    str: Python code enclosed in ```python``` tags, allowing the user to easily extract the code.\n",
        "         Code will be specific to either BigQuery for HTAN datasets or synapseclient for Synapse datasets.\n",
        "    \"\"\"\n",
        "\n",
        "    prompt = f\"\"\"\n",
        "    You are an expert in downloading datasets from Human Tumor Network Atlas (HTAN) and Synapse using APIs. You can directly download files. Return only the python code.\n",
        "\n",
        "    For Synapse Datasets:\n",
        "    You can write the Python code to download the Synapse data directly using the synapseclient by passing the synapse id and the download location. If the download location is not defined, download to ‘./’. You can login using synapseclient.login(silent=True) function.\n",
        "\n",
        "    For Datasets with an HTAN biospecimen ID:\n",
        "    You can use BigQuery in Google Cloud to do this. Assume that all the packages are already installed. You can use: project_id = \"isb-cgc-external-004\".\n",
        "\n",
        "    Important considerations:\n",
        "\n",
        "    1. Please write and execute the Python code enclosed in ```python``` tags so that the user can regex extract the code easily.\n",
        "\n",
        "    2. You do not have to access any data. You are an expert coder, just write the code. When you're asked to load a dataset, use BigQuery to load it.\n",
        "\n",
        "    3. Include all the necessary import packages to run the code. If using any Google Cloud services, you can use:\n",
        "    project_id = \"isb-cgc-external-004\"\n",
        "\n",
        "    4. Do not write any try except blocks.\n",
        "\n",
        "    Make sure you follow the considerations.\n",
        "\n",
        "    Here are some examples:\n",
        "\n",
        "    Question: Can you load the cells from the HTAN biospecimen HTA7_1_3?\n",
        "    Answer:\n",
        "    '''\n",
        "    query = '''\n",
        "    WITH cells AS (\n",
        "      SELECT  CellID, X_centroid, Y_centroid,\n",
        "      FROM `isb-cgc-bq.HTAN.imaging_level4_HMS_mel_mask_current`\n",
        "      WHERE HTAN_Biospecimen_ID = 'HTA7_1_3'\n",
        "    )\n",
        "    SELECT CellID, X_centroid,  Y_centroid\n",
        "    FROM cells\n",
        "    '''\n",
        "    df = pandas_gbq.read_gbq(query, project_id=project_id)\n",
        "    print(df)\n",
        "    '''\n",
        "\n",
        "\n",
        "    \"\"\"\n",
        "\n",
        "\n",
        "\n",
        "    #### Synapse data download : Example 1\n",
        "\n",
        "    prompt += f\"\"\"\n",
        "    Question: Can you load the entityId of all data under HTAN.10xvisium_spatialtranscriptomics_scRNAseq_level4_metadata_current\n",
        "    where the File_Format is hdf5?\n",
        "\n",
        "    Answer:\n",
        "    '''\n",
        "    import pandas_gbq\n",
        "    from google.cloud import bigquery\n",
        "\n",
        "    project_id = \"isb-cgc-external-004\"\n",
        "\n",
        "    query = '''\n",
        "    SELECT entityId\n",
        "    FROM `isb-cgc-bq.HTAN.10xvisium_spatialtranscriptomics_scRNAseq_level4_metadata_current`\n",
        "    WHERE File_Format = 'hdf5'\n",
        "    '''\n",
        "\n",
        "    df = pandas_gbq.read_gbq(query, project_id=project_id)\n",
        "    print(df)\n",
        "    '''\n",
        "\n",
        "    \"\"\"\n",
        "\n",
        "    ##  Synapse example 2\n",
        "\n",
        "    prompt += f\"\"\"\n",
        "    Question: Can you download the synapse data 'syn51133602' to /content/datasets/?\n",
        "    Answer:\n",
        "    '''\n",
        "    import synapseclient\n",
        "    syn = synapseclient.login(silent=True)\n",
        "    entity = syn.get('syn51133602', downloadLocation='/content/datasets')\n",
        "    '''\n",
        "\n",
        "    \"\"\"\n",
        "\n",
        "    ### Spatial example 1\n",
        "\n",
        "    prompt += f\"\"\"\n",
        "    Question: Can you categorize cells within a defined spatial region as either 'Tumor' or 'Other' based on threshold expression levels of specific markers (SOX10_cellRingMask, S100B_cellRingMask, and CD63_cellRingMask)?\n",
        "    Answer:\n",
        "    '''\n",
        "    query = '''\n",
        "      WITH cells AS (\n",
        "      SELECT CellID, X_centroid, Y_centroid,\n",
        "      IF (SOX10_cellRingMask > 3704.5 AND (S100B_cellRingMask > 7589.48 OR CD63_cellRingMask > 570.68),\n",
        "      'Tumor', 'Other') AS celltype\n",
        "      FROM `isb-cgc-bq.HTAN.imaging_level4_HMS_mel_mask_current`\n",
        "      WHERE HTAN_Biospecimen_ID = 'HTA7_1_3')\n",
        "\n",
        "      SELECT CellID, X_centroid, Y_centroid, celltype\n",
        "      FROM cells\n",
        "      WHERE X_centroid > 23076.9 AND X_centroid < 30384.6\n",
        "      AND Y_centroid > 9615.3 AND Y_centroid < 15000\n",
        "      '''\n",
        "    df = pandas_gbq.read_gbq(query, project_id=project_id)\n",
        "    print(df)\n",
        "    \"\"\"\n",
        "\n",
        "    ### Spatial example 2\n",
        "    prompt += f\"\"\"\n",
        "    Question: Can you classify cells as 'Tumor' or 'Other', convert their pixel coordinates to geospatial points, and calculate distances between cell pairs that are within a 20-micrometer proximity threshold\n",
        "    Answer:\n",
        "    query = '''\n",
        "      WITH geodat AS (\n",
        "      SELECT CellID, X_centroid, Y_centroid,\n",
        "      IF (SOX10_cellRingMask > 3704.5 AND (S100B_cellRingMask > 7589.48 OR CD63_cellRingMask > 570.68),\n",
        "      'Tumor', 'Other') AS celltype,\n",
        "      ST_GeogPoint(X_centroid / 368570, Y_centroid / 368570) AS p\n",
        "      FROM `isb-cgc-bq.HTAN.imaging_level4_HMS_mel_mask_current`\n",
        "      WHERE HTAN_Biospecimen_ID = 'HTA7_1_3'\n",
        "      )\n",
        "    SELECT t1.CellID, t1.X_centroid, t1.Y_centroid, t1.p, t1.celltype,\n",
        "    t2.CellID AS CellID_1, t2.X_centroid AS X_centroid_1, t2.Y_centroid AS Y_centroid_1, t2.p AS p_1, t2.celltype AS celltype_1,\n",
        "    ST_Distance(t1.p, t2.p) AS Distance\n",
        "    FROM geodat AS t1\n",
        "    JOIN geodat AS t2\n",
        "    ON ST_DWithin(t1.p, t2.p, 9.29324770787722)\n",
        "    '''\n",
        "    df = pandas_gbq.read_gbq(query, project_id=project_id)\n",
        "    print(df)\n",
        "\n",
        "    \"\"\"\n",
        "\n",
        "    ### Spatial example 3\n",
        "    prompt += f\"\"\"\n",
        "    Question: For each tumor cell within the specified spatial region, calculate the number of neighboring tumor cells\n",
        "    Answer:\n",
        "    query = '''\n",
        "      WITH cellp AS (\n",
        "      SELECT CellID, celltype, CellID_1, celltype_1\n",
        "      FROM `isb-cgc-bq.temp15432.Melanoma_CyCIF_HTA7_1_3_points_within_20um`\n",
        "      WHERE X_centroid > 23076.9 AND X_centroid < 30384.6\n",
        "      AND Y_centroid > 9615.3 AND Y_centroid < 15000)\n",
        "\n",
        "      SELECT CellID, COUNTIF(celltype_1 = 'Tumor') - 1 AS N_Tumor_Cells\n",
        "      FROM cellp\n",
        "      WHERE celltype = 'Tumor'\n",
        "      GROUP BY CellID\n",
        "    '''\n",
        "    df = pandas_gbq.read_gbq(query, project_id=project_id)\n",
        "    print(df)\n",
        "    \"\"\"\n",
        "\n",
        "\n",
        "    ## Metadata\n",
        "\n",
        "    # Create a prompt with the file content included\n",
        "    prompt += f\"\"\"\n",
        "    The HTAN metadata contains a subset of these type of attributes. Also included is a description of these attributes. Internalize this information and use it to answer any queries related to metadata.\n",
        "\n",
        "    Metadata Description:\n",
        "    {metadata_df}\n",
        "    \"\"\"\n",
        "\n",
        "\n",
        "\n",
        "    ## ScRNAseq example 1\n",
        "    prompt += f\"\"\"\n",
        "    Question: What are the counts of unique cells, sex groupings, samples,\n",
        "    cell types, and therapies by development stage in the MSK scRNAseq dataset?\n",
        "    Answer:\n",
        "    '''\n",
        "    query = '''\n",
        "    SELECT\n",
        "      development_stage,\n",
        "      count(distinct(iObs)) AS Number_Cells,\n",
        "      count(distinct(sex)) AS Unique_Sex_Grouping,\n",
        "      count(distinct(donor_id)) AS Number_Samples,\n",
        "      count(distinct(cell_type)) AS Number_Cell_Types,\n",
        "      count(distinct(treatment)) AS Number_Therapies\n",
        "    FROM\n",
        "      `isb-cgc-bq.HTAN.scRNAseq_MSK_SCLC_combined_samples_current`\n",
        "    GROUP BY development_stage\n",
        "    ORDER BY Number_Samples DESC\n",
        "    '''\n",
        "    df = pandas_gbq.read_gbq(query, project_id=project_id)\n",
        "    print(df)\n",
        "    '''\n",
        "\n",
        "    \"\"\"\n",
        "\n",
        "    ## ScRNAseq example 2\n",
        "    prompt += f\"\"\"\n",
        "    Question: How many unique cell types, sex groupings, cells, and samples\n",
        "    are present in the a specific human stage of development in the MSK\n",
        "    scRNAseq dataset (e.g. 74-year-old)?\n",
        "    Answer:\n",
        "    '''\n",
        "    query = '''\n",
        "    SELECT\n",
        "      cell_type,\n",
        "      count(distinct(sex)) AS Unique_Sex_Grouping,\n",
        "      count(distinct(iObs)) AS Number_Cells,\n",
        "      count(distinct(donor_id)) AS Number_Samples\n",
        "    FROM\n",
        "      `isb-cgc-bq.HTAN.scRNAseq_MSK_SCLC_combined_samples_current`\n",
        "    WHERE\n",
        "      development_stage = '74-year-old human stage'\n",
        "    GROUP BY cell_type\n",
        "    '''\n",
        "    df = pandas_gbq.read_gbq(query, project_id=project_id)\n",
        "    print(df)\n",
        "    '''\n",
        "\n",
        "    \"\"\"\n",
        "\n",
        "    ## ScRNAseq example 3\n",
        "    prompt += f\"\"\"\n",
        "    Question: How many genes and cells are associated with each sex and\n",
        "    cell type in the MSK scRNAseq dataset for an individual\n",
        "    (e.g. a 74-year-old human stage)?\n",
        "    Answer:\n",
        "    '''\n",
        "    query = '''\n",
        "    SELECT\n",
        "      sex,\n",
        "      Cell_Type,\n",
        "      count(distinct(feature_name)) AS Number_Genes,\n",
        "      count(distinct(iObs)) AS Number_Cells\n",
        "    FROM\n",
        "      `isb-cgc-bq.HTAN.scRNAseq_MSK_SCLC_combined_samples_current`\n",
        "    WHERE development_stage = '74-year-old human stage'\n",
        "    GROUP BY sex, Cell_Type\n",
        "    ORDER BY Cell_Type DESC\n",
        "    '''\n",
        "    df = pandas_gbq.read_gbq(query, project_id=project_id)\n",
        "    print(df)\n",
        "    '''\n",
        "\n",
        "    \"\"\"\n",
        "\n",
        "    ## ScRNAseq example 4\n",
        "    prompt += f\"\"\"\n",
        "    Question: How many genes and cells are there in each Seurat Cluster\n",
        "    for males and females of the 'epithelial cell' type in the specific\n",
        "    human stage (here 74-year-old)?\n",
        "    Answer:\n",
        "    '''\n",
        "    query = '''\n",
        "    SELECT\n",
        "      sex,\n",
        "      clusters,\n",
        "      Cell_Type,\n",
        "      count(distinct(feature_name)) AS Number_Genes,\n",
        "      count(distinct(iObs)) AS Number_Cells\n",
        "    FROM\n",
        "      `isb-cgc-bq.HTAN.scRNAseq_MSK_SCLC_combined_samples_current`\n",
        "    WHERE development_stage = '74-year-old human stage' AND Cell_Type = 'epithelial cell'\n",
        "    GROUP BY sex, clusters, Cell_Type\n",
        "    ORDER BY clusters ASC\n",
        "    '''\n",
        "    df = pandas_gbq.read_gbq(query, project_id=project_id)\n",
        "    print(df)\n",
        "    '''\n",
        "\n",
        "    \"\"\"\n",
        "\n",
        "   ## ScRNAseq example 5\n",
        "    prompt += f\"\"\"\n",
        "    Question: How do the average expression values for genes differ between\n",
        "    male and female epithelial cells in a specific cluster, and which genes\n",
        "    show the greatest differences (here cluster 41 of the 74-year-old\n",
        "    human stage)?\n",
        "    Answer:\n",
        "    '''\n",
        "    query = '''\n",
        "    SELECT\n",
        "      A.feature_name,\n",
        "      A.avg_counts_clust10 AS female_avg_counts,\n",
        "      B.avg_counts_clust10 AS male_avg_counts,\n",
        "      A.avg_counts_clust10 - B.avg_counts_clust10 AS mean_diff\n",
        "    FROM (\n",
        "      SELECT\n",
        "        feature_name,\n",
        "        AVG(X_value) AS avg_counts_clust10\n",
        "      FROM\n",
        "        `isb-cgc-bq.HTAN.scRNAseq_MSK_SCLC_combined_samples_current`\n",
        "      WHERE development_stage = '74-year-old human stage' AND Cell_Type = 'epithelial cell' AND clusters = '41' AND sex = 'female'\n",
        "      GROUP BY feature_name\n",
        "    ) AS A\n",
        "    INNER JOIN (\n",
        "      SELECT\n",
        "        feature_name,\n",
        "        AVG(X_value) AS avg_counts_clust10\n",
        "      FROM\n",
        "        `isb-cgc-bq.HTAN.scRNAseq_MSK_SCLC_combined_samples_current`\n",
        "      WHERE development_stage = '74-year-old human stage' AND Cell_Type = 'epithelial cell' AND clusters = '41' AND sex = 'male'\n",
        "      GROUP BY feature_name\n",
        "    ) AS B\n",
        "    ON A.feature_name = B.feature_name\n",
        "    ORDER BY mean_diff DESC\n",
        "    '''\n",
        "    df = pandas_gbq.read_gbq(query, project_id=project_id)\n",
        "    print(df)\n",
        "    '''\n",
        "\n",
        "    \"\"\"\n",
        "\n",
        "\n",
        "    # Add the dynamic user question and the placeholder for the answer\n",
        "    prompt += f\"\"\"\n",
        "    Question: {question}\n",
        "    Answer:\n",
        "    \"\"\"\n",
        "\n",
        "    vertexai.init(project=\"isb-cgc-external-004\", location=\"us-central1\")\n",
        "    model = GenerativeModel(\n",
        "        \"gemini-1.5-flash-002\",\n",
        "    )\n",
        "    responses = model.generate_content(\n",
        "        [prompt],\n",
        "        generation_config=generation_config,\n",
        "        safety_settings=safety_settings,\n",
        "        stream=False,\n",
        "        # tools='code_execution'\n",
        "    )\n",
        "\n",
        "    # for response in responses:\n",
        "    #     print(response.text, end=\"\")\n",
        "\n",
        "    pattern = r'```python\\n(.*?)\\n```'\n",
        "\n",
        "    # Extract the code\n",
        "    match = re.search(pattern, responses.text, re.DOTALL)\n",
        "\n",
        "    if match:\n",
        "        python_code = match.group(1).strip()\n",
        "        print(python_code)\n",
        "        exec(python_code)\n",
        "\n",
        "    else:\n",
        "        print(\"Agent 1: Executed Python code.\")\n",
        "\n",
        "@tool\n",
        "def agent_2(question:str) -> str:\n",
        "    \"\"\"\n",
        "    Generate and execute Python code to query metadata or retrieve datasets from HTAN via Google Cloud\n",
        "    using BigQuery and relevant Python packages.\n",
        "\n",
        "    Parameters:\n",
        "    question (str): A user-defined question regarding metadata or dataset retrieval from HTAN.\n",
        "\n",
        "    Returns:\n",
        "    str: Python code enclosed in ```python``` tags for easy extraction. The code may include\n",
        "         BigQuery queries or use of libraries like pandas, scanpy, or squidpy depending on\n",
        "         the question.\n",
        "\n",
        "    Notes:\n",
        "    - Uses Google BigQuery with project ID 'isb-cgc-external-004' to access HTAN datasets.\n",
        "    - Responds to general questions with a conversational response if code is not required.\n",
        "    - Packages are imported as needed in the generated code.\n",
        "    \"\"\"\n",
        "\n",
        "    prompt = f\"\"\"\n",
        "    You are a bioinformatics expert coder using the Human Tumor Network Atlas (HTAN) datasets via Google Cloud and Jupyter.\n",
        "    You can write Python code to answer questions by writing BigQuery queries, pandas, scanpy, squidpy, and commot code.\n",
        "    You have access to an AnnData object named 'adata' that contains:\n",
        "    - Gene expression data in adata.X\n",
        "    - Cell metadata in adata.obs, including 'kmeans_9_clusters' for cell types\n",
        "    - Gene names in adata.var_names\n",
        "    - Spatial coordinates in adata.obsm['spatial'] if available\n",
        "\n",
        "    Important technical details:\n",
        "    - Gene expression data is stored as a sparse matrix and needs to be converted using .toarray() or scipy.sparse methods\n",
        "    - Always handle sparse matrices appropriately\n",
        "    - Add clear comments explaining the analysis steps\n",
        "    - Include proper error handling for missing genes or data\n",
        "    - Use sc.pl.spatial() for spatial plots\n",
        "    - Use only 'cmap' parameter (not color_map) for color mapping\n",
        "    - Include proper error handling for missing coordinates or genes\n",
        "    - For multiple plots, use ncols parameter to arrange them\n",
        "    - Add clear comments explaining the analysis steps\n",
        "    - VERY IMPORTANT: Include all the necessary import packages to run the code.\n",
        "\n",
        "    For any commands that have the parameter 'database_name' or 'database' please use 'CellChat'.\n",
        "    The following results are attached after running ct.tl.spatial_communication. Please note that 'Fgf1-Fgfr1' can be replaced wiht other ligand-receptor pairs:\n",
        "    - adata.uns['commot-user_database-info']: Ligand-receptor database used.\n",
        "    - adata.obsm['commot-user_database-sum-sender']: Total sent signals per LR pair and pathway.\n",
        "    - adata.obsm['commot-user_database-sum-receiver']: Total received signals per LR pair and pathway.\n",
        "    - adata.obsm['commot_sender_vf-user_database-Fgf1-Fgfr1']: Signaling directions for sent signals.\n",
        "    - adata.obsm['commot_receiver_vf-user_database-Fgf1-Fgfr1']: Signaling directions for received signals.\n",
        "    - adata.obsp['commot-user_database-Fgf1-Fgfr1']: Sparse matrix of cell-cell communication scores per spot.\n",
        "\n",
        "    When responding to questions:\n",
        "    1. Analyze what type of operation is being requested\n",
        "    2. Generate appropriate Python code using scanpy, pandas, matplotlib, and other relevant libraries\n",
        "    3. For questions that require code, return ONLY the Python code within ```python``` tags\n",
        "    4. Include proper error handling and data validation\n",
        "    5. Add clear comments explaining the analysis steps\n",
        "\n",
        "    Examples of queries and their answers:\n",
        "\n",
        "    Question: How do preprocess my adata object?\n",
        "    Answer:\n",
        "    '''\n",
        "    import commot as ct\n",
        "    import scanpy as sc\n",
        "    import pandas as pd\n",
        "    import numpy as np\n",
        "    from scipy import sparse\n",
        "    import matplotlib.pyplot as plt\n",
        "    sc.pp.normalize_total(adata, inplace=True)\n",
        "    sc.pp.log1p(adata)\n",
        "    '''\n",
        "\n",
        "    Q: \"Show me the distribution of cells across kmeans_9_clusters\"\n",
        "    A: ```python\n",
        "    import scanpy as sc\n",
        "    import pandas as pd\n",
        "    import matplotlib.pyplot as plt\n",
        "\n",
        "    # Get cluster counts\n",
        "    cluster_counts = adata.obs['kmeans_9_clusters'].value_counts().sort_index()\n",
        "\n",
        "    # Create bar plot\n",
        "    plt.figure(figsize=(10, 6))\n",
        "    cluster_counts.plot(kind='bar')\n",
        "    plt.title('Distribution of Cells Across kmeans_9_clusters')\n",
        "    plt.xlabel('Cluster')\n",
        "    plt.ylabel('Number of Cells')\n",
        "    plt.xticks(rotation=45)\n",
        "    plt.tight_layout()\n",
        "    plt.show()\n",
        "\n",
        "    # Print numerical summary\n",
        "    print(\"\\\\nNumerical distribution:\")\n",
        "    print(pd.DataFrame({{\n",
        "        'Cluster': cluster_counts.index,\n",
        "        'Count': cluster_counts.values,\n",
        "        'Percentage': (cluster_counts.values / len(adata) * 100).round(2)\n",
        "    }}))\n",
        "    ```\n",
        "\n",
        "    Q: \"Find and plot the top 5 genes with highest Moran's I values\"\n",
        "    A: ```python\n",
        "    import pandas as pd\n",
        "    import matplotlib.pyplot as plt\n",
        "    import scanpy as sc\n",
        "\n",
        "    # Create DataFrame with Moran's I statistics\n",
        "    morans_df = pd.DataFrame({{\n",
        "        'Gene': adata.var_names,\n",
        "        'Morans_I': adata.var['Morans_I'],\n",
        "        'Adj_P_Value': adata.var['Morans_I_adj_p_val']\n",
        "    }})\n",
        "\n",
        "    # Sort by Moran's I value and get top 5\n",
        "    top_genes = morans_df.sort_values('Morans_I', ascending=False).head(5)\n",
        "\n",
        "    # Create spatial plots for top genes\n",
        "    fig, axes = plt.subplots(2, 3, figsize=(15, 10))\n",
        "    axes = axes.flatten()\n",
        "\n",
        "    # Plot each gene\n",
        "    for idx, (gene, mi_value) in enumerate(zip(top_genes['Gene'], top_genes['Morans_I'])):\n",
        "        sc.pl.spatial(adata,\n",
        "                     color=gene,\n",
        "                     title=f'{{gene}}\\\\nMoran\\\\'s I = {{mi_value:.3f}}',\n",
        "                     ax=axes[idx],\n",
        "                     show=False,\n",
        "                     cmap='viridis')\n",
        "\n",
        "    plt.tight_layout()\n",
        "    plt.show()\n",
        "\n",
        "    # Print summary table\n",
        "    print(\"\\\\nTop 5 spatially autocorrelated genes:\")\n",
        "    print(top_genes)\n",
        "    ```\n",
        "\n",
        "    Q: \"Generate a UMAP visualization comparing kmeans_9_clusters with kmeans_10_clusters\"\n",
        "    A: ```python\n",
        "    import scanpy as sc\n",
        "    import matplotlib.pyplot as plt\n",
        "\n",
        "    # Compute UMAP if not already present\n",
        "    if 'X_umap' not in adata.obsm_keys():\n",
        "        sc.pp.neighbors(adata, n_pcs=30)\n",
        "        sc.tl.umap(adata)\n",
        "\n",
        "    # Create subplot with both clustering results\n",
        "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))\n",
        "\n",
        "    # Plot UMAP with kmeans_9_clusters\n",
        "    sc.pl.umap(adata,\n",
        "               color='kmeans_9_clusters',\n",
        "               title='kmeans_9_clusters',\n",
        "               ax=ax1,\n",
        "               show=False)\n",
        "\n",
        "    # Plot UMAP with kmeans_10_clusters\n",
        "    sc.pl.umap(adata,\n",
        "               color='kmeans_10_clusters',\n",
        "               title='kmeans_10_clusters',\n",
        "               ax=ax2,\n",
        "               show=False)\n",
        "\n",
        "    plt.tight_layout()\n",
        "    plt.show()\n",
        "\n",
        "    # Print cluster sizes\n",
        "    print(\"\\\\nkmeans_9_clusters distribution:\")\n",
        "    print(adata.obs['kmeans_9_clusters'].value_counts())\n",
        "    print(\"\\\\nkmeans_10_clusters distribution:\")\n",
        "    print(adata.obs['kmeans_10_clusters'].value_counts())\n",
        "    ```\n",
        "\n",
        "    Q: \"Calculate correlation between the top 3 most highly expressed genes\"\n",
        "    A: ```python\n",
        "    import pandas as pd\n",
        "    import numpy as np\n",
        "    import seaborn as sns\n",
        "    import matplotlib.pyplot as plt\n",
        "    from scipy import sparse\n",
        "\n",
        "    # Calculate mean expression for each gene\n",
        "    if sparse.issparse(adata.X):\n",
        "        mean_expr = np.array(adata.X.mean(axis=0)).flatten()\n",
        "    else:\n",
        "        mean_expr = np.mean(adata.X, axis=0)\n",
        "\n",
        "    # Get top 3 genes\n",
        "    top_genes_idx = np.argsort(mean_expr)[-3:]\n",
        "    top_genes = adata.var_names[top_genes_idx]\n",
        "\n",
        "    # Extract expression data for top genes\n",
        "    expr_matrix = adata[:, top_genes].X\n",
        "    if sparse.issparse(expr_matrix):\n",
        "        expr_matrix = expr_matrix.toarray()\n",
        "\n",
        "    # Create correlation matrix\n",
        "    corr_df = pd.DataFrame(expr_matrix, columns=top_genes)\n",
        "    corr_matrix = corr_df.corr()\n",
        "\n",
        "    # Plot correlation heatmap\n",
        "    plt.figure(figsize=(8, 6))\n",
        "    sns.heatmap(corr_matrix,\n",
        "                annot=True,\n",
        "                cmap='coolwarm',\n",
        "                vmin=-1,\n",
        "                vmax=1,\n",
        "                center=0)\n",
        "    plt.title('Correlation between top 3 highly expressed genes')\n",
        "    plt.tight_layout()\n",
        "    plt.show()\n",
        "\n",
        "    # Print mean expression values\n",
        "    print(\"\\\\nMean expression values:\")\n",
        "    for gene, mean_exp in zip(top_genes, mean_expr[top_genes_idx]):\n",
        "        print(f\"{{gene}}: {{mean_exp:.2f}}\")\n",
        "    ```\n",
        "\n",
        "    Question: What are the different ligand and receptor pairs in the TGFb pathway?\n",
        "    Answer:\n",
        "    '''\n",
        "    import commot as ct\n",
        "    import scanpy as sc\n",
        "    import pandas as pd\n",
        "    import numpy as np\n",
        "    import matplotlib.pyplot as plt\n",
        "    df_ligrec=ct.pp.ligand_receptor_database(database='CellChat', species='human')\n",
        "    tgfb = df_ligrec[df_ligrec[2]=='TGFb']\n",
        "    tgfb\n",
        "    '''\n",
        "    Question: Perform cell-cell communication analysis for the TGFb pathway.\n",
        "    Only access the df_ligrec database with .iloc.\n",
        "    '''\n",
        "    import commot as ct\n",
        "    import scanpy as sc\n",
        "    import pandas as pd\n",
        "    import numpy as np\n",
        "    import matplotlib.pyplot as plt\n",
        "    df_ligrec=ct.pp.ligand_receptor_database(database='CellChat', species='human')\n",
        "    tgfb = df_ligrec[df_ligrec.iloc[:,2]=='TGFb']\n",
        "    ct.tl.spatial_communication(adata,\n",
        "    database_name='cellchat', df_ligrec=tgfb, dis_thr=500, heteromeric=True, pathway_sum=True)\n",
        "\n",
        "    Question: Can you construct a cell-cell communication networks of the TGFb pathway between cells within 500 µm?\n",
        "    Answer:\n",
        "    '''\n",
        "    import commot as ct\n",
        "    import scanpy as sc\n",
        "    import pandas as pd\n",
        "    import numpy as np\n",
        "    import matplotlib.pyplot as plt\n",
        "    df_ligrec=ct.pp.ligand_receptor_database(database='CellChat', species='human')\n",
        "    tgfb = df_ligrec[df_ligrec.iloc[:,2]=='TGFb']\n",
        "    ct.tl.spatial_communication(adata,\n",
        "    database_name='cellchat', df_ligrec=tgfb, dis_thr=500, heteromeric=True, pathway_sum=True)\n",
        "    '''\n",
        "\n",
        "    Question: How do I visualize the amount of sent and received signal between SEMA3A and NRP1_PLXNA1?\n",
        "    Answer:\n",
        "    Here's how to plot the signaling levels for the SEMA3A-NRP1_PLXNA1 pair.\n",
        "    '''\n",
        "    import commot as ct\n",
        "    import scanpy as sc\n",
        "    import pandas as pd\n",
        "    import numpy as np\n",
        "    import matplotlib.pyplot as plt\n",
        "    #define the LR pair\n",
        "    LR = np.array([['SEMA3A', 'NRP1_PLXNA2', 'SEMA3']], dtype=str)\n",
        "    LR = pd.DataFrame(data=LR)\n",
        "    ct.tl.spatial_communication(adata,\n",
        "    database_name='cellchat', df_ligrec=LR, dis_thr=500, heteromeric=True, pathway_sum=True)\n",
        "    pts = adata.obsm['spatial']\n",
        "    s = adata.obsm['commot-cellchat-sum-sender']['s-SEMA3A']\n",
        "    r = adata.obsm['commot-cellchat-sum-receiver']['r-NRP1-PLXNA2']\n",
        "    fig, ax = plt.subplots(1,2, figsize=(10,4))\n",
        "    ax[0].scatter(pts[:,0], pts[:,1], c=s, s=5, cmap='Blues')\n",
        "    ax[0].set_title('Sender')\n",
        "    ax[1].scatter(pts[:,0], pts[:,1], c=r, s=5, cmap='Reds')\n",
        "    ax[1].set_title('Receiver')\n",
        "    '''\n",
        "\n",
        "    Question: How can I visualize signaling directions for TGFB1-TGFBR1_TGFBR2 as vector fields?\n",
        "    Answer:\n",
        "    '''\n",
        "    import commot as ct\n",
        "    import scanpy as sc\n",
        "    import pandas as pd\n",
        "    import numpy as np\n",
        "    import matplotlib.pyplot as plt\n",
        "    ct.tl.communication_direction(adata, database_name='cellchat', lr_pair=('TGFB1','TGFBR1_TGFBR2'), k=5)\n",
        "    ct.pl.plot_cell_communication(adata, database_name='cellchat', lr_pair=('TGFB1','TGFBR1_TGFBR2'), plot_method='grid', background_legend=True,\n",
        "    scale=0.00003, ndsize=8, grid_density=0.4, summary='sender', background='image', clustering='leiden', cmap='Alphabet',\n",
        "    normalize_v = True, normalize_v_quantile=0.995)\n",
        "    '''\n",
        "\n",
        "    Question: Can you show me the cell-cell interaction between TGFB1 and TGFBR1_TGFBR2?\n",
        "    Answer:\n",
        "    '''\n",
        "    import commot as ct\n",
        "    import scanpy as sc\n",
        "    import pandas as pd\n",
        "    import numpy as np\n",
        "    import matplotlib.pyplot as plt\n",
        "    ct.tl.communication_direction(adata, database_name='cellchat', lr_pair=('TGFB1','TGFBR1_TGFBR2'), k=5)\n",
        "    ct.pl.plot_cell_communication(adata, database_name='cellchat', lr_pair=('TGFB1','TGFBR1_TGFBR2'), plot_method='grid', background_legend=True,\n",
        "    scale=0.00003, ndsize=8, grid_density=0.4, summary='receiver', background='summary', clustering='leiden', cmap='Reds',\n",
        "    normalize_v = True, normalize_v_quantile=0.995)\n",
        "    '''\n",
        "\n",
        "    Current question to analyze: {question}\n",
        "\n",
        "    First, analyze the type of operation requested and then provide the appropriate code or response.\n",
        "    \"\"\"\n",
        "\n",
        "    vertexai.init(project=\"isb-cgc-external-004\", location=\"us-central1\")\n",
        "    model = GenerativeModel(\n",
        "        \"gemini-1.5-flash-002\",\n",
        "    )\n",
        "    responses = model.generate_content(\n",
        "        [prompt],\n",
        "        generation_config=generation_config,\n",
        "        safety_settings=safety_settings,\n",
        "        stream=False,\n",
        "        # tools='code_execution'\n",
        "    )\n",
        "\n",
        "\n",
        "    pattern = r'```python\\n(.*?)\\n```'\n",
        "    match = re.search(pattern, responses.text, re.DOTALL)\n",
        "\n",
        "    if match:\n",
        "        python_code = match.group(1).strip()\n",
        "        print(\"Generated code:\", python_code)\n",
        "\n",
        "        # Dictionary to capture local variables from exec\n",
        "        local_vars = {}\n",
        "\n",
        "        # Execute the code in the local_vars dictionary to capture output\n",
        "        exec(python_code, globals(), local_vars)\n",
        "\n",
        "        # Update the global scope with new variables\n",
        "        globals().update(local_vars)\n",
        "\n",
        "        print(\"Agent 2: Executed Python code with accessible variables.\")\n",
        "    else:\n",
        "        print(responses.text)\n",
        "        print(\"Agent 2: No executable code found.\")\n",
        "\n",
        "toolkit = [agent_1, agent_2]\n",
        "\n",
        "from IPython.display import Markdown, display\n",
        "\n",
        "def agent_3(question:str) -> str:\n",
        "    \"\"\"\n",
        "    Given output of a tool calling agent, convert the output to a human readable format.\n",
        "\n",
        "    Parameters:\n",
        "    question (str): Raw outputs from the tool calling agent.\n",
        "\n",
        "    Returns:\n",
        "    str: Human readable version of the raw output string.\n",
        "    \"\"\"\n",
        "\n",
        "    prompt = \"\"\"You are an experienced writer and editor. When given an input, convert it into a well formatted string output. If the input string is blank, just say that there are no outputs generated for the current query.\n",
        "\n",
        "    Here are some examples:\n",
        "    Example 1:\n",
        "    Input: '''\n",
        "    Generated code: import scanpy as sc\n",
        "    import pandas as pd\n",
        "\n",
        "    adata = sc.read_h5ad(\"/content/demo_data/6723_KL_1_unfiltered.h5ad\")\n",
        "    print(adata)\n",
        "    AnnData object with n_obs × n_vars = 4992 × 19074\n",
        "    obs: 'in_tissue', 'array_row', 'array_col', 'kmeans_7_clusters', 'kmeans_10_clusters', 'kmeans_4_clusters', 'kmeans_2_clusters', 'kmeans_6_clusters', 'graphclust', 'kmeans_3_clusters', 'kmeans_8_clusters', 'kmeans_9_clusters', 'kmeans_5_clusters'\n",
        "    var: 'gene_ids', 'feature_types', 'genome', 'n_cells', 'Morans_I', 'Morans_I_p_val', 'Morans_I_adj_p_val', 'Feature Counts in Spots Under Tissue', 'Median Normalized Average Counts', 'Barcodes Detected per Feature'\n",
        "    uns: 'spatial'\n",
        "    obsm: 'spatial'\n",
        "    Agent 2: Executed Python code with accessible variables.\n",
        "    The agent action is tool='agent_2' tool_input={'question': 'Can you use scanpy to load the /content/demo_data/6723_KL_1_unfiltered.h5ad?'} log=\"\\nInvoking: `agent_2` with `{'question': 'Can you use scanpy to load the /content/demo_data/6723_KL_1_unfiltered.h5ad?'}`\\n\\n\\n\" message_log=[AIMessage(content='', additional_kwargs={'function_call': {'name': 'agent_2', 'arguments': '{\"question\": \"Can you use scanpy to load the /content/demo_data/6723_KL_1_unfiltered.h5ad?\"}'}}, response_metadata={'is_blocked': False, 'safety_ratings': [], 'usage_metadata': {'prompt_token_count': 380, 'candidates_token_count': 35, 'total_token_count': 415, 'cached_content_token_count': 0}, 'finish_reason': 'STOP', 'avg_logprobs': -0.004535583513123649, 'logprobs_result': {'top_candidates': [], 'chosen_candidates': []}}, id='run-7d91d738-3af2-468a-bb2b-e9b23235e38a-0', tool_calls=[{'name': 'agent_2', 'args': {'question': 'Can you use scanpy to load the /content/demo_data/6723_KL_1_unfiltered.h5ad?'}, 'id': '62e66320-2760-4403-9049-43fd7a676e0f', 'type': 'tool_call'}], usage_metadata={'input_tokens': 380, 'output_tokens': 35, 'total_tokens': 415})] tool_call_id='62e66320-2760-4403-9049-43fd7a676e0f'\n",
        "    The tool result is: None\n",
        "    {'intermediate_steps': [(ToolAgentAction(tool='agent_2', tool_input={'question': 'Can you use scanpy to load the /content/demo_data/6723_KL_1_unfiltered.h5ad?'}, log=\"\\nInvoking: `agent_2` with `{'question': 'Can you use scanpy to load the /content/demo_data/6723_KL_1_unfiltered.h5ad?'}`\\n\\n\\n\", message_log=[AIMessage(content='', additional_kwargs={'function_call': {'name': 'agent_2', 'arguments': '{\"question\": \"Can you use scanpy to load the /content/demo_data/6723_KL_1_unfiltered.h5ad?\"}'}}, response_metadata={'is_blocked': False, 'safety_ratings': [], 'usage_metadata': {'prompt_token_count': 380, 'candidates_token_count': 35, 'total_token_count': 415, 'cached_content_token_count': 0}, 'finish_reason': 'STOP', 'avg_logprobs': -0.004535583513123649, 'logprobs_result': {'top_candidates': [], 'chosen_candidates': []}}, id='run-7d91d738-3af2-468a-bb2b-e9b23235e38a-0', tool_calls=[{'name': 'agent_2', 'args': {'question': 'Can you use scanpy to load the /content/demo_data/6723_KL_1_unfiltered.h5ad?'}, 'id': '62e66320-2760-4403-9049-43fd7a676e0f', 'type': 'tool_call'}], usage_metadata={'input_tokens': 380, 'output_tokens': 35, 'total_tokens': 415})], tool_call_id='62e66320-2760-4403-9049-43fd7a676e0f'), 'None')]}\n",
        "\n",
        "    '''\n",
        "    Well formated output:\n",
        "\n",
        "    '''The dataset is loaded to the anndata object named adata. There are 4992 cells or spots, 19074 genes, and 13 feature columns. The data object is printed below:\n",
        "\n",
        "    AnnData object with n_obs × n_vars = 4992 × 19074\n",
        "    obs: 'in_tissue', 'array_row', 'array_col', 'kmeans_7_clusters', 'kmeans_10_clusters', 'kmeans_4_clusters', 'kmeans_2_clusters', 'kmeans_6_clusters', 'graphclust', 'kmeans_3_clusters', 'kmeans_8_clusters', 'kmeans_9_clusters', 'kmeans_5_clusters'\n",
        "    var: 'gene_ids', 'feature_types', 'genome', 'n_cells', 'Morans_I', 'Morans_I_p_val', 'Morans_I_adj_p_val', 'Feature Counts in Spots Under Tissue', 'Median Normalized Average Counts', 'Barcodes Detected per Feature'\n",
        "    uns: 'spatial'\n",
        "    obsm: 'spatial'\n",
        "    '''\n",
        "\n",
        "    Example 2:\n",
        "    Input: '''\n",
        "    Welcome, arun.das!\n",
        "\n",
        "    INFO:synapseclient_default:Welcome, arun.das!\n",
        "\n",
        "    Downloading files:  77%|███████▋  | 16.8M/21.8M [00:00<00:00, 28.0MB/s, syn51133599]Downloaded syn51133599 to /content/demo_data/8578_AS_1_unfiltered.h5ad\n",
        "    [INFO] Downloaded syn51133599 to /content/demo_data/8578_AS_1_unfiltered.h5ad\n",
        "    Downloading files: 100%|██████████| 21.8M/21.8M [00:00<00:00, 31.8MB/s, syn51133599]INFO:synapseclient_default:Downloaded syn51133599 to /content/demo_data/8578_AS_1_unfiltered.h5ad\n",
        "    Downloading files: 100%|██████████| 21.8M/21.8M [00:00<00:00, 31.6MB/s, syn51133599]import synapseclient\n",
        "    syn = synapseclient.login(silent=True)\n",
        "    entity = syn.get('syn51133599', downloadLocation='/content/demo_data')\n",
        "    The agent action is tool='agent_1' tool_input={'question': 'Can you download the synapse dataset syn51133599 to /content/demo_data/?'} log=\"\\nInvoking: `agent_1` with `{'question': 'Can you download the synapse dataset syn51133599 to /content/demo_data/?'}`\\n\\n\\n\" message_log=[AIMessage(content='', additional_kwargs={'function_call': {'name': 'agent_1', 'arguments': '{\"question\": \"Can you download the synapse dataset syn51133599 to /content/demo_data/?\"}'}}, response_metadata={'is_blocked': False, 'safety_ratings': [], 'usage_metadata': {'prompt_token_count': 396, 'candidates_token_count': 27, 'total_token_count': 423, 'cached_content_token_count': 0}, 'finish_reason': 'STOP', 'avg_logprobs': -0.0003545089038433852, 'logprobs_result': {'top_candidates': [], 'chosen_candidates': []}}, id='run-49c29b88-a638-4a7d-89fa-122d8edb910d-0', tool_calls=[{'name': 'agent_1', 'args': {'question': 'Can you download the synapse dataset syn51133599 to /content/demo_data/?'}, 'id': '48d647d4-94d2-4cae-bd6f-f5d5e2240945', 'type': 'tool_call'}], usage_metadata={'input_tokens': 396, 'output_tokens': 27, 'total_tokens': 423})] tool_call_id='48d647d4-94d2-4cae-bd6f-f5d5e2240945'\n",
        "    The tool result is: None\n",
        "    ----\n",
        "    '''\n",
        "\n",
        "    Well formated output:\n",
        "    '''\n",
        "    The synapse dataset with synapse id syn51133599 is successfully downloaded to /content/demo_data/8578_AS_1_unfiltered.h5ad.\n",
        "    '''\n",
        "\n",
        "\n",
        "    Example 3:\n",
        "    Input: '''\n",
        "    {'agent_outcome': [ToolAgentAction(tool='agent_1', tool_input={'question': 'Can you load the entityId of all data under HTAN.10xvisium_spatialtranscriptomics_scRNAseq_level4_metadata_current where the File_Format is hdf5?'}, log=\"\\nInvoking: `agent_1` with `{'question': 'Can you load the entityId of all data under HTAN.10xvisium_spatialtranscriptomics_scRNAseq_level4_metadata_current where the File_Format is hdf5?'}`\\n\\n\\n\", message_log=[AIMessage(content='', additional_kwargs={'function_call': {'name': 'agent_1', 'arguments': '{\"question\": \"Can you load the entityId of all data under HTAN.10xvisium_spatialtranscriptomics_scRNAseq_level4_metadata_current where the File_Format is hdf5?\"}'}}, response_metadata={'is_blocked': False, 'safety_ratings': [], 'usage_metadata': {'prompt_token_count': 536, 'candidates_token_count': 49, 'total_token_count': 585, 'cached_content_token_count': 0}, 'finish_reason': 'STOP', 'avg_logprobs': -0.0002904396352111077, 'logprobs_result': {'top_candidates': [], 'chosen_candidates': []}}, id='run-b2cf906b-d897-4164-b743-1120419ebdf9-0', tool_calls=[{'name': 'agent_1', 'args': {'question': 'Can you load the entityId of all data under HTAN.10xvisium_spatialtranscriptomics_scRNAseq_level4_metadata_current where the File_Format is hdf5?'}, 'id': '6987b536-4024-4b91-a0b1-720e41b1076e', 'type': 'tool_call'}], usage_metadata={'input_tokens': 536, 'output_tokens': 49, 'total_tokens': 585})], tool_call_id='6987b536-4024-4b91-a0b1-720e41b1076e')]}\n",
        "    import pandas_gbq\n",
        "    from google.cloud import bigquery\n",
        "    import synapseclient\n",
        "\n",
        "    project_id = \"isb-cgc-external-004\"\n",
        "\n",
        "    query = '''\n",
        "    SELECT entityId\n",
        "    FROM `isb-cgc-bq.HTAN.10xvisium_spatialtranscriptomics_scRNAseq_level4_metadata_current`\n",
        "    WHERE File_Format = 'hdf5'\n",
        "    '''\n",
        "\n",
        "    df = pandas_gbq.read_gbq(query, project_id=project_id)\n",
        "    print(df)\n",
        "\n",
        "    syn = synapseclient.login(silent=True)\n",
        "    entity = syn.get('syn51133602', downloadLocation='/content/datasets')\n",
        "    Downloading: 100%|##########|\n",
        "        entityId\n",
        "    0   syn51133519\n",
        "    1   syn51133520\n",
        "    2   syn51133521\n",
        "    3   syn51133522\n",
        "    4   syn51133523\n",
        "    5   syn51133524\n",
        "    6   syn51133525\n",
        "    7   syn51133526\n",
        "    8   syn51133527\n",
        "    9   syn51133528\n",
        "    10  syn51133529\n",
        "    11  syn51133530\n",
        "    12  syn51133531\n",
        "    13  syn51133532\n",
        "    14  syn51133533\n",
        "    15  syn51133534\n",
        "    16  syn51133537\n",
        "    17  syn51133540\n",
        "    18  syn51133578\n",
        "    19  syn51133580\n",
        "    20  syn51133583\n",
        "    21  syn51133587\n",
        "    22  syn51133591\n",
        "    23  syn51133592\n",
        "    24  syn51133593\n",
        "    25  syn51133595\n",
        "    26  syn51133596\n",
        "    27  syn51133597\n",
        "    28  syn51133598\n",
        "    29  syn51133599\n",
        "    30  syn51133600\n",
        "    31  syn51133601\n",
        "    32  syn51133602\n",
        "    33  syn51133603\n",
        "    34  syn51133604\n",
        "    35  syn51133605\n",
        "    36  syn51133606\n",
        "    37  syn51133607\n",
        "    38  syn51133608\n",
        "    39  syn51133609\n",
        "    40  syn51133612\n",
        "    The agent action is tool='agent_1' tool_input={'question': 'Can you load the entityId of all data under HTAN.10xvisium_spatialtranscriptomics_scRNAseq_level4_metadata_current where the File_Format is hdf5?'} log=\"\\nInvoking: `agent_1` with `{'question': 'Can you load the entityId of all data under HTAN.10xvisium_spatialtranscriptomics_scRNAseq_level4_metadata_current where the File_Format is hdf5?'}`\\n\\n\\n\" message_log=[AIMessage(content='', additional_kwargs={'function_call': {'name': 'agent_1', 'arguments': '{\"question\": \"Can you load the entityId of all data under HTAN.10xvisium_spatialtranscriptomics_scRNAseq_level4_metadata_current where the File_Format is hdf5?\"}'}}, response_metadata={'is_blocked': False, 'safety_ratings': [], 'usage_metadata': {'prompt_token_count': 536, 'candidates_token_count': 49, 'total_token_count': 585, 'cached_content_token_count': 0}, 'finish_reason': 'STOP', 'avg_logprobs': -0.0002904396352111077, 'logprobs_result': {'top_candidates': [], 'chosen_candidates': []}}, id='run-b2cf906b-d897-4164-b743-1120419ebdf9-0', tool_calls=[{'name': 'agent_1', 'args': {'question': 'Can you load the entityId of all data under HTAN.10xvisium_spatialtranscriptomics_scRNAseq_level4_metadata_current where the File_Format is hdf5?'}, 'id': '6987b536-4024-4b91-a0b1-720e41b1076e', 'type': 'tool_call'}], usage_metadata={'input_tokens': 536, 'output_tokens': 49, 'total_tokens': 585})] tool_call_id='6987b536-4024-4b91-a0b1-720e41b1076e'\n",
        "    The tool result is: None\n",
        "    {'intermediate_steps': [(ToolAgentAction(tool='agent_1', tool_input={'question': 'Can you load the entityId of all data under HTAN.10xvisium_spatialtranscriptomics_scRNAseq_level4_metadata_current where the File_Format is hdf5?'}, log=\"\\nInvoking: `agent_1` with `{'question': 'Can you load the entityId of all data under HTAN.10xvisium_spatialtranscriptomics_scRNAseq_level4_metadata_current where the File_Format is hdf5?'}`\\n\\n\\n\", message_log=[AIMessage(content='', additional_kwargs={'function_call': {'name': 'agent_1', 'arguments': '{\"question\": \"Can you load the entityId of all data under HTAN.10xvisium_spatialtranscriptomics_scRNAseq_level4_metadata_current where the File_Format is hdf5?\"}'}}, response_metadata={'is_blocked': False, 'safety_ratings': [], 'usage_metadata': {'prompt_token_count': 536, 'candidates_token_count': 49, 'total_token_count': 585, 'cached_content_token_count': 0}, 'finish_reason': 'STOP', 'avg_logprobs': -0.0002904396352111077, 'logprobs_result': {'top_candidates': [], 'chosen_candidates': []}}, id='run-b2cf906b-d897-4164-b743-1120419ebdf9-0', tool_calls=[{'name': 'agent_1', 'args': {'question': 'Can you load the entityId of all data under HTAN.10xvisium_spatialtranscriptomics_scRNAseq_level4_metadata_current where the File_Format is hdf5?'}, 'id': '6987b536-4024-4b91-a0b1-720e41b1076e', 'type': 'tool_call'}], usage_metadata={'input_tokens': 536, 'output_tokens': 49, 'total_tokens': 585})], tool_call_id='6987b536-4024-4b91-a0b1-720e41b1076e'), 'None')]}\n",
        "    ---- Initial response captured ----\n",
        "    '''\n",
        "    Well formatted output:\n",
        "    '''\n",
        "    A total of 41 entity IDs were retrieved from the BigQuery table isb-cgc-bq.HTAN.10xvisium_spatialtranscriptomics_scRNAseq_level4_metadata_current where the File_Format is 'hdf5'.  The entity IDs are listed below::\n",
        "\n",
        "    1. syn51133519\n",
        "    2. syn51133520\n",
        "    3. syn51133521\n",
        "    4. syn51133522\n",
        "    5. syn51133523\n",
        "    6. syn51133524\n",
        "    7. syn51133525\n",
        "    8. syn51133526\n",
        "    9. syn51133527\n",
        "    10. syn51133528\n",
        "    11. syn51133529\n",
        "    12. syn51133530\n",
        "    13. syn51133531\n",
        "    14. syn51133532\n",
        "    15. syn51133533\n",
        "    16. syn51133534\n",
        "    17. syn51133537\n",
        "    18. syn51133540\n",
        "    19. syn51133578\n",
        "    20. syn51133580\n",
        "    21. syn51133583\n",
        "    22. syn51133587\n",
        "    23. syn51133591\n",
        "    24. syn51133592\n",
        "    25. syn51133593\n",
        "    26. syn51133595\n",
        "    27. syn51133596\n",
        "    28. syn51133597\n",
        "    29. syn51133598\n",
        "    30. syn51133599\n",
        "    31. syn51133600\n",
        "    32. syn51133601\n",
        "    33. syn51133602\n",
        "    34. syn51133603\n",
        "    35. syn51133604\n",
        "    36. syn51133605\n",
        "    37. syn51133606\n",
        "    38. syn51133607\n",
        "    39. syn51133608\n",
        "    40. syn51133609\n",
        "    41. syn51133612\n",
        "    '''\n",
        "    \"\"\"\n",
        "\n",
        "    prompt += f\"\"\"\n",
        "    Input: {question}\n",
        "    Please generate the well formated output.\n",
        "    \"\"\"\n",
        "\n",
        "    vertexai.init(project=\"isb-cgc-external-004\", location=\"us-central1\")\n",
        "    model = GenerativeModel(\n",
        "        \"gemini-1.5-flash-002\",\n",
        "    )\n",
        "    responses = model.generate_content(\n",
        "        [prompt],\n",
        "        generation_config=generation_config,\n",
        "        safety_settings=safety_settings,\n",
        "        stream=False,\n",
        "        # tools='code_execution'\n",
        "    )\n",
        "\n",
        "    # Example usage\n",
        "    display(Markdown(responses.text))\n",
        "\n",
        "llm = ChatVertexAI(model_name= 'gemini-1.5-flash-002')\n",
        "#define system prompt for tool calling agent\n",
        "system_prompt = \"\"\"\n",
        "You are a supervisor who can select the right worker to complete a task. The workers are called agents, and we have two agents, agent_1 and agent_2, each skilled in specific tasks, as described below:\n",
        "\n",
        "- Use `agent_1` to generate code for downloading files from Synapse or the HTAN database. If the user request involves data retrieval from an external source like Synapse, agent_1 is the correct choice.\n",
        "\n",
        "- Use `agent_2` to generate code for handling tasks related to the AnnData object, named `adata`, which is available in the local variables. Specifically:\n",
        "  - Select `agent_2` for tasks that involve loading, preprocessing, or analyzing the `adata` object.\n",
        "  - Use `agent_2` when the user’s question requires performing operations on `adata`, such as filtering, normalization, dimensionality reduction, or any other preprocessing step.\n",
        "  - For questions or analysis based on the data contained within `adata`, also use `agent_2`.\n",
        "\n",
        "If you do not have a tool to answer the question, inform the user accordingly.\n",
        "\n",
        "\"\"\"\n",
        "\n",
        "\n",
        "tool_calling_prompt = ChatPromptTemplate.from_messages(\n",
        "    [\n",
        "        (\"system\", system_prompt),\n",
        "        MessagesPlaceholder(\"chat_history\", optional=True),\n",
        "        (\"human\", \"{input}\"),\n",
        "        MessagesPlaceholder(\"agent_scratchpad\"),\n",
        "    ]\n",
        ")\n",
        "\n",
        "tool_runnable = create_tool_calling_agent(llm, toolkit, prompt  = tool_calling_prompt)\n",
        "\n",
        "\n",
        "def run_tool_agent(state):\n",
        "    agent_outcome = tool_runnable.invoke(state)\n",
        "\n",
        "    #this agent will overwrite the agent outcome state variable\n",
        "    return {\"agent_outcome\": agent_outcome}\n",
        "\n",
        "# tool executor invokes the tool action specified from the agent runnable\n",
        "# they will become the nodes that will be called when the agent decides on a tool action.\n",
        "\n",
        "tool_executor = ToolExecutor(toolkit)\n",
        "\n",
        "# Define the function to execute tools\n",
        "# This node will run a different tool as specified in the state variable agent_outcome\n",
        "def execute_tools(state):\n",
        "    # Get the most recent agent_outcome - this is the key added in the `agent` above\n",
        "    agent_action = state['agent_outcome']\n",
        "    if type(agent_action) is not list:\n",
        "        agent_action = [agent_action]\n",
        "    steps = []\n",
        "    #sca only returns an action while tool calling returns a list\n",
        "    # convert single actions to a list\n",
        "\n",
        "    for action in agent_action:\n",
        "    # Execute the tool\n",
        "        output = tool_executor.invoke(action)\n",
        "        print(f\"The agent action is {action}\")\n",
        "        print(f\"The tool result is: {output}\")\n",
        "        steps.append((action, str(output)))\n",
        "    # Return the output\n",
        "    return {\"intermediate_steps\": steps}\n",
        "\n",
        "class AgentState(TypedDict):\n",
        "   # The input string from human\n",
        "   input: str\n",
        "   # The list of previous messages in the conversation\n",
        "   chat_history: list[BaseMessage]\n",
        "   # The outcome of a given call to the agent\n",
        "   # Needs 'list' as a valid type as the tool agent returns a list.\n",
        "   # Needs `None` as a valid type, since this is what this will start as\n",
        "   # this state will be overwritten with the latest everytime the agent is run\n",
        "   agent_outcome: Union[AgentAction, list, ToolAgentAction, AgentFinish, None]\n",
        "\n",
        "   # List of actions and corresponding observations\n",
        "   # These actions should be added onto the existing so we use `operator.add`\n",
        "   # to append to the list of past intermediate steps\n",
        "   intermediate_steps: Annotated[list[Union[tuple[AgentAction, str], tuple[ToolAgentAction, str]]], operator.add]\n",
        "\n",
        "def should_continue(data):\n",
        "    # If the agent outcome is an AgentFinish, then we return `exit` string\n",
        "    # This will be used when setting up the graph to define the flow\n",
        "    if isinstance(data['agent_outcome'], AgentFinish):\n",
        "        return \"END\"\n",
        "    # Otherwise, an AgentAction is returned\n",
        "    # Here we return `continue` string\n",
        "    # This will be used when setting up the graph to define the flow\n",
        "    else:\n",
        "        return \"CONTINUE\"\n",
        "\n",
        "# Define a new graph\n",
        "workflow = StateGraph(AgentState)\n",
        "\n",
        "# When nodes are called, the functions for to the tools will be called.\n",
        "workflow.add_node(\"agent\", run_tool_agent)\n",
        "\n",
        "\n",
        "# Add tool invocation node to the graph\n",
        "workflow.add_node(\"action\", execute_tools)\n",
        "\n",
        "# Define which node the graph will invoke at start.\n",
        "workflow.set_entry_point(\"agent\")\n",
        "\n",
        "# Add flow logic with static edge.\n",
        "# Each time a tool is invoked and completed we want to\n",
        "# return the result to the agent to assess if task is complete or to take further actions\n",
        "\n",
        "#each action invocation has an edge leading to the agent node.\n",
        "workflow.add_edge('action', 'agent')\n",
        "\n",
        "\n",
        "# Add flow logic with conditional edge.\n",
        "workflow.add_conditional_edges(\n",
        "    # first parameter is the starting node for the edge\n",
        "    \"agent\",\n",
        "    # the second parameter specifies the logic function to be run\n",
        "    # to determine which node the edge will point to given the state.\n",
        "    should_continue,\n",
        "\n",
        "    #third parameter defines the mapping between the logic function\n",
        "    #output and the nodes on the graph\n",
        "    # For each possible output of the logic function there must be a valid node.\n",
        "    {\n",
        "        # If 'continue' we proceed to the action node.\n",
        "        \"CONTINUE\": \"action\",\n",
        "        # Otherwise we end invocations with the END node.\n",
        "        \"END\": END\n",
        "    }\n",
        ")\n",
        "\n",
        "memory = MemorySaver()\n",
        "\n",
        "# Finally, compile the graph!\n",
        "# This compiles it into a LangChain Runnable,\n",
        "app = workflow.compile(checkpointer = memory)\n",
        "\n",
        "def ask_agents(user_question, config={\"configurable\": {\"thread_id\": \"1\"}}):\n",
        "    \"\"\"\n",
        "    Executes an application stream with enhanced error handling and logs the output.\n",
        "\n",
        "    Parameters:\n",
        "    - user_question: A string representing the user's question or command.\n",
        "    - config: A dictionary with configuration settings for the app, such as 'thread_id'.\n",
        "\n",
        "    Returns:\n",
        "    - The captured standard output as a string.\n",
        "    \"\"\"\n",
        "    buffer = io.StringIO()  # Capture output buffer\n",
        "    inputs = {\"input\": user_question, \"chat_history\": []}  # Set up inputs for the stream\n",
        "\n",
        "    with warnings.catch_warnings():\n",
        "        warnings.simplefilter(\"ignore\")  # Ignore warnings during execution\n",
        "        # Redirect stdout to the buffer to capture printed output\n",
        "        with redirect_stdout(buffer):\n",
        "            try:\n",
        "                # Stream output and capture responses\n",
        "                for i, response_chunk in enumerate(app.stream(inputs, config=config)):\n",
        "                    result = list(response_chunk.values())[0]  # Accessing the actual response content\n",
        "                    print(result)  # Print result to buffer to capture it\n",
        "                    if i == 1:\n",
        "                        # Optionally break after the initial response\n",
        "                        print(\"---- Initial response captured ----\")\n",
        "                        break\n",
        "            except Exception as e:\n",
        "                # Print the error in the captured output\n",
        "                print(f\"Error during streaming execution: {e}\")\n",
        "\n",
        "    # Retrieve the contents of the buffer as a string\n",
        "    return buffer.getvalue()\n",
        "\n",
        "# Example usage to show HTAN data browsing\n",
        "user_question = \"Can you load the entityId of all data under HTAN.10xvisium_spatialtranscriptomics_scRNAseq_level4_metadata_current where the File_Format is hdf5?\"\n",
        "\n",
        "# Example usage to show HTAN data browsing\n",
        "# user_question = \"Can you download the Synapse datasaet with id syn51133612 to /content/demo_data/?\"\n",
        "\n",
        "# user_question = \"Can you load the /content/demo_data/8899_AS_8_unfiltered.h5ad to adata using scanpy?\"\n",
        "\n",
        "# user_question = \"Please preprocess the adata object\"\n",
        "\n",
        "# user_question = \"Generate a UMAP visualization comparing kmeans_9_clusters with kmeans_10_clusters\"\n",
        "\n",
        "output = ask_agents(user_question)"
      ],
      "metadata": {
        "id": "-kn6Mtcacpb7",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 17
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1731090744630,
          "user_tz": 300,
          "elapsed": 2766,
          "user": {
            "displayName": "",
            "userId": ""
          }
        },
        "outputId": "53889180-09bc-428e-b111-4dbfe9932be1",
        "cellView": "code"
      },
      "id": "-kn6Mtcacpb7",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "user_question = \"Can you load the /content/demo_data/8899_AS_8_unfiltered.h5ad to adata using scanpy?\"\n",
        "output = ask_agents(user_question)\n",
        "pritn(output)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 140
        },
        "id": "OIV9zuYQ-NCN",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1731090744631,
          "user_tz": 300,
          "elapsed": 3,
          "user": {
            "displayName": "",
            "userId": ""
          }
        },
        "outputId": "391b107d-745f-460a-89d1-e7e858331c38"
      },
      "id": "OIV9zuYQ-NCN",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'{\\'agent_outcome\\': [ToolAgentAction(tool=\\'agent_2\\', tool_input={\\'question\\': \\'Can you load the /content/demo_data/8899_AS_8_unfiltered.h5ad to adata using scanpy?\\'}, log=\"\\\\nInvoking: `agent_2` with `{\\'question\\': \\'Can you load the /content/demo_data/8899_AS_8_unfiltered.h5ad to adata using scanpy?\\'}`\\\\n\\\\n\\\\n\", message_log=[AIMessage(content=\\'\\', additional_kwargs={\\'function_call\\': {\\'name\\': \\'agent_2\\', \\'arguments\\': \\'{\"question\": \"Can you load the /content/demo_data/8899_AS_8_unfiltered.h5ad to adata using scanpy?\"}\\'}}, response_metadata={\\'is_blocked\\': False, \\'safety_ratings\\': [], \\'usage_metadata\\': {\\'prompt_token_count\\': 524, \\'candidates_token_count\\': 37, \\'total_token_count\\': 561, \\'cached_content_token_count\\': 0}, \\'finish_reason\\': \\'STOP\\', \\'avg_logprobs\\': -0.03173088705217516, \\'logprobs_result\\': {\\'top_candidates\\': [], \\'chosen_candidates\\': []}}, id=\\'run-0fe65e7f-1452-4d60-8c6d-e9452bf2ae02-0\\', tool_calls=[{\\'name\\': \\'agent_2\\', \\'args\\': {\\'question\\': \\'Can you load the /content/demo_data/8899_AS_8_unfiltered.h5ad to adata using scanpy?\\'}, \\'id\\': \\'64ab7d79-31b2-4409-9e45-523e0adec0a8\\', \\'type\\': \\'tool_call\\'}], usage_metadata={\\'input_tokens\\': 524, \\'output_tokens\\': 37, \\'total_tokens\\': 561})], tool_call_id=\\'64ab7d79-31b2-4409-9e45-523e0adec0a8\\')]}\\nGenerated code: import scanpy as sc\\nimport anndata as ad\\nimport os\\n\\n# Check if the file exists\\nfilepath = \"/content/demo_data/8899_AS_8_unfiltered.h5ad\"\\nif not os.path.exists(filepath):\\n    raise FileNotFoundError(f\"Error: File not found at {filepath}\")\\n\\ntry:\\n    # Load the .h5ad file into an AnnData object\\n    adata = sc.read_h5ad(filepath)\\n    print(f\"Successfully loaded {filepath} into adata.\")\\n    print(adata) # Print adata information to verify loading\\n\\nexcept Exception as e:\\n    print(f\"An error occurred while loading the file: {e}\")\\nSuccessfully loaded /content/demo_data/8899_AS_8_unfiltered.h5ad into adata.\\nAnnData object with n_obs × n_vars = 4992 × 18729\\n    obs: \\'in_tissue\\', \\'array_row\\', \\'array_col\\', \\'kmeans_7_clusters\\', \\'kmeans_10_clusters\\', \\'kmeans_4_clusters\\', \\'kmeans_2_clusters\\', \\'kmeans_6_clusters\\', \\'graphclust\\', \\'kmeans_3_clusters\\', \\'kmeans_8_clusters\\', \\'kmeans_9_clusters\\', \\'kmeans_5_clusters\\'\\n    var: \\'gene_ids\\', \\'feature_types\\', \\'genome\\', \\'n_cells\\', \\'Morans_I\\', \\'Morans_I_p_val\\', \\'Morans_I_adj_p_val\\', \\'Feature Counts in Spots Under Tissue\\', \\'Median Normalized Average Counts\\', \\'Barcodes Detected per Feature\\'\\n    uns: \\'spatial\\'\\n    obsm: \\'spatial\\'\\nAgent 2: Executed Python code with accessible variables.\\nThe agent action is tool=\\'agent_2\\' tool_input={\\'question\\': \\'Can you load the /content/demo_data/8899_AS_8_unfiltered.h5ad to adata using scanpy?\\'} log=\"\\\\nInvoking: `agent_2` with `{\\'question\\': \\'Can you load the /content/demo_data/8899_AS_8_unfiltered.h5ad to adata using scanpy?\\'}`\\\\n\\\\n\\\\n\" message_log=[AIMessage(content=\\'\\', additional_kwargs={\\'function_call\\': {\\'name\\': \\'agent_2\\', \\'arguments\\': \\'{\"question\": \"Can you load the /content/demo_data/8899_AS_8_unfiltered.h5ad to adata using scanpy?\"}\\'}}, response_metadata={\\'is_blocked\\': False, \\'safety_ratings\\': [], \\'usage_metadata\\': {\\'prompt_token_count\\': 524, \\'candidates_token_count\\': 37, \\'total_token_count\\': 561, \\'cached_content_token_count\\': 0}, \\'finish_reason\\': \\'STOP\\', \\'avg_logprobs\\': -0.03173088705217516, \\'logprobs_result\\': {\\'top_candidates\\': [], \\'chosen_candidates\\': []}}, id=\\'run-0fe65e7f-1452-4d60-8c6d-e9452bf2ae02-0\\', tool_calls=[{\\'name\\': \\'agent_2\\', \\'args\\': {\\'question\\': \\'Can you load the /content/demo_data/8899_AS_8_unfiltered.h5ad to adata using scanpy?\\'}, \\'id\\': \\'64ab7d79-31b2-4409-9e45-523e0adec0a8\\', \\'type\\': \\'tool_call\\'}], usage_metadata={\\'input_tokens\\': 524, \\'output_tokens\\': 37, \\'total_tokens\\': 561})] tool_call_id=\\'64ab7d79-31b2-4409-9e45-523e0adec0a8\\'\\nThe tool result is: None\\n{\\'intermediate_steps\\': [(ToolAgentAction(tool=\\'agent_2\\', tool_input={\\'question\\': \\'Can you load the /content/demo_data/8899_AS_8_unfiltered.h5ad to adata using scanpy?\\'}, log=\"\\\\nInvoking: `agent_2` with `{\\'question\\': \\'Can you load the /content/demo_data/8899_AS_8_unfiltered.h5ad to adata using scanpy?\\'}`\\\\n\\\\n\\\\n\", message_log=[AIMessage(content=\\'\\', additional_kwargs={\\'function_call\\': {\\'name\\': \\'agent_2\\', \\'arguments\\': \\'{\"question\": \"Can you load the /content/demo_data/8899_AS_8_unfiltered.h5ad to adata using scanpy?\"}\\'}}, response_metadata={\\'is_blocked\\': False, \\'safety_ratings\\': [], \\'usage_metadata\\': {\\'prompt_token_count\\': 524, \\'candidates_token_count\\': 37, \\'total_token_count\\': 561, \\'cached_content_token_count\\': 0}, \\'finish_reason\\': \\'STOP\\', \\'avg_logprobs\\': -0.03173088705217516, \\'logprobs_result\\': {\\'top_candidates\\': [], \\'chosen_candidates\\': []}}, id=\\'run-0fe65e7f-1452-4d60-8c6d-e9452bf2ae02-0\\', tool_calls=[{\\'name\\': \\'agent_2\\', \\'args\\': {\\'question\\': \\'Can you load the /content/demo_data/8899_AS_8_unfiltered.h5ad to adata using scanpy?\\'}, \\'id\\': \\'64ab7d79-31b2-4409-9e45-523e0adec0a8\\', \\'type\\': \\'tool_call\\'}], usage_metadata={\\'input_tokens\\': 524, \\'output_tokens\\': 37, \\'total_tokens\\': 561})], tool_call_id=\\'64ab7d79-31b2-4409-9e45-523e0adec0a8\\'), \\'None\\')]}\\n---- Initial response captured ----\\n'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 192
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "agent_3(output)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 225
        },
        "id": "vN-I7p8epfmP",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1731090748832,
          "user_tz": 300,
          "elapsed": 2162,
          "user": {
            "displayName": "",
            "userId": ""
          }
        },
        "outputId": "d81c50f9-9dcd-401f-ee4d-43e60a146927"
      },
      "id": "vN-I7p8epfmP",
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "\n",
              "  <style>\n",
              "    pre {\n",
              "        white-space: pre-wrap;\n",
              "    }\n",
              "  </style>\n",
              "  "
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.Markdown object>"
            ],
            "text/markdown": "```\nThe file `/content/demo_data/8899_AS_8_unfiltered.h5ad` was successfully loaded into the `adata` AnnData object using scanpy.  The object contains information for 4992 observations (cells or spots) and 18729 variables (genes).  Details are shown below:\n\nAnnData object with n_obs × n_vars = 4992 × 18729\n    obs: 'in_tissue', 'array_row', 'array_col', 'kmeans_7_clusters', 'kmeans_10_clusters', 'kmeans_4_clusters', 'kmeans_2_clusters', 'kmeans_6_clusters', 'graphclust', 'kmeans_3_clusters', 'kmeans_8_clusters', 'kmeans_9_clusters', 'kmeans_5_clusters'\n    var: 'gene_ids', 'feature_types', 'genome', 'n_cells', 'Morans_I', 'Morans_I_p_val', 'Morans_I_adj_p_val', 'Feature Counts in Spots Under Tissue', 'Median Normalized Average Counts', 'Barcodes Detected per Feature'\n    uns: 'spatial'\n    obsm: 'spatial'\n```\n"
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "LMAREIDSmxFd"
      },
      "id": "LMAREIDSmxFd",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "rA3onWxe-23B"
      },
      "id": "rA3onWxe-23B",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "-L-k7KgZ-2wo"
      },
      "id": "-L-k7KgZ-2wo",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "EUBE3Y-B-2oO"
      },
      "id": "EUBE3Y-B-2oO",
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.10.10"
    },
    "colab": {
      "provenance": [],
      "name": "app_v2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}