{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# Unity + Koog: Drive your game from a Kotlin Agent\n",
    "\n",
    "This notebook walks you through building a Unity-savvy AI agent with Koog using the Model Context Protocol (MCP). We'll connect to a Unity MCP server, discover tools, plan with an LLM, and execute actions against your open scene.\n",
    "\n",
    "> Prerequisites\n",
    "> - A Unity project with the Unity-MCP server plugin installed\n",
    "> - JDK 17+\n",
    "> - An OpenAI API key in the OPENAI_API_KEY environment variable\n"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "%useLatestDescriptors\n",
    "%use koog\n"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "lateinit var process: Process\n"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## 1) Provide your OpenAI API key\n",
    "We read the API key from the `OPENAI_API_KEY` environment variable so you can keep secrets out of the notebook.\n"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "val token = System.getenv(\"OPENAI_API_KEY\") ?: error(\"OPENAI_API_KEY environment variable not set\")\n",
    "val executor = simpleOpenAIExecutor(token)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## 2) Configure the Unity agent\n",
    "We define a compact system prompt and agent settings for Unity.\n"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "val agentConfig = AIAgentConfig(\n",
    "    prompt = prompt(\"cook_agent_system_prompt\") {\n",
    "        system {\n",
    "            \"You are a Unity assistant. You can execute different tasks by interacting with tools from the Unity engine.\"\n",
    "        }\n",
    "    },\n",
    "    model = OpenAIModels.Chat.GPT4o,\n",
    "    maxAgentIterations = 1000\n",
    ")"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": ""
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## 3) Start the Unity MCP server\n",
    "We'll launch the Unity MCP server from your Unity project directory and connect over stdio.\n"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "// https://github.com/IvanMurzak/Unity-MCP\n",
    "val pathToUnityProject = \"path/to/unity/project\"\n",
    "val process = ProcessBuilder(\n",
    "    \"$pathToUnityProject/com.ivanmurzak.unity.mcp.server/bin~/Release/net9.0/com.IvanMurzak.Unity.MCP.Server\",\n",
    "    \"60606\"\n",
    ").start()"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## 4) Connect from Koog and run the agent\n",
    "We discover tools from the Unity MCP server, build a small plan-first strategy, and run an agent that uses only tools to modify your open scene.\n"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import kotlinx.coroutines.runBlocking\n",
    "\n",
    "runBlocking {\n",
    "    // Create the ToolRegistry with tools from the MCP server\n",
    "    val toolRegistry = McpToolRegistryProvider.fromTransport(\n",
    "        transport = McpToolRegistryProvider.defaultStdioTransport(process)\n",
    "    )\n",
    "\n",
    "    toolRegistry.tools.forEach {\n",
    "        println(it.name)\n",
    "        println(it.descriptor)\n",
    "    }\n",
    "\n",
    "    val strategy = strategy<String, String>(\"unity_interaction\") {\n",
    "        val nodePlanIngredients by nodeLLMRequest(allowToolCalls = false)\n",
    "        val interactionWithUnity by subgraphWithTask<String, String>(\n",
    "            // work with plan\n",
    "            tools = toolRegistry.tools,\n",
    "        ) { input ->\n",
    "            \"Start interacting with Unity according to the plan: $input\"\n",
    "        }\n",
    "\n",
    "        edge(\n",
    "            nodeStart forwardTo nodePlanIngredients transformed {\n",
    "                \"Create detailed plan for \" + agentInput + \"\" +\n",
    "                    \"using the following tools: ${toolRegistry.tools.joinToString(\"\\n\") {\n",
    "                        it.name + \"\\ndescription:\" + it.descriptor\n",
    "                    }}\"\n",
    "            }\n",
    "        )\n",
    "        edge(nodePlanIngredients forwardTo interactionWithUnity onAssistantMessage { true })\n",
    "        edge(interactionWithUnity forwardTo nodeFinish transformed { it.result })\n",
    "    }\n",
    "\n",
    "    val agent = AIAgent(\n",
    "        promptExecutor = executor,\n",
    "        strategy = strategy,\n",
    "        agentConfig = agentConfig,\n",
    "        toolRegistry = toolRegistry,\n",
    "        installFeatures = {\n",
    "            install(Tracing)\n",
    "\n",
    "            install(EventHandler) {\n",
    "                onAgentStarting { eventContext ->\n",
    "                    println(\"OnAgentStarting first (strategy: ${strategy.name})\")\n",
    "                }\n",
    "\n",
    "                onAgentStarting { eventContext ->\n",
    "                    println(\"OnAgentStarting second (strategy: ${strategy.name})\")\n",
    "                }\n",
    "\n",
    "                onAgentCompleted { eventContext ->\n",
    "                    println(\n",
    "                        \"OnAgentCompleted (agent id: ${eventContext.agentId}, result: ${eventContext.result})\"\n",
    "                    )\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    )\n",
    "\n",
    "    val result = agent.run(\n",
    "        \" extend current opened scene for the towerdefence game. \" +\n",
    "            \"Add more placements for the towers, change the path for the enemies\"\n",
    "    )\n",
    "\n",
    "    result\n",
    "}"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## 5) Shut down the MCP process\n",
    "Always clean up the external Unity MCP server process at the end of your run.\n"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "// Shutdown the Unity MCP process\n",
    "process.destroy()"
   ]
  }
 ],
 "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
}