{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e6da6609-8887-416f-ba9a-d08892fb5cee",
   "metadata": {},
   "source": [
    "# Prototype Connect Four battle\n",
    "\n",
    "Pit LLMs against each other in a game of Connect Four"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c154063-8f36-426f-ae54-346acc5ba58d",
   "metadata": {
    "collapsed": false,
    "is_executing": true,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "from dotenv import load_dotenv\n",
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7446303a-d699-4816-b8dd-52da09add974",
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28740858-38f3-4aaf-9185-ec636a45ba75",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some constants\n",
    "\n",
    "RED = 1\n",
    "YELLOW = -1\n",
    "EMPTY = 0\n",
    "show = {EMPTY:\"⚪️\", RED: \"🔴\", YELLOW: \"🟡\"}\n",
    "pieces = {EMPTY: \"empty\", RED: \"red\", YELLOW: \"yellow\"}\n",
    "cols = \"ABCDEFG\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "668559eb-3542-4320-8029-eb20fda90fde",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The Game Board\n",
    "\n",
    "class Board:\n",
    "\n",
    "    def __init__(self):\n",
    "        self.cells = [[EMPTY for _ in range(7)] for _ in range(6)]\n",
    "        self.player = RED\n",
    "        self.winner = EMPTY\n",
    "\n",
    "    def __repr__(self):\n",
    "        result = \"\"\n",
    "        for y in range(6):\n",
    "            for x in range(7):\n",
    "                result += show[self.cells[5-y][x]]\n",
    "            result += \"\\n\"\n",
    "        if self.winner:\n",
    "            result += f\"\\n{show[self.winner]} wins\\n\"\n",
    "        else:\n",
    "            result += f\"\\n{show[self.player]} to play\\n\"\n",
    "        return result\n",
    "\n",
    "    def json(self):\n",
    "        result = \"{\\n\"\n",
    "        result += '    \"Column names\": [\"A\", \"B\", \"C\", \"D\", \"E\", \"F\", \"G\"],\\n'\n",
    "        for y in range(6):\n",
    "            result += f'    \"Row {6-y}\": [' \n",
    "            for x in range(7):\n",
    "                result += f'\"{pieces[self.cells[5-y][x]]}\", '\n",
    "            result = result[:-2] + '],\\n'\n",
    "        result = result[:-2]+'\\n}'\n",
    "        return result  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec586c38-f396-4cb2-95ef-c8e6d067d19b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# See the game board\n",
    "\n",
    "Board()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a17c172-f711-4ebd-9aae-3f1f35665433",
   "metadata": {},
   "outputs": [],
   "source": [
    "# And the json representation\n",
    "\n",
    "print(Board().json())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5ad4649-0a48-47b1-bb04-ee28fd6808b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some useful methods\n",
    "\n",
    "def height(self, x):\n",
    "    height = 0\n",
    "    while height<6 and self.cells[height][x] != EMPTY:\n",
    "        height += 1\n",
    "    return height\n",
    "\n",
    "def legal_moves(self):\n",
    "    return [cols[x] for x in range(7) if self.height(x)<6]\n",
    "\n",
    "def move(self, x):\n",
    "    self.cells[self.height(x)][x] = self.player\n",
    "    self.player = -1 * self.player\n",
    "\n",
    "Board.height = height\n",
    "Board.legal_moves = legal_moves\n",
    "Board.move = move"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3da4a6df-0217-492b-a050-9488849ff2a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "b = Board()\n",
    "b.move(3)\n",
    "b.move(3)\n",
    "b.move(2)\n",
    "b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f5d0b85-aaf8-4d55-bf18-8602dade4cf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "b.legal_moves()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dffa507-e865-4379-a765-cd27032ff657",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test for winning move\n",
    "\n",
    "def winning_line(self, x, y, dx, dy):\n",
    "    color = self.cells[y][x]\n",
    "    for pointer in range(1, 4):\n",
    "        xp = x + dx * pointer\n",
    "        yp = y + dy * pointer\n",
    "        if not (0 <= xp <= 6 and 0 <= yp <= 5) or self.cells[yp][xp] != color:\n",
    "            return EMPTY\n",
    "    return color\n",
    "\n",
    "def winning_cell(self, x, y):\n",
    "    for dx, dy in ((0, 1), (1, 1), (1, 0), (1, -1)):\n",
    "        if winner := self.winning_line(x, y, dx, dy):\n",
    "            return winner\n",
    "    return EMPTY\n",
    "\n",
    "def wins(self):\n",
    "    for y in range(6):\n",
    "        for x in range(7):\n",
    "            if winner := self.winning_cell(x, y):\n",
    "                return winner\n",
    "    return EMPTY\n",
    "\n",
    "def move(self, x):\n",
    "    self.cells[self.height(x)][x] = self.player\n",
    "    if winner := self.wins():\n",
    "        self.winner = winner\n",
    "    else:\n",
    "        self.player = -1 * self.player\n",
    "    return self\n",
    "\n",
    "Board.winning_line = winning_line\n",
    "Board.winning_cell = winning_cell\n",
    "Board.wins = wins\n",
    "Board.move = move"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65dc68e5-54a3-4e9f-8cee-008bbe313b20",
   "metadata": {},
   "outputs": [],
   "source": [
    "b = Board()\n",
    "b.move(2).move(3).move(2).move(3).move(2).move(3).move(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29c68043-cd96-49aa-9a30-ad1ac7f9c00b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# And now - a player that calls gpt-4o-mini\n",
    "\n",
    "class Player:\n",
    "\n",
    "    def __init__(self, model, color):\n",
    "        self.color = color\n",
    "        self.model = model\n",
    "        self.llm = OpenAI()\n",
    "\n",
    "    def system(self, board):\n",
    "        legal_moves = \", \".join(board.legal_moves())\n",
    "        return f\"\"\"You are an expert player of the board game Connect 4.\n",
    "Players take turns to drop counters into one of 7 columns labelled A, B, C, D, E, F, G.\n",
    "The winner is the first player to get 4 coins in a row in a straight or diagonal line.\n",
    "You are playing with the {pieces[self.color]} coins.\n",
    "And your opponent is playing with the {pieces[self.color * -1]} coins.\n",
    "You will be presented with the board and asked to pick a column to drop your piece.\n",
    "You must pick one of the following legal moves: {legal_moves}. You must pick one of those letters.\n",
    "You should respond in JSON, and only in JSON, according to this spec:\n",
    "\n",
    "{{\n",
    "    \"evaluation\": \"brief assessment of the board\",\n",
    "    \"threats\": \"any threats from your opponent or weaknesses in your position\",\n",
    "    \"opportunities\": \"any opportunities to gain the upper hand or strengths in your position\",\n",
    "    \"strategy\": \"the thought process behind your next move\",\n",
    "    \"move_column\": \"one letter from this list of legal moves: {legal_moves}\"\n",
    "}}\"\"\"\n",
    "\n",
    "    def user(self, board):\n",
    "        legal_moves = \", \".join(board.legal_moves())\n",
    "        return f\"\"\"It is your turn to make a move as {pieces[self.color]}.\n",
    "The current board position is:\n",
    "\n",
    "{board.json()}\n",
    "\n",
    "Now with this in mind, make your decision. Respond only in JSON strictly according to this spec:\n",
    "\n",
    "{{\n",
    "    \"evaluation\": \"brief assessment of the board\",\n",
    "    \"threats\": \"any threats from your opponent or weaknesses in your position\",\n",
    "    \"opportunities\": \"any opportunities to gain the upper hand or strengths in your position\",\n",
    "    \"strategy\": \"the thought process behind your next move\",\n",
    "    \"move_column\": \"one of {legal_moves} which are the legal moves\"\n",
    "}}\n",
    "\n",
    "You must pick one of these letters for your move_column: {legal_moves}\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "    def process_move(self, reply):\n",
    "        print(reply)\n",
    "        try:\n",
    "            result = json.loads(reply)\n",
    "            move = result.get(\"move_column\") or \"\"\n",
    "            move = move.upper()\n",
    "            col = cols.find(move)\n",
    "            if not (0 <= col <= 6) or board.height(col)==6:\n",
    "                raise ValueError(\"Illegal move\")\n",
    "            board.move(col)\n",
    "        except Exception as e:\n",
    "            print(f\"Exception {e}\")\n",
    "            board.winner = -1 * board.player\n",
    "            \n",
    "    \n",
    "    def move(self, board):\n",
    "        system = self.system(board)\n",
    "        user = self.user(board)\n",
    "        reply = self.llm.chat.completions.create(\n",
    "            model=self.model,\n",
    "            messages=[\n",
    "                {\"role\": \"system\", \"content\": system},\n",
    "                {\"role\": \"user\", \"content\": user}\n",
    "            ],\n",
    "            response_format={\"type\": \"json_object\"}\n",
    "            \n",
    "        )\n",
    "        self.process_move(reply.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a3384f5-a1a7-45d2-8201-523e41db173c",
   "metadata": {},
   "source": [
    "# Let's do this!\n",
    "\n",
    "Wrap it in a loop, and we're off!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "163e2af5-83c1-444d-8194-92e8b495c0af",
   "metadata": {},
   "outputs": [],
   "source": [
    "board = Board()\n",
    "red = Player(\"gpt-4o-mini\", RED)\n",
    "yellow = Player(\"gpt-4o\", YELLOW)\n",
    "while not board.winner:\n",
    "    red.move(board)\n",
    "    print(board)\n",
    "    if not board.winner:\n",
    "        yellow.move(board)\n",
    "        print(board)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f73d9d16-9658-4e85-8d46-f75e6adf7051",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}