{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install encord encord-agents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SETUP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from google.colab import userdata"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## VPC/US Clients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"ENCORD_SSH_KEY\"] = userdata.get(\"encord_vpc_key\")\n",
    "DOMAIN = \"https://api.annotation.encord.ai\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## EMEA Clients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"ENCORD_SSH_KEY\"] = userdata.get(\"encord_key\")\n",
    "DOMAIN = \"https://api.encord.com\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from encord import EncordUserClient\n",
    "\n",
    "os.environ[\"ENCORD_DOMAIN\"] = DOMAIN\n",
    "\n",
    "# In the SECRETS tab on the left (key icon), add your Encord SSH key\n",
    "# You can get your Encord Key here: https://docs.encord.com/platform-documentation/Annotate/annotate-api-keys\n",
    "encord_client = EncordUserClient.create_with_ssh_private_key(ssh_private_key=userdata.get(\"encord_key\"), domain=DOMAIN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change this to be your Project Hash\n",
    "PROJECT_HASH = \"00000000-0000-0000-0000-000000000000\"\n",
    "STAGE_NAME = \"Custom Router\"\n",
    "\n",
    "project = encord_client.get_project(PROJECT_HASH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Filter Labels based on Object Attributes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ONTOLOGY_OBJECT_NAME = \"Bird\"\n",
    "ONTOLOGY_OBJECT_SUBFIELD_NAME = \"Type\"\n",
    "ATTRIBUTE_TO_FILTER_ON = \"Bluejay\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define ontology (& attributes you'd like)\n",
    "ontology = project.ontology_structure\n",
    "ontology_object = ontology.get_child_by_title(ONTOLOGY_OBJECT_NAME)\n",
    "question = ontology_object.get_child_by_title(ONTOLOGY_OBJECT_SUBFIELD_NAME)\n",
    "\n",
    "field_filter = ontology_object.get_child_by_title(ATTRIBUTE_TO_FILTER_ON)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Agent Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from encord_agents.core.utils import get_user_client\n",
    "from encord_agents.tasks import Depends, Runner\n",
    "\n",
    "runner = Runner(project_hash=PROJECT_HASH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from encord.objects import LabelRowV2\n",
    "from encord.workflow.stages.agent import AgentStage, AgentTask\n",
    "\n",
    "\n",
    "# main agent\n",
    "@runner.stage(STAGE_NAME, overwrite=True)\n",
    "def move_to_next_stage(lr: LabelRowV2):\n",
    "    for object in lr.get_object_instances(ontology_object):\n",
    "        try:\n",
    "            if object.get_answer(question).feature_node_hash == field_filter.feature_node_hash:\n",
    "                return \"Pass\"\n",
    "        except Exception:\n",
    "            continue\n",
    "    return \"Second Review\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the Agent\n",
    "# After 50 label updates, tasks are moved in the workflow.\n",
    "# Set to refresh every 3600 seconds to continuously check for new tasks\n",
    "runner(refresh_every=3600, task_batch_size=50, max_tasks_per_stage=None)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}