{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# Build a Simple Vacuum Cleaner Agent\n",
    "\n",
    "In this notebook, we'll explore how to implement a basic reflex agent using the new Kotlin agents framework.\n",
    "Our example will be the classic \"vacuum world\" problem —\n",
    "a simple environment with two locations that can be clean or dirty, and an agent that needs to clean them.\n",
    "\n",
    "First, let's understand our environment model:"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-06-19T09:00:53.703690Z",
     "start_time": "2025-06-19T09:00:53.462997Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import kotlin.random.Random\n",
    "\n",
    "/**\n",
    " * Represents a simple vacuum world with two locations (A and B).\n",
    " *\n",
    " * The environment tracks:\n",
    " * - The current location of the vacuum agent ('A' or 'B')\n",
    " * - The cleanliness status of each location (true = dirty, false = clean)\n",
    " */\n",
    "class VacuumEnv {\n",
    "    var location: Char = 'A'\n",
    "        private set\n",
    "\n",
    "    private val status = mutableMapOf(\n",
    "        'A' to Random.nextBoolean(),\n",
    "        'B' to Random.nextBoolean()\n",
    "    )\n",
    "\n",
    "    fun percept(): Pair<Char, Boolean> = location to status.getValue(location)\n",
    "\n",
    "    fun clean(): String {\n",
    "        status[location] = false\n",
    "        return \"cleaned\"\n",
    "    }\n",
    "\n",
    "    fun moveLeft(): String {\n",
    "        location = 'A'\n",
    "        return \"move to A\"\n",
    "    }\n",
    "\n",
    "    fun moveRight(): String {\n",
    "        location = 'B'\n",
    "        return \"move to B\"\n",
    "    }\n",
    "\n",
    "    fun isClean(): Boolean = status.values.all { it }\n",
    "\n",
    "    fun worldLayout(): String = \"${status.keys}\"\n",
    "\n",
    "    override fun toString(): String = \"location=$location, dirtyA=${status['A']}, dirtyB=${status['B']}\"\n",
    "}"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "The VacuumEnv class models our simple world:\n",
    "- Two locations are represented by characters 'A' and 'B'\n",
    "- Each location can be either clean or dirty (randomly initialized)\n",
    "- The agent can be at either location at any given time\n",
    "- The agent can perceive its current location and whether it's dirty\n",
    "- The agent can take actions: move to a specific location or clean the current location\n",
    "\n",
    "## Creating Tools for Vacuum Agent\n",
    "Now, let's define the tools our AI agent will use to interact with the environment:"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-06-19T09:00:53.780156Z",
     "start_time": "2025-06-19T09:00:53.707889Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import ai.koog.agents.core.tools.annotations.LLMDescription\n",
    "import ai.koog.agents.core.tools.annotations.Tool\n",
    "import ai.koog.agents.core.tools.reflect.ToolSet\n",
    "\n",
    "\n",
    "/**\n",
    " * Provides tools for the LLM agent to control the vacuum robot.\n",
    " * All methods either mutate or read from the VacuumEnv passed to the constructor.\n",
    " */\n",
    "@LLMDescription(\"Tools for controlling a two-cell vacuum world\")\n",
    "class VacuumTools(private val env: VacuumEnv) : ToolSet {\n",
    "\n",
    "    @Tool\n",
    "    @LLMDescription(\"Returns current location and whether it is dirty\")\n",
    "    fun sense(): String {\n",
    "        val (loc, dirty) = env.percept()\n",
    "        return \"location=$loc, dirty=$dirty, locations=${env.worldLayout()}\"\n",
    "    }\n",
    "\n",
    "    @Tool\n",
    "    @LLMDescription(\"Cleans the current cell\")\n",
    "    fun clean(): String = env.clean()\n",
    "\n",
    "    @Tool\n",
    "    @LLMDescription(\"Moves the agent to cell A\")\n",
    "    fun moveLeft(): String = env.moveLeft()\n",
    "\n",
    "    @Tool\n",
    "    @LLMDescription(\"Moves the agent to cell B\")\n",
    "    fun moveRight(): String = env.moveRight()\n",
    "}"
   ],
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "The `VacuumTools` class creates an interface between our LLM agent and the environment:\n",
    "\n",
    "- It implements `ToolSet` from the Kotlin AI Agents framework\n",
    "- Each tool is annotated with `@Tool` and has a description for the LLM\n",
    "- The tools allow the agent to sense its environment and take actions\n",
    "- Each method returns a string that describes the outcome of the action\n",
    "\n",
    "## Setting Up the Agent\n",
    "Next, we'll configure and create our AI agent:"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-06-19T09:00:54.113519Z",
     "start_time": "2025-06-19T09:00:53.810674Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import ai.koog.agents.core.agent.AIAgent\n",
    "import ai.koog.agents.core.agent.config.AIAgentConfig\n",
    "import ai.koog.agents.core.tools.ToolRegistry\n",
    "import ai.koog.agents.core.tools.reflect.asTools\n",
    "import ai.koog.agents.ext.agent.chatAgentStrategy\n",
    "import ai.koog.agents.ext.tool.AskUser\n",
    "import ai.koog.agents.ext.tool.SayToUser\n",
    "import ai.koog.prompt.dsl.prompt\n",
    "import ai.koog.prompt.executor.clients.openai.OpenAIModels\n",
    "import ai.koog.prompt.executor.llms.all.simpleOpenAIExecutor\n",
    "import ai.koog.prompt.params.LLMParams\n",
    "\n",
    "\n",
    "val env = VacuumEnv()\n",
    "val apiToken = System.getenv(\"OPENAI_API_KEY\") ?: error(\"OPENAI_API_KEY environment variable not set\")\n",
    "val executor = simpleOpenAIExecutor(apiToken = apiToken)\n",
    "\n",
    "val toolRegistry = ToolRegistry {\n",
    "    tool(SayToUser)\n",
    "    tool(AskUser)\n",
    "    tools(VacuumTools(env).asTools())\n",
    "}\n",
    "\n",
    "val systemVacuumPrompt = \"\"\"\n",
    "    You are a reflex vacuum-cleaner agent living in a two-cell world labelled A and B.\n",
    "    Your goal: make both cells clean, using the provided tools.\n",
    "    First, call sense() to inspect where you are. Then decide: if dirty → clean(); else moveLeft()/moveRight().\n",
    "    Continue until both cells are clean, then tell the user \"done\".\n",
    "    Use sayToUser to inform the user about each step.\n",
    "\"\"\".trimIndent()\n",
    "\n",
    "val agentConfig = AIAgentConfig(\n",
    "    prompt = prompt(\"chat\", params = LLMParams(temperature = 1.0)) {\n",
    "        system(systemVacuumPrompt)\n",
    "    },\n",
    "    model = OpenAIModels.Chat.GPT4o,\n",
    "    maxAgentIterations = 50,\n",
    ")\n",
    "\n",
    "val agent = AIAgent(\n",
    "    promptExecutor = executor,\n",
    "    strategy = chatAgentStrategy(),\n",
    "    agentConfig = agentConfig,\n",
    "    toolRegistry = toolRegistry\n",
    ")"
   ],
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "In this setup:\n",
    "\n",
    "1. We create an instance of our environment\n",
    "2. We set up a connection to OpenAI's GPT-4o model\n",
    "3. We register the tools our agent can use\n",
    "4. We define a system prompt that gives the agent its goal and behavior rules\n",
    "5. We create the agent using the `AIAgent` constructor with a chat strategy\n",
    "\n",
    "## Running the Agent\n",
    "\n",
    "Finally, let's run our agent:"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-06-19T09:01:00.816892Z",
     "start_time": "2025-06-19T09:00:54.119810Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import kotlinx.coroutines.runBlocking\n",
    "\n",
    "runBlocking {\n",
    "    agent.run(\"Start cleaning, please\")\n",
    "}"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Agent says: Currently in cell A. It's already clean.\n",
      "Agent says: Moved to cell B. It's already clean.\n"
     ]
    }
   ],
   "execution_count": 4
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "When we run this code:\n",
    "\n",
    "1. The agent receives the initial prompt to start cleaning\n",
    "2. It uses its tools to sense the environment and make decisions\n",
    "3. It continues cleaning until both cells are clean\n",
    "4. Throughout the process, it keeps the user informed about what it's doing\n"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-06-19T09:01:00.854101Z",
     "start_time": "2025-06-19T09:01:00.822551Z"
    }
   },
   "cell_type": "code",
   "source": [
    "// Finally we can validate that the work is finished by printing the env state\n",
    "\n",
    "env"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "location=B, dirtyA=false, dirtyB=false"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 5
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Kotlin",
   "language": "kotlin",
   "name": "kotlin"
  },
  "language_info": {
   "name": "kotlin",
   "version": "1.9.23",
   "mimetype": "text/x-kotlin",
   "file_extension": ".kt",
   "pygments_lexer": "kotlin",
   "codemirror_mode": "text/x-kotlin",
   "nbconvert_exporter": ""
  },
  "ktnbPluginMetadata": {
   "projectDependencies": [
    "koog-agents.examples.main"
   ],
   "projectLibraries": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}