{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# Building a Number‑Guessing Agent with Koog\n",
    "\n",
    "Let’s build a small but fun agent that guesses a number you’re thinking of. We’ll lean on Koog’s tool-calling to ask targeted questions and converge using a classic binary search strategy. The result is an idiomatic Kotlin Notebook that you can drop straight into docs.\n",
    "\n",
    "We’ll keep the code minimal and the flow transparent: a few tiny tools, a compact prompt, and an interactive CLI loop."
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Setup\n",
    "\n",
    "This notebook assumes:\n",
    "- You’re running in a Kotlin Notebook with Koog available.\n",
    "- The environment variable `OPENAI_API_KEY` is set. The agent uses it via `simpleOpenAIExecutor(System.getenv(\"OPENAI_API_KEY\"))`.\n",
    "\n",
    "Load the Koog kernel:"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "%useLatestDescriptors\n",
    "%use koog"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Tools: asking targeted questions\n",
    "\n",
    "Tools are small, well-described functions the LLM can call. We’ll provide three:\n",
    "- `lessThan(value)`: “Is your number less than value?”\n",
    "- `greaterThan(value)`: “Is your number greater than value?”\n",
    "- `proposeNumber(value)`: “Is your number equal to value?” (used once the range is tight)\n",
    "\n",
    "Each tool returns a simple \"YES\"/\"NO\" string. The helper `ask` implements a minimal Y/n loop and validates input. Descriptions via `@LLMDescription` help the model select tools correctly."
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import ai.koog.agents.core.tools.annotations.Tool\n",
    "\n",
    "class GuesserTool : ToolSet {\n",
    "\n",
    "    @Tool\n",
    "    @LLMDescription(\"Asks the user if his number is STRICTLY less than a given value.\")\n",
    "    fun lessThan(\n",
    "        @LLMDescription(\"A value to compare the guessed number with.\") value: Int\n",
    "    ): String = ask(\"Is your number less than $value?\", value)\n",
    "\n",
    "    @Tool\n",
    "    @LLMDescription(\"Asks the user if his number is STRICTLY greater than a given value.\")\n",
    "    fun greaterThan(\n",
    "        @LLMDescription(\"A value to compare the guessed number with.\") value: Int\n",
    "    ): String = ask(\"Is your number greater than $value?\", value)\n",
    "\n",
    "    @Tool\n",
    "    @LLMDescription(\"Asks the user if his number is EXACTLY equal to the given number. Only use this tool once you've narrowed down your answer.\")\n",
    "    fun proposeNumber(\n",
    "        @LLMDescription(\"A value to compare the guessed number with.\") value: Int\n",
    "    ): String = ask(\"Is your number equal to $value?\", value)\n",
    "\n",
    "    fun ask(question: String, value: Int): String {\n",
    "        print(\"$question [Y/n]: \")\n",
    "        val input = readln()\n",
    "        println(input)\n",
    "\n",
    "        return when (input.lowercase()) {\n",
    "            \"\", \"y\", \"yes\" -> \"YES\"\n",
    "            \"n\", \"no\" -> \"NO\"\n",
    "            else -> {\n",
    "                println(\"Invalid input! Please, try again.\")\n",
    "                ask(question, value)\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Tool Registry\n",
    "\n",
    "Expose your tools to the agent. We also add a built‑in `SayToUser` tool so the agent can surface messages directly to the user."
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "val toolRegistry = ToolRegistry {\n",
    "    tool(SayToUser)\n",
    "    tools(GuesserTool())\n",
    "}"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Agent configuration\n",
    "\n",
    "A short, tool‑forward system prompt is all we need. We’ll suggest a binary search strategy and keep `temperature = 0.0` for stable, deterministic behavior. Here we use OpenAI’s reasoning model `GPT4oMini` for crisp planning."
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "val agent = AIAgent(\n",
    "    executor = simpleOpenAIExecutor(System.getenv(\"OPENAI_API_KEY\")),\n",
    "    llmModel = OpenAIModels.Reasoning.GPT4oMini,\n",
    "    systemPrompt = \"\"\"\n",
    "            You are a number guessing agent. Your goal is to guess a number that the user is thinking of.\n",
    "            \n",
    "            Follow these steps:\n",
    "            1. Start by asking the user to think of a number between 1 and 100.\n",
    "            2. Use the less_than and greater_than tools to narrow down the range.\n",
    "                a. If it's neither greater nor smaller, use the propose_number tool.\n",
    "            3. Once you're confident about the number, use the propose_number tool to check if your guess is correct.\n",
    "            4. If your guess is correct, congratulate the user. If not, continue guessing.\n",
    "            \n",
    "            Be efficient with your guessing strategy. A binary search approach works well.\n",
    "        \"\"\".trimIndent(),\n",
    "    temperature = 0.0,\n",
    "    toolRegistry = toolRegistry\n",
    ")"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Run it\n",
    "\n",
    "- Think of a number between 1 and 100.\n",
    "- Type `start` to begin.\n",
    "- Answer the agent’s questions with `Y`/`Enter` for yes or `n` for no. The agent should zero in on your number in ~7 steps."
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import kotlinx.coroutines.runBlocking\n",
    "\n",
    "println(\"Number Guessing Game started!\")\n",
    "println(\"Think of a number between 1 and 100, and I'll try to guess it.\")\n",
    "println(\"Type 'start' to begin the game.\")\n",
    "\n",
    "val initialMessage = readln()\n",
    "runBlocking {\n",
    "    agent.run(initialMessage)\n",
    "}"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## How it works\n",
    "\n",
    "- The agent reads the system prompt and plans a binary search.\n",
    "- On each iteration it calls one of your tools: `lessThan`, `greaterThan`, or (when certain) `proposeNumber`.\n",
    "- The helper `ask` collects your Y/n input and returns a clean \"YES\"/\"NO\" signal back to the model.\n",
    "- When it gets confirmation, it congratulates you via `SayToUser`."
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Extend it\n",
    "\n",
    "- Change the range (e.g., 1..1000) by tweaking the system prompt.\n",
    "- Add a `between(low, high)` tool to reduce calls further.\n",
    "- Swap models or executors (e.g., use an Ollama executor and a local model) while keeping the same tools.\n",
    "- Persist guesses or outcomes to a store for analytics."
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Troubleshooting\n",
    "\n",
    "- Missing key: ensure `OPENAI_API_KEY` is set in your environment.\n",
    "- Kernel not found: make sure `%useLatestDescriptors` and `%use koog` executed successfully.\n",
    "- Tool not called: confirm the `ToolRegistry` includes `GuesserTool()` and the names in the prompt match your tool functions."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Kotlin",
   "language": "kotlin",
   "name": "kotlin"
  },
  "language_info": {
   "name": "kotlin",
   "version": "2.2.20-Beta2",
   "mimetype": "text/x-kotlin",
   "file_extension": ".kt",
   "pygments_lexer": "kotlin",
   "codemirror_mode": "text/x-kotlin",
   "nbconvert_exporter": ""
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}