{
 "cells": [
  {
   "cell_type": "markdown",
   "source": "# Building an AI Chess Player with Koog Framework\n\nThis tutorial demonstrates how to build an intelligent chess-playing agent using the Koog framework. We'll explore key concepts including tool integration, agent strategies, memory optimization, and interactive AI decision-making.\n\n## What You'll Learn\n\n- How to model domain-specific data structures for complex games\n- Creating custom tools that agents can use to interact with the environment\n- Implementing efficient agent strategies with memory management\n- Building interactive AI systems with choice selection capabilities\n- Optimizing agent performance for turn-based games\n\n## Setup\n\nFirst, let's import the Koog framework and set up our development environment:",
   "metadata": {}
  },
  {
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-08-18T21:15:01.940208Z",
     "start_time": "2025-08-18T21:15:01.028852Z"
    }
   },
   "cell_type": "code",
   "source": "%useLatestDescriptors\n%use koog",
   "outputs": [],
   "execution_count": 1
  },
  {
   "cell_type": "markdown",
   "source": "## Modeling the Chess Domain\n\nCreating a robust domain model is essential for any game AI. In chess, we need to represent players, pieces, and their relationships. Let's start by defining our core data structures:\n\n### Core Enums and Types",
   "metadata": {}
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:15:03.785575Z",
     "start_time": "2025-08-18T21:15:03.493469Z"
    }
   },
   "cell_type": "code",
   "source": "enum class Player {\n    White, Black, None;\n\n    fun opponent(): Player = when (this) {\n        White -> Black\n        Black -> White\n        None -> throw IllegalArgumentException(\"No opponent for None player\")\n    }\n}\n\nenum class PieceType(val id: Char) {\n    King('K'), Queen('Q'), Rook('R'),\n    Bishop('B'), Knight('N'), Pawn('P'), None('*');\n\n    companion object {\n        fun fromId(id: String): PieceType {\n            require(id.length == 1) { \"Invalid piece id: $id\" }\n\n            return entries.first { it.id == id.single() }\n        }\n    }\n}\n\nenum class Side {\n    King, Queen\n}",
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "The `Player` enum represents the two sides in chess, with an `opponent()` method for easy switching between players. The `PieceType` enum maps each chess piece to its standard notation character, enabling easy parsing of chess moves.\n",
    "\n",
    "The `Side` enum helps distinguish between kingside and queenside castling moves.\n",
    "\n",
    "### Piece and Position Modeling"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:15:07.432666Z",
     "start_time": "2025-08-18T21:15:07.111867Z"
    }
   },
   "cell_type": "code",
   "source": "data class Piece(val pieceType: PieceType, val player: Player) {\n    init {\n        require((pieceType == PieceType.None) == (player == Player.None)) {\n            \"Invalid piece: $pieceType $player\"\n        }\n    }\n\n    fun toChar(): Char = when (player) {\n        Player.White -> pieceType.id.uppercaseChar()\n        Player.Black -> pieceType.id.lowercaseChar()\n        Player.None -> pieceType.id\n    }\n\n    fun isNone(): Boolean = pieceType == PieceType.None\n\n    companion object {\n        val None = Piece(PieceType.None, Player.None)\n    }\n}\n\ndata class Position(val row: Int, val col: Char) {\n    init {\n        require(row in 1..8 && col in 'a'..'h') { \"Invalid position: $col$row\" }\n    }\n\n    constructor(position: String) : this(\n        position[1].digitToIntOrNull() ?: throw IllegalArgumentException(\"Incorrect position: $position\"),\n        position[0],\n    ) {\n        require(position.length == 2) { \"Invalid position: $position\" }\n    }\n}\n\nclass ChessBoard {\n    private val backRow = listOf(\n        PieceType.Rook, PieceType.Knight, PieceType.Bishop,\n        PieceType.Queen, PieceType.King,\n        PieceType.Bishop, PieceType.Knight, PieceType.Rook\n    )\n\n    private val board: List<MutableList<Piece>> = listOf(\n        backRow.map { Piece(it, Player.Black) }.toMutableList(),\n        List(8) { Piece(PieceType.Pawn, Player.Black) }.toMutableList(),\n        List(8) { Piece.None }.toMutableList(),\n        List(8) { Piece.None }.toMutableList(),\n        List(8) { Piece.None }.toMutableList(),\n        List(8) { Piece.None }.toMutableList(),\n        List(8) { Piece(PieceType.Pawn, Player.White) }.toMutableList(),\n        backRow.map { Piece(it, Player.White) }.toMutableList()\n    )\n\n    override fun toString(): String = board\n        .withIndex().joinToString(\"\\n\") { (index, row) ->\n            \"${8 - index} ${row.map { it.toChar() }.joinToString(\" \")}\"\n        } + \"\\n  a b c d e f g h\"\n\n    fun getPiece(position: Position): Piece = board[8 - position.row][position.col - 'a']\n    fun setPiece(position: Position, piece: Piece) {\n        board[8 - position.row][position.col - 'a'] = piece\n    }\n}",
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "The `Piece` data class combines a piece type with its owner, using uppercase letters for white pieces and lowercase for black pieces in the visual representation. The `Position` class encapsulates chess coordinates (e.g., \"e4\") with built-in validation.\n",
    "\n",
    "## Game State Management\n",
    "\n",
    "### ChessBoard Implementation\n",
    "\n",
    "The `ChessBoard` class manages the 8×8 grid and piece positions. Key design decisions include:\n",
    "\n",
    "- **Internal Representation**: Uses a list of mutable lists for efficient access and modification\n",
    "- **Visual Display**: The `toString()` method provides a clear ASCII representation with rank numbers and file letters\n",
    "- **Position Mapping**: Converts between chess notation (a1-h8) and internal array indices\n",
    "\n",
    "### ChessGame Logic"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:15:20.075419Z",
     "start_time": "2025-08-18T21:15:19.918855Z"
    }
   },
   "cell_type": "code",
   "source": "/**\n * Simple chess game without checks for valid moves.\n * Stores a correct state of the board if the entered moves are valid\n */\nclass ChessGame {\n    private val board: ChessBoard = ChessBoard()\n    private var currentPlayer: Player = Player.White\n    val moveNotation: String = \"\"\"\n        0-0 - short castle\n        0-0-0 - long castle\n        <piece>-<from>-<to> - usual move. e.g. p-e2-e4\n        <piece>-<from>-<to>-<promotion> - promotion move. e.g. p-e7-e8-q.\n        Piece names:\n            p - pawn\n            n - knight\n            b - bishop\n            r - rook\n            q - queen\n            k - king\n    \"\"\".trimIndent()\n\n    fun move(move: String) {\n        when {\n            move == \"0-0\" -> castleMove(Side.King)\n            move == \"0-0-0\" -> castleMove(Side.Queen)\n            move.split(\"-\").size == 3 -> {\n                val (_, from, to) = move.split(\"-\")\n                usualMove(Position(from), Position(to))\n            }\n\n            move.split(\"-\").size == 4 -> {\n                val (piece, from, to, promotion) = move.split(\"-\")\n\n                require(PieceType.fromId(piece) == PieceType.Pawn) { \"Only pawn can be promoted\" }\n\n                usualMove(Position(from), Position(to))\n                board.setPiece(Position(to), Piece(PieceType.fromId(promotion), currentPlayer))\n            }\n\n            else -> throw IllegalArgumentException(\"Invalid move: $move\")\n        }\n\n        updateCurrentPlayer()\n    }\n\n    fun getBoard(): String = board.toString()\n    fun currentPlayer(): String = currentPlayer.name.lowercase()\n\n    private fun updateCurrentPlayer() {\n        currentPlayer = currentPlayer.opponent()\n    }\n\n    private fun usualMove(from: Position, to: Position) {\n        if (board.getPiece(from).pieceType == PieceType.Pawn && from.col != to.col && board.getPiece(to).isNone()) {\n            // the move is en passant\n            board.setPiece(Position(from.row, to.col), Piece.None)\n        }\n\n        movePiece(from, to)\n    }\n\n    private fun castleMove(side: Side) {\n        val row = if (currentPlayer == Player.White) 1 else 8\n        val kingFrom = Position(row, 'e')\n        val (rookFrom, kingTo, rookTo) = if (side == Side.King) {\n            Triple(Position(row, 'h'), Position(row, 'g'), Position(row, 'f'))\n        } else {\n            Triple(Position(row, 'a'), Position(row, 'c'), Position(row, 'd'))\n        }\n\n        movePiece(kingFrom, kingTo)\n        movePiece(rookFrom, rookTo)\n    }\n\n    private fun movePiece(from: Position, to: Position) {\n        board.setPiece(to, board.getPiece(from))\n        board.setPiece(from, Piece.None)\n    }\n}",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "The `ChessGame` class orchestrates the game logic and maintains state. Notable features include:\n",
    "\n",
    "- **Move Notation Support**: Accepts standard chess notation for regular moves, castling (0-0, 0-0-0), and pawn promotion\n",
    "- **Special Move Handling**: Implements en passant capture and castling logic\n",
    "- **Turn Management**: Automatically alternates between players after each move\n",
    "- **Validation**: While it doesn't validate move legality (trusting the AI to make valid moves), it handles move parsing and state updates correctly\n",
    "\n",
    "The `moveNotation` string provides clear documentation for the AI agent on acceptable move formats.\n",
    "\n",
    "## Integrating with Koog Framework\n",
    "\n",
    "### Creating Custom Tools"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:15:24.784111Z",
     "start_time": "2025-08-18T21:15:24.621251Z"
    }
   },
   "cell_type": "code",
   "source": "import kotlinx.serialization.Serializable\n\nclass Move(val game: ChessGame) : SimpleTool<Move.Args>() {\n    @Serializable\n    data class Args(val notation: String) : ToolArgs\n\n    override val argsSerializer = Args.serializer()\n\n    override val descriptor = ToolDescriptor(\n        name = \"move\",\n        description = \"Moves a piece according to the notation:\\n${game.moveNotation}\",\n        requiredParameters = listOf(\n            ToolParameterDescriptor(\n                name = \"notation\",\n                description = \"The notation of the piece to move\",\n                type = ToolParameterType.String,\n            )\n        )\n    )\n\n    override suspend fun doExecute(args: Args): String {\n        game.move(args.notation)\n        println(game.getBoard())\n        println(\"-----------------\")\n        return \"Current state of the game:\\n${game.getBoard()}\\n${game.currentPlayer()} to move! Make the move!\"\n    }\n}",
   "outputs": [],
   "execution_count": 5
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "The `Move` tool demonstrates the Koog framework's tool integration pattern:\n",
    "\n",
    "1. **Extends SimpleTool**: Inherits the basic tool functionality with type-safe argument handling\n",
    "2. **Serializable Arguments**: Uses Kotlin serialization to define the tool's input parameters\n",
    "3. **Rich Documentation**: The `ToolDescriptor` provides the LLM with detailed information about the tool's purpose and parameters\n",
    "4. **Execution Logic**: The `doExecute` method handles the actual move execution and provides formatted feedback\n",
    "\n",
    "Key design aspects:\n",
    "- **Context Injection**: The tool receives the `ChessGame` instance, allowing it to modify game state\n",
    "- **Feedback Loop**: Returns the current board state and prompts the next player, maintaining conversational flow\n",
    "- **Error Handling**: Relies on the game class for move validation and error reporting\n",
    "\n",
    "## Agent Strategy Design\n",
    "\n",
    "### Memory Optimization Technique"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:15:30.111002Z",
     "start_time": "2025-08-18T21:15:29.891773Z"
    }
   },
   "cell_type": "code",
   "source": "import ai.koog.agents.core.environment.ReceivedToolResult\n\n/**\n * Chess position is (almost) completely defined by the board state,\n * So we can trim the history of the LLM to only contain the system prompt and the last move.\n */\ninline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.nodeTrimHistory(\n    name: String? = null\n): AIAgentNodeDelegate<T, T> = node(name) { result ->\n    llm.writeSession {\n        rewritePrompt { prompt ->\n            val messages = prompt.messages\n\n            prompt.copy(messages = listOf(messages.first(), messages.last()))\n        }\n    }\n\n    result\n}\n\nval strategy = strategy<String, String>(\"chess_strategy\") {\n    val nodeCallLLM by nodeLLMRequest(\"sendInput\")\n    val nodeExecuteTool by nodeExecuteTool(\"nodeExecuteTool\")\n    val nodeSendToolResult by nodeLLMSendToolResult(\"nodeSendToolResult\")\n    val nodeTrimHistory by nodeTrimHistory<ReceivedToolResult>()\n\n    edge(nodeStart forwardTo nodeCallLLM)\n    edge(nodeCallLLM forwardTo nodeExecuteTool onToolCall { true })\n    edge(nodeCallLLM forwardTo nodeFinish onAssistantMessage { true })\n    edge(nodeExecuteTool forwardTo nodeTrimHistory)\n    edge(nodeTrimHistory forwardTo nodeSendToolResult)\n    edge(nodeSendToolResult forwardTo nodeFinish onAssistantMessage { true })\n    edge(nodeSendToolResult forwardTo nodeExecuteTool onToolCall { true })\n}",
   "outputs": [],
   "execution_count": 6
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "The `nodeTrimHistory` function implements a crucial optimization for chess games. Since chess positions are largely determined by the current board state rather than the full move history, we can significantly reduce token usage by keeping only:\n",
    "\n",
    "1. **System Prompt**: Contains the agent's core instructions and behavior guidelines\n",
    "2. **Latest Message**: The most recent board state and game context\n",
    "\n",
    "This approach:\n",
    "- **Reduces Token Consumption**: Prevents exponential growth of conversation history\n",
    "- **Maintains Context**: Preserves essential game state information\n",
    "- **Improves Performance**: Faster processing with shorter prompts\n",
    "- **Enables Long Games**: Allows for extended gameplay without hitting token limits"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "The chess strategy demonstrates Koog's graph-based agent architecture:\n",
    "\n",
    "**Node Types:**\n",
    "- `nodeCallLLM`: Processes input and generates responses/tool calls\n",
    "- `nodeExecuteTool`: Executes the Move tool with the provided parameters\n",
    "- `nodeTrimHistory`: Optimizes conversation memory as described above\n",
    "- `nodeSendToolResult`: Sends tool execution results back to the LLM\n",
    "\n",
    "**Control Flow:**\n",
    "- **Linear Path**: Start → LLM Request → Tool Execution → History Trim → Send Result\n",
    "- **Decision Points**: LLM responses can either finish the conversation or trigger another tool call\n",
    "- **Memory Management**: History trimming occurs after each tool execution\n",
    "\n",
    "This strategy ensures efficient, stateful gameplay while maintaining conversational coherence.\n",
    "\n",
    "### Setting up the AI Agent"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:15:59.583574Z",
     "start_time": "2025-08-18T21:15:59.503444Z"
    }
   },
   "cell_type": "code",
   "source": "val baseExecutor = simpleOpenAIExecutor(System.getenv(\"OPENAI_API_KEY\"))",
   "outputs": [],
   "execution_count": 7
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "This section initializes our OpenAI executor. The `simpleOpenAIExecutor` creates a connection to OpenAI's API using your API key from environment variables.\n",
    "\n",
    "**Configuration Notes:**\n",
    "- Store your OpenAI API key in the `OPENAI_API_KEY` environment variable\n",
    "- The executor handles authentication and API communication automatically\n",
    "- Different executor types are available for various LLM providers\n",
    "\n",
    "### Agent Assembly"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:16:05.795181Z",
     "start_time": "2025-08-18T21:16:05.727710Z"
    }
   },
   "cell_type": "code",
   "source": "val game = ChessGame()\nval toolRegistry = ToolRegistry { tools(listOf(Move(game))) }\n\n// Create a chat agent with a system prompt and the tool registry\nval agent = AIAgent(\n    executor = baseExecutor,\n    strategy = strategy,\n    llmModel = OpenAIModels.Reasoning.O3Mini,\n    systemPrompt = \"\"\"\n            You are an agent who plays chess.\n            You should always propose a move in response to the \"Your move!\" message.\n\n            DO NOT HALLUCINATE!!!\n            DO NOT PLAY ILLEGAL MOVES!!!\n            YOU CAN SEND A MESSAGE ONLY IF IT IS A RESIGNATION OR A CHECKMATE!!!\n        \"\"\".trimMargin(),\n    temperature = 0.0,\n    toolRegistry = toolRegistry,\n    maxIterations = 200,\n)",
   "outputs": [],
   "execution_count": 8
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "Here we assemble all components into a functional chess-playing agent:\n",
    "\n",
    "**Key Configuration:**\n",
    "- **Model Choice**: Using `OpenAIModels.Reasoning.O3Mini` for high-quality chess play\n",
    "- **Temperature**: Set to 0.0 for deterministic, strategic moves\n",
    "- **System Prompt**: Carefully crafted instructions emphasizing legal moves and proper behavior\n",
    "- **Tool Registry**: Provides the agent access to the Move tool\n",
    "- **Max Iterations**: Set to 200 to allow for complete games\n",
    "\n",
    "**System Prompt Design:**\n",
    "- Emphasizes move proposal responsibility\n",
    "- Prohibits hallucination and illegal moves\n",
    "- Restricts messaging to only resignations or checkmate declarations\n",
    "- Creates focused, game-oriented behavior\n",
    "\n",
    "### Running the Basic Agent"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:16:33.014845Z",
     "start_time": "2025-08-18T21:16:12.019504Z"
    }
   },
   "cell_type": "code",
   "source": "import kotlinx.coroutines.runBlocking\n\nprintln(\"Chess Game started!\")\n\nval initialMessage = \"Starting position is ${game.getBoard()}. White to move!\"\n\nrunBlocking {\n    agent.run(initialMessage)\n}",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Chess Game started!\n",
      "8 r n b q k b n r\n",
      "7 p p p p p p p p\n",
      "6 * * * * * * * *\n",
      "5 * * * * * * * *\n",
      "4 * * * * P * * *\n",
      "3 * * * * * * * *\n",
      "2 P P P P * P P P\n",
      "1 R N B Q K B N R\n",
      "  a b c d e f g h\n",
      "-----------------\n",
      "8 r n b q k b n r\n",
      "7 p p p p * p p p\n",
      "6 * * * * * * * *\n",
      "5 * * * * p * * *\n",
      "4 * * * * P * * *\n",
      "3 * * * * * * * *\n",
      "2 P P P P * P P P\n",
      "1 R N B Q K B N R\n",
      "  a b c d e f g h\n",
      "-----------------\n",
      "8 r n b q k b n r\n",
      "7 p p p p * p p p\n",
      "6 * * * * * * * *\n",
      "5 * * * * p * * *\n",
      "4 * * * * P * * *\n",
      "3 * * * * * N * *\n",
      "2 P P P P * P P P\n",
      "1 R N B Q K B * R\n",
      "  a b c d e f g h\n",
      "-----------------\n",
      "8 r n b q k b * r\n",
      "7 p p p p * p p p\n",
      "6 * * * * * n * *\n",
      "5 * * * * p * * *\n",
      "4 * * * * P * * *\n",
      "3 * * * * * N * *\n",
      "2 P P P P * P P P\n",
      "1 R N B Q K B * R\n",
      "  a b c d e f g h\n",
      "-----------------\n",
      "8 r n b q k b * r\n",
      "7 p p p p * p p p\n",
      "6 * * * * * n * *\n",
      "5 * * * * p * * *\n",
      "4 * * * * P * * *\n",
      "3 * * N * * N * *\n",
      "2 P P P P * P P P\n",
      "1 R * B Q K B * R\n",
      "  a b c d e f g h\n",
      "-----------------\n"
     ]
    },
    {
     "ename": "org.jetbrains.kotlinx.jupyter.exceptions.ReplInterruptedException",
     "evalue": "The execution was interrupted",
     "output_type": "error",
     "traceback": [
      "The execution was interrupted"
     ]
    }
   ],
   "execution_count": 9
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "This basic agent plays autonomously, making moves automatically. The game output shows the sequence of moves and board states as the AI plays against itself.\n",
    "\n",
    "## Advanced Feature: Interactive Choice Selection\n",
    "\n",
    "The next sections demonstrate a more sophisticated approach where users can participate in the AI's decision-making process by choosing from multiple AI-generated moves.\n",
    "\n",
    "### Custom Choice Selection Strategy"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:16:37.788531Z",
     "start_time": "2025-08-18T21:16:37.671267Z"
    }
   },
   "cell_type": "code",
   "source": "import ai.koog.agents.core.feature.choice.ChoiceSelectionStrategy\n\n/**\n * `AskUserChoiceStrategy` allows users to interactively select a choice from a list of options\n * presented by a language model. The strategy uses customizable methods to display the prompt\n * and choices and read user input to determine the selected choice.\n *\n * @property promptShowToUser A function that formats and displays a given `Prompt` to the user.\n * @property choiceShowToUser A function that formats and represents a given `LLMChoice` to the user.\n * @property print A function responsible for displaying messages to the user, e.g., for showing prompts or feedback.\n * @property read A function to capture user input.\n */\nclass AskUserChoiceSelectionStrategy(\n    private val promptShowToUser: (Prompt) -> String = { \"Current prompt: $it\" },\n    private val choiceShowToUser: (LLMChoice) -> String = { \"$it\" },\n    private val print: (String) -> Unit = ::println,\n    private val read: () -> String? = ::readlnOrNull\n) : ChoiceSelectionStrategy {\n    override suspend fun choose(prompt: Prompt, choices: List<LLMChoice>): LLMChoice {\n        print(promptShowToUser(prompt))\n\n        print(\"Available LLM choices\")\n\n        choices.withIndex().forEach { (index, choice) ->\n            print(\"Choice number ${index + 1}: ${choiceShowToUser(choice)}\")\n        }\n\n        var choiceNumber = ask(choices.size)\n        while (choiceNumber == null) {\n            print(\"Invalid response.\")\n            choiceNumber = ask(choices.size)\n        }\n\n        return choices[choiceNumber - 1]\n    }\n\n    private fun ask(numChoices: Int): Int? {\n        print(\"Please choose a choice. Enter a number between 1 and $numChoices: \")\n\n        return read()?.toIntOrNull()?.takeIf { it in 1..numChoices }\n    }\n}",
   "outputs": [],
   "execution_count": 10
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "The `AskUserChoiceSelectionStrategy` implements Koog's `ChoiceSelectionStrategy` interface to enable human participation in AI decision-making:\n",
    "\n",
    "**Key Features:**\n",
    "- **Customizable Display**: Functions for formatting prompts and choices\n",
    "- **Interactive Input**: Uses standard input/output for user interaction\n",
    "- **Validation**: Ensures user input is within valid range\n",
    "- **Flexible I/O**: Configurable print and read functions for different environments\n",
    "\n",
    "**Use Cases:**\n",
    "- Human-AI collaboration in gameplay\n",
    "- AI decision transparency and explainability\n",
    "- Training and debugging scenarios\n",
    "- Educational demonstrations\n",
    "\n",
    "### Enhanced Strategy with Choice Selection"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:16:45.596851Z",
     "start_time": "2025-08-18T21:16:45.453678Z"
    }
   },
   "cell_type": "code",
   "source": "inline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.nodeTrimHistory(\n    name: String? = null\n): AIAgentNodeDelegate<T, T> = node(name) { result ->\n    llm.writeSession {\n        rewritePrompt { prompt ->\n            val messages = prompt.messages\n\n            prompt.copy(messages = listOf(messages.first(), messages.last()))\n        }\n    }\n\n    result\n}\n\nval strategy = strategy<String, String>(\"chess_strategy\") {\n    val nodeCallLLM by nodeLLMRequest(\"sendInput\")\n    val nodeExecuteTool by nodeExecuteTool(\"nodeExecuteTool\")\n    val nodeSendToolResult by nodeLLMSendToolResult(\"nodeSendToolResult\")\n    val nodeTrimHistory by nodeTrimHistory<ReceivedToolResult>()\n\n    edge(nodeStart forwardTo nodeCallLLM)\n    edge(nodeCallLLM forwardTo nodeExecuteTool onToolCall { true })\n    edge(nodeCallLLM forwardTo nodeFinish onAssistantMessage { true })\n    edge(nodeExecuteTool forwardTo nodeTrimHistory)\n    edge(nodeTrimHistory forwardTo nodeSendToolResult)\n    edge(nodeSendToolResult forwardTo nodeFinish onAssistantMessage { true })\n    edge(nodeSendToolResult forwardTo nodeExecuteTool onToolCall { true })\n}\n\nval askChoiceStrategy = AskUserChoiceSelectionStrategy(promptShowToUser = { prompt ->\n    val lastMessage = prompt.messages.last()\n    if (lastMessage is Message.Tool.Call) {\n        lastMessage.content\n    } else {\n        \"\"\n    }\n})",
   "outputs": [],
   "execution_count": 11
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:17:31.737498Z",
     "start_time": "2025-08-18T21:17:31.715531Z"
    }
   },
   "cell_type": "code",
   "source": "val promptExecutor = PromptExecutorWithChoiceSelection(baseExecutor, askChoiceStrategy)",
   "outputs": [],
   "execution_count": 12
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "The first interactive approach uses `PromptExecutorWithChoiceSelection`, which wraps the base executor with choice selection capability. The custom display function extracts move information from tool calls to show users what the AI wants to do.\n",
    "\n",
    "**Architecture Changes:**\n",
    "- **Wrapped Executor**: `PromptExecutorWithChoiceSelection` adds choice functionality to any base executor\n",
    "- **Context-Aware Display**: Shows the last tool call content instead of the full prompt\n",
    "- **Higher Temperature**: Increased to 1.0 for more diverse move options\n",
    "\n",
    "### Advanced Strategy: Manual Choice Selection"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:17:33.776618Z",
     "start_time": "2025-08-18T21:17:33.740475Z"
    }
   },
   "cell_type": "code",
   "source": "val game = ChessGame()\nval toolRegistry = ToolRegistry { tools(listOf(Move(game))) }\n\nval agent = AIAgent(\n    executor = promptExecutor,\n    strategy = strategy,\n    llmModel = OpenAIModels.Reasoning.O3Mini,\n    systemPrompt = \"\"\"\n            You are an agent who plays chess.\n            You should always propose a move in response to the \"Your move!\" message.\n\n            DO NOT HALLUCINATE!!!\n            DO NOT PLAY ILLEGAL MOVES!!!\n            YOU CAN SEND A MESSAGE ONLY IF IT IS A RESIGNATION OR A CHECKMATE!!!\n        \"\"\".trimMargin(),\n    temperature = 1.0,\n    toolRegistry = toolRegistry,\n    maxIterations = 200,\n    numberOfChoices = 3,\n)",
   "outputs": [],
   "execution_count": 13
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "The advanced strategy integrates choice selection directly into the agent's execution graph:\n",
    "\n",
    "**New Nodes:**\n",
    "- `nodeLLMSendResultsMultipleChoices`: Handles multiple LLM choices simultaneously\n",
    "- `nodeSelectLLMChoice`: Integrates the choice selection strategy into the workflow\n",
    "\n",
    "**Enhanced Control Flow:**\n",
    "- Tool results are wrapped in lists to support multiple choices\n",
    "- User selection occurs before continuing with the chosen path\n",
    "- The selected choice is unwrapped and continues through the normal flow\n",
    "\n",
    "**Benefits:**\n",
    "- **Greater Control**: Fine-grained integration with agent workflow\n",
    "- **Flexibility**: Can be combined with other agent features\n",
    "- **Transparency**: Users see exactly what the AI is considering\n",
    "\n",
    "### Running Interactive Agents"
   ]
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:18:04.659502Z",
     "start_time": "2025-08-18T21:17:36.326235Z"
    }
   },
   "cell_type": "code",
   "source": "println(\"Chess Game started!\")\n\nval initialMessage = \"Starting position is ${game.getBoard()}. White to move!\"\n\nrunBlocking {\n    agent.run(initialMessage)\n}",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Chess Game started!\n",
      "\n",
      "Available LLM choices\n",
      "Choice number 1: [Call(id=call_K46Upz7XoBIG5RchDh7bZE8F, tool=move, content={\"notation\": \"p-e2-e4\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:40.368252Z, totalTokensCount=773, inputTokensCount=315, outputTokensCount=458, additionalInfo={}))]\n",
      "Choice number 2: [Call(id=call_zJ6OhoCHrVHUNnKaxZkOhwoU, tool=move, content={\"notation\": \"p-e2-e4\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:40.368252Z, totalTokensCount=773, inputTokensCount=315, outputTokensCount=458, additionalInfo={}))]\n",
      "Choice number 3: [Call(id=call_nwX6ZMJ3F5AxiNUypYlI4BH4, tool=move, content={\"notation\": \"p-e2-e4\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:40.368252Z, totalTokensCount=773, inputTokensCount=315, outputTokensCount=458, additionalInfo={}))]\n",
      "Please choose a choice. Enter a number between 1 and 3: \n",
      "8 r n b q k b n r\n",
      "7 p p p p p p p p\n",
      "6 * * * * * * * *\n",
      "5 * * * * * * * *\n",
      "4 * * * * P * * *\n",
      "3 * * * * * * * *\n",
      "2 P P P P * P P P\n",
      "1 R N B Q K B N R\n",
      "  a b c d e f g h\n",
      "-----------------\n",
      "\n",
      "Available LLM choices\n",
      "Choice number 1: [Call(id=call_2V93GXOcIe0fAjUAIFEk9h5S, tool=move, content={\"notation\": \"p-e7-e5\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:47.949303Z, totalTokensCount=1301, inputTokensCount=341, outputTokensCount=960, additionalInfo={}))]\n",
      "Choice number 2: [Call(id=call_INM59xRzKMFC1w8UAV74l9e1, tool=move, content={\"notation\": \"p-e7-e5\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:47.949303Z, totalTokensCount=1301, inputTokensCount=341, outputTokensCount=960, additionalInfo={}))]\n",
      "Choice number 3: [Call(id=call_r4QoiTwn0F3jizepHH5ia8BU, tool=move, content={\"notation\": \"p-e7-e5\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:47.949303Z, totalTokensCount=1301, inputTokensCount=341, outputTokensCount=960, additionalInfo={}))]\n",
      "Please choose a choice. Enter a number between 1 and 3: \n",
      "8 r n b q k b n r\n",
      "7 p p p p * p p p\n",
      "6 * * * * * * * *\n",
      "5 * * * * p * * *\n",
      "4 * * * * P * * *\n",
      "3 * * * * * * * *\n",
      "2 P P P P * P P P\n",
      "1 R N B Q K B N R\n",
      "  a b c d e f g h\n",
      "-----------------\n",
      "\n",
      "Available LLM choices\n",
      "Choice number 1: [Call(id=call_f9XTizn41svcrtvnmkCfpSUQ, tool=move, content={\"notation\": \"n-g1-f3\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:55.467712Z, totalTokensCount=917, inputTokensCount=341, outputTokensCount=576, additionalInfo={}))]\n",
      "Choice number 2: [Call(id=call_c0Dfce5RcSbN3cOOm5ESYriK, tool=move, content={\"notation\": \"n-g1-f3\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:55.467712Z, totalTokensCount=917, inputTokensCount=341, outputTokensCount=576, additionalInfo={}))]\n",
      "Choice number 3: [Call(id=call_Lr4Mdro1iolh0fDyAwZsutrW, tool=move, content={\"notation\": \"n-g1-f3\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:17:55.467712Z, totalTokensCount=917, inputTokensCount=341, outputTokensCount=576, additionalInfo={}))]\n",
      "Please choose a choice. Enter a number between 1 and 3: \n",
      "8 r n b q k b n r\n",
      "7 p p p p * p p p\n",
      "6 * * * * * * * *\n",
      "5 * * * * p * * *\n",
      "4 * * * * P * * *\n",
      "3 * * * * * N * *\n",
      "2 P P P P * P P P\n",
      "1 R N B Q K B * R\n",
      "  a b c d e f g h\n",
      "-----------------\n"
     ]
    },
    {
     "ename": "org.jetbrains.kotlinx.jupyter.exceptions.ReplInterruptedException",
     "evalue": "The execution was interrupted",
     "output_type": "error",
     "traceback": [
      "The execution was interrupted"
     ]
    }
   ],
   "execution_count": 14
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:18:07.230374Z",
     "start_time": "2025-08-18T21:18:07.079347Z"
    }
   },
   "cell_type": "code",
   "source": "import ai.koog.agents.core.feature.choice.nodeLLMSendResultsMultipleChoices\nimport ai.koog.agents.core.feature.choice.nodeSelectLLMChoice\n\ninline fun <reified T> AIAgentSubgraphBuilderBase<*, *>.nodeTrimHistory(\n    name: String? = null\n): AIAgentNodeDelegate<T, T> = node(name) { result ->\n    llm.writeSession {\n        rewritePrompt { prompt ->\n            val messages = prompt.messages\n\n            prompt.copy(messages = listOf(messages.first(), messages.last()))\n        }\n    }\n\n    result\n}\n\nval strategy = strategy<String, String>(\"chess_strategy\") {\n    val nodeCallLLM by nodeLLMRequest(\"sendInput\")\n    val nodeExecuteTool by nodeExecuteTool(\"nodeExecuteTool\")\n    val nodeSendToolResult by nodeLLMSendResultsMultipleChoices(\"nodeSendToolResult\")\n    val nodeSelectLLMChoice by nodeSelectLLMChoice(askChoiceStrategy, \"chooseLLMChoice\")\n    val nodeTrimHistory by nodeTrimHistory<ReceivedToolResult>()\n\n    edge(nodeStart forwardTo nodeCallLLM)\n    edge(nodeCallLLM forwardTo nodeExecuteTool onToolCall { true })\n    edge(nodeCallLLM forwardTo nodeFinish onAssistantMessage { true })\n    edge(nodeExecuteTool forwardTo nodeTrimHistory)\n    edge(nodeTrimHistory forwardTo nodeSendToolResult transformed { listOf(it) })\n    edge(nodeSendToolResult forwardTo nodeSelectLLMChoice)\n    edge(nodeSelectLLMChoice forwardTo nodeFinish transformed { it.first() } onAssistantMessage { true })\n    edge(nodeSelectLLMChoice forwardTo nodeExecuteTool transformed { it.first() } onToolCall { true })\n}",
   "outputs": [],
   "execution_count": 15
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:18:07.861656Z",
     "start_time": "2025-08-18T21:18:07.812415Z"
    }
   },
   "cell_type": "code",
   "source": "val game = ChessGame()\nval toolRegistry = ToolRegistry { tools(listOf(Move(game))) }\n\nval agent = AIAgent(\n    executor = baseExecutor,\n    strategy = strategy,\n    llmModel = OpenAIModels.Reasoning.O3Mini,\n    systemPrompt = \"\"\"\n            You are an agent who plays chess.\n            You should always propose a move in response to the \"Your move!\" message.\n\n            DO NOT HALLUCINATE!!!\n            DO NOT PLAY ILLEGAL MOVES!!!\n            YOU CAN SEND A MESSAGE ONLY IF IT IS A RESIGNATION OR A CHECKMATE!!!\n        \"\"\".trimMargin(),\n    temperature = 1.0,\n    toolRegistry = toolRegistry,\n    maxIterations = 200,\n    numberOfChoices = 3,\n)",
   "outputs": [],
   "execution_count": 16
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:18:39.588667Z",
     "start_time": "2025-08-18T21:18:08.817688Z"
    }
   },
   "cell_type": "code",
   "source": "println(\"Chess Game started!\")\n\nval initialMessage = \"Starting position is ${game.getBoard()}. White to move!\"\n\nrunBlocking {\n    agent.run(initialMessage)\n}",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Chess Game started!\n",
      "8 r n b q k b n r\n",
      "7 p p p p p p p p\n",
      "6 * * * * * * * *\n",
      "5 * * * * * * * *\n",
      "4 * * * * P * * *\n",
      "3 * * * * * * * *\n",
      "2 P P P P * P P P\n",
      "1 R N B Q K B N R\n",
      "  a b c d e f g h\n",
      "-----------------\n",
      "\n",
      "Available LLM choices\n",
      "Choice number 1: [Call(id=call_gqMIar0z11CyUl5nup3zbutj, tool=move, content={\"notation\": \"p-e7-e5\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:17.313548Z, totalTokensCount=917, inputTokensCount=341, outputTokensCount=576, additionalInfo={}))]\n",
      "Choice number 2: [Call(id=call_6niUGnZPPJILRFODIlJsCKax, tool=move, content={\"notation\": \"p-e7-e5\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:17.313548Z, totalTokensCount=917, inputTokensCount=341, outputTokensCount=576, additionalInfo={}))]\n",
      "Choice number 3: [Call(id=call_q1b8ZmIBph0EoVaU3Ic9A09j, tool=move, content={\"notation\": \"p-e7-e5\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:17.313548Z, totalTokensCount=917, inputTokensCount=341, outputTokensCount=576, additionalInfo={}))]\n",
      "Please choose a choice. Enter a number between 1 and 3: \n",
      "8 r n b q k b n r\n",
      "7 p p p p * p p p\n",
      "6 * * * * * * * *\n",
      "5 * * * * p * * *\n",
      "4 * * * * P * * *\n",
      "3 * * * * * * * *\n",
      "2 P P P P * P P P\n",
      "1 R N B Q K B N R\n",
      "  a b c d e f g h\n",
      "-----------------\n",
      "\n",
      "Available LLM choices\n",
      "Choice number 1: [Call(id=call_pdBIX7MVi82MyWwawTm1Q2ef, tool=move, content={\"notation\": \"n-g1-f3\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:24.505344Z, totalTokensCount=1237, inputTokensCount=341, outputTokensCount=896, additionalInfo={}))]\n",
      "Choice number 2: [Call(id=call_oygsPHaiAW5OM6pxhXhtazgp, tool=move, content={\"notation\": \"n-g1-f3\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:24.505344Z, totalTokensCount=1237, inputTokensCount=341, outputTokensCount=896, additionalInfo={}))]\n",
      "Choice number 3: [Call(id=call_GJTEsZ8J8cqOKZW4Tx54RqCh, tool=move, content={\"notation\": \"n-g1-f3\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:24.505344Z, totalTokensCount=1237, inputTokensCount=341, outputTokensCount=896, additionalInfo={}))]\n",
      "Please choose a choice. Enter a number between 1 and 3: \n",
      "8 r n b q k b n r\n",
      "7 p p p p * p p p\n",
      "6 * * * * * * * *\n",
      "5 * * * * p * * *\n",
      "4 * * * * P * * *\n",
      "3 * * * * * N * *\n",
      "2 P P P P * P P P\n",
      "1 R N B Q K B * R\n",
      "  a b c d e f g h\n",
      "-----------------\n",
      "\n",
      "Available LLM choices\n",
      "Choice number 1: [Call(id=call_5C7HdlTU4n3KdXcyNogE4rGb, tool=move, content={\"notation\": \"n-g8-f6\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:34.646667Z, totalTokensCount=1621, inputTokensCount=341, outputTokensCount=1280, additionalInfo={}))]\n",
      "Choice number 2: [Call(id=call_EjCcyeMLQ88wMa5yh3vmeJ2w, tool=move, content={\"notation\": \"n-g8-f6\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:34.646667Z, totalTokensCount=1621, inputTokensCount=341, outputTokensCount=1280, additionalInfo={}))]\n",
      "Choice number 3: [Call(id=call_NBMMSwmFIa8M6zvfbPw85NKh, tool=move, content={\"notation\": \"n-g8-f6\"}, metaInfo=ResponseMetaInfo(timestamp=2025-08-18T21:18:34.646667Z, totalTokensCount=1621, inputTokensCount=341, outputTokensCount=1280, additionalInfo={}))]\n",
      "Please choose a choice. Enter a number between 1 and 3: \n",
      "8 r n b q k b * r\n",
      "7 p p p p * p p p\n",
      "6 * * * * * n * *\n",
      "5 * * * * p * * *\n",
      "4 * * * * P * * *\n",
      "3 * * * * * N * *\n",
      "2 P P P P * P P P\n",
      "1 R N B Q K B * R\n",
      "  a b c d e f g h\n",
      "-----------------\n"
     ]
    },
    {
     "ename": "org.jetbrains.kotlinx.jupyter.exceptions.ReplInterruptedException",
     "evalue": "The execution was interrupted",
     "output_type": "error",
     "traceback": [
      "The execution was interrupted"
     ]
    }
   ],
   "execution_count": 17
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "The interactive examples show how users can guide the AI's decision-making process. In the output, you can see:\n",
    "\n",
    "1. **Multiple Choices**: The AI generates 3 different move options\n",
    "2. **User Selection**: Users input numbers 1-3 to choose their preferred move\n",
    "3. **Game Continuation**: The selected move is executed and the game continues\n",
    "\n",
    "## Conclusion\n",
    "\n",
    "This tutorial demonstrates several key aspects of building intelligent agents with the Koog framework:\n",
    "\n",
    "### Key Takeaways\n",
    "\n",
    "1. **Domain Modeling**: Well-structured data models are crucial for complex applications\n",
    "2. **Tool Integration**: Custom tools enable agents to interact with external systems effectively\n",
    "3. **Memory Management**: Strategic history trimming optimizes performance for long interactions\n",
    "4. **Strategy Graphs**: Koog's graph-based approach provides flexible control flow\n",
    "5. **Interactive AI**: Choice selection enables human-AI collaboration and transparency\n",
    "\n",
    "### Framework Features Explored\n",
    "\n",
    "- ✅ Custom tool creation and integration\n",
    "- ✅ Agent strategy design and graph-based control flow\n",
    "- ✅ Memory optimization techniques\n",
    "- ✅ Interactive choice selection\n",
    "- ✅ Multiple LLM response handling\n",
    "- ✅ Stateful game management\n",
    "\n",
    "The Koog framework provides the foundation for building sophisticated AI agents that can handle complex, multi-turn interactions while maintaining efficiency and transparency."
   ]
  }
 ],
 "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
}