llm_node.py

"""
llm_node.py — Headless LLM inference node.
Runs standalone. No game imports. Pure AI.

Start:  python llm_node.py --server 192.168.1.10:5555 --port 6001 --id node_01
"""

import argparse
import asyncio
import json
import math
import random
import re
import socket
import time
import uuid
from dataclasses import dataclass

import aiohttp

# ── Constants ─────────────────────────────────────────────────────────────────
SEMAPHORE_WRITERS = 5       # Max concurrent phi3 calls
BATCH_SIZE        = 20      # Monsters per prompt
EVAL_RANGE        = 30.0    # Blocks — monsters inside this trigger evaluation
STATE_LOCK_MIN    = 60.0
STATE_LOCK_MAX    = 900.0
HEARTBEAT_EVERY   = 5.0     # Seconds between register pulses to server
OLLAMA_URL        = "http://localhost:11434/api/generate"

_PERSONALITY_BIAS = {
    "aggressive": ["ATTACK", "STALK",  "WANDER", "IGNORE", "FLEE"],
    "cunning":    ["STALK",  "ATTACK", "WANDER", "IGNORE", "FLEE"],
    "timid":      ["FLEE",   "IGNORE", "WANDER", "STALK",  "ATTACK"],
    "curious":    ["WANDER", "STALK",  "IGNORE", "FLEE",   "ATTACK"],
}

VALID_COMMANDS = {"ATTACK", "FLEE", "STALK", "WANDER", "IGNORE"}


# ── Data classes ──────────────────────────────────────────────────────────────
@dataclass
class MonsterAccount:
    """Lightweight monster record. The node owns these."""
    id:                  str
    personality:         str
    x:                   float
    y:                   float
    z:                   float
    state:               str   = "IDLE"
    thought:             str   = ""
    is_thinking:         bool  = False
    state_locked_until:  float = 0.0


@dataclass
class MonsterGroup:
    """Unit of work pushed into the semaphore queue."""
    accounts:     list   # list[MonsterAccount]
    player_ctx:   list   # [{player_name, hp, item, distance}] parallel to accounts


# ── Inference helpers ─────────────────────────────────────────────────────────
def _build_prompt(group: MonsterGroup) -> str:
    lines = []
    for i, (acc, ctx) in enumerate(zip(group.accounts, group.player_ctx)):
        lines.append(
            f"M{i}({acc.personality}): {ctx['player_name']} "
            f"{ctx['distance']}b {ctx['hp']}HP {ctx['item']}"
        )
    return (
        "Monster AI. 2-6 words each. Embed action in *stars*.\n"
        "Actions: ATTACK FLEE STALK WANDER IGNORE\n\n"
        + "\n".join(lines)
        + "\n\nReply:\n"
          "M0: I *ATTACK* the weak prey\n"
          "M1: I *FLEE* into shadows\n"
          "M2: I *STALK* from darkness\n"
          "M?: ..."
    )


def _parse_response(text: str, expected: int) -> dict:
    results = {}
    for match in re.finditer(r'M(\d+)\s*:\s*(.+)', text, re.IGNORECASE):
        idx     = int(match.group(1))
        thought = match.group(2).strip()

        thought = thought.strip('"\'')
        paren = thought.find('(')
        if paren > 0:
            thought = thought[:paren].strip()

        if not (0 <= idx < expected):
            continue
        kw = re.search(r'\*(\w+)\*', thought)
        decision = kw.group(1).upper() if kw else None
        if decision not in VALID_COMMANDS:
            decision = None
        results[idx] = {"decision": decision, "thought": thought}
    return results


async def _call_phi3(prompt: str) -> str:
    payload = {
        "model":   "phi3",
        "prompt":  prompt,
        "stream":  False,
        "options": {"temperature": 0.5, "num_predict": BATCH_SIZE * 6}
    }
    try:
        async with aiohttp.ClientSession() as s:
            async with s.post(OLLAMA_URL, json=payload, timeout=10.0) as r:
                if r.status == 200:
                    return (await r.json()).get("response", "")
    except Exception as e:
        print(f"[Node] phi3 error: {e}")
    return ""


# ── Node ──────────────────────────────────────────────────────────────────────
class LLMNode:
    def __init__(self, node_id: str, server_addr: tuple, listen_port: int):
        self.node_id     = node_id
        self.server_addr = server_addr          # (ip, port) of game server
        self.listen_port = listen_port

        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sock.bind(('0.0.0.0', listen_port))
        self.sock.setblocking(False)

        # Monster accounts this node owns
        self.accounts: dict[str, MonsterAccount] = {}

        # Latest player snapshot from server
        self.players: list[dict] = []

        # Semaphore queue
        self.group_queue: asyncio.Queue = asyncio.Queue()
        self.semaphore = asyncio.Semaphore(SEMAPHORE_WRITERS)

        self._last_heartbeat = 0.0

    # ── Public entry point ────────────────────────────────────────────────────
    async def run(self):
        print(f"[Node {self.node_id}] Online — listening :{self.listen_port} "
              f"| server {self.server_addr} | {SEMAPHORE_WRITERS} writers")

        # Launch 5 permanent writer workers
        workers = [
            asyncio.create_task(self._writer_worker(slot_id=i))
            for i in range(SEMAPHORE_WRITERS)
        ]

        await asyncio.gather(
            self._listen_loop(),
            self._tick_loop(),
            self._heartbeat_loop(),
            *workers
        )

    # ── Writer workers (the semaphore pool) ───────────────────────────────────
    async def _writer_worker(self, slot_id: int):
        """
        Permanent coroutine. One of SEMAPHORE_WRITERS running concurrently.
        Grabs a MonsterGroup, acquires a semaphore slot, inferences, ships
        commands, releases. Only 5 ever hold the semaphore simultaneously.
        """
        while True:
            group = await self.group_queue.get()

            async with self.semaphore:              # ← flag ON
                await self._process_group(group, slot_id)
                                                    # ← flag OFF (context exit)
            self.group_queue.task_done()

    async def _process_group(self, group: MonsterGroup, slot_id: int):
        prompt    = _build_prompt(group)
        raw       = await _call_phi3(prompt)
        decisions = _parse_response(raw, len(group.accounts))
        current_t = time.perf_counter()

        for i, acc in enumerate(group.accounts):
            result   = decisions.get(i, {})
            decision = result.get("decision") or _PERSONALITY_BIAS.get(
                acc.personality, ["STALK"]
            )[0]
            thought  = result.get("thought", f"I *{decision}*")

            acc.state              = decision
            acc.thought            = thought
            acc.is_thinking        = False
            acc.state_locked_until = current_t + random.uniform(STATE_LOCK_MIN, STATE_LOCK_MAX)

            # Ship command to game server
            self._send({
                "type":       "command",
                "node_id":    self.node_id,
                "monster_id": acc.id,
                "command":    decision,
                "thought":    thought,
                "target":     group.player_ctx[i]["player_name"],
            })

        print(f"[Node/{slot_id}] Batch of {len(group.accounts)} processed.")

    # ── Tick: build groups and enqueue ────────────────────────────────────────
    async def _tick_loop(self):
        """
        Every second: find monsters that need re-evaluation,
        bundle them into MonsterGroups, push into the queue.
        The semaphore workers pick them up automatically.
        """
        while True:
            await asyncio.sleep(1.0)

            if not self.players:
                continue

            current_t = time.perf_counter()
            pending   = []

            for acc in self.accounts.values():
                if acc.is_thinking:
                    continue
                if current_t < acc.state_locked_until:
                    continue

                # Find nearest player
                nearest = min(
                    self.players,
                    key=lambda p: math.hypot(p['x'] - acc.x, p['z'] - acc.z),
                    default=None
                )
                if nearest is None:
                    continue

                dist = math.hypot(nearest['x'] - acc.x, nearest['z'] - acc.z)
                if dist > EVAL_RANGE:
                    continue

                acc.is_thinking = True
                pending.append((acc, {
                    "player_name": f"Player_{nearest['id'][:4]}",
                    "hp":          nearest.get('hp', 100),
                    "item":        str(nearest.get('item', 'Nothing')),
                    "distance":    round(dist, 1),
                }))

            # Slice into BATCH_SIZE groups and enqueue
            for i in range(0, len(pending), BATCH_SIZE):
                chunk = pending[i:i + BATCH_SIZE]
                await self.group_queue.put(MonsterGroup(
                    accounts   = [p[0] for p in chunk],
                    player_ctx = [p[1] for p in chunk],
                ))


    # ── UDP listener ──────────────────────────────────────────────────────────
    async def _listen_loop(self):
        loop = asyncio.get_running_loop()

        class _UDPProtocol(asyncio.DatagramProtocol):
            def __init__(self, node):
                self.node = node

            def datagram_received(self, data, addr):
                try:
                    packet = json.loads(data.decode())
                    asyncio.ensure_future(self.node._handle_packet(packet))
                except Exception:
                    pass

            def error_received(self, exc):
                print(f"[Node] UDP error: {exc}")

        await loop.create_datagram_endpoint(
            lambda: _UDPProtocol(self),
            sock=self.sock
        )

        # Keep the coroutine alive — the protocol callbacks drive everything now
        while True:
            await asyncio.sleep(3600)

    async def _handle_packet(self, packet: dict):
        ptype = packet.get("type")

        if ptype == "assign":
            for m in packet.get("monsters", []):
                if m["id"] not in self.accounts:
                    self.accounts[m["id"]] = MonsterAccount(
                        id          = m["id"],
                        personality = m.get("personality", "curious"),
                        x           = m.get("x", 0.0),
                        y           = m.get("y", 0.0),
                        z           = m.get("z", 0.0),
                    )
            print(f"[Node] {len(self.accounts)} accounts total.")

        elif ptype == "player_snapshot":
            self.players = packet.get("players", [])

        elif ptype == "despawn":
            mid = packet.get("monster_id")
            self.accounts.pop(mid, None)

    # ── Heartbeat ─────────────────────────────────────────────────────────────
    async def _heartbeat_loop(self):
        while True:
            self._send({
                "type":     "register",
                "node_id":  self.node_id,
                "port":     self.listen_port,
                "capacity": len(self.accounts),
            })
            await asyncio.sleep(HEARTBEAT_EVERY)

    # ── Helpers ───────────────────────────────────────────────────────────────
    def _send(self, packet: dict):
        try:
            self.sock.sendto(
                json.dumps(packet).encode(),
                self.server_addr
            )
        except Exception as e:
            print(f"[Node] Send error: {e}")


# ── Entry point ───────────────────────────────────────────────────────────────
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--server", default="127.0.0.1:5555")
    parser.add_argument("--port",   type=int, default=6001)
    parser.add_argument("--id",     default=f"node_{uuid.uuid4().hex[:4]}")
    args = parser.parse_args()

    host, port = args.server.split(":")
    node = LLMNode(
        node_id     = args.id,
        server_addr = (host, int(port)),
        listen_port = args.port,
    )
    asyncio.run(node.run())