updated with latest versions

Author: kiranandcode

tests/test_maker.py

@@ -1,25 +1,24 @@
 import itertools
 import logging
-import random
-import re
 import sys
-from collections import Counter
+from abc import ABC, abstractmethod
 from typing import Optional
 
-from openai import OpenAI
+import pydantic
+from litellm import ConfigDict
+from PIL import Image, ImageDraw
 from pydantic.dataclasses import dataclass
 
-from effectful.handlers import futures
-from effectful.handlers.futures import Executor, ThreadPoolFuturesInterpretation
 from effectful.handlers.llm import Template
-from effectful.handlers.llm.providers import LLMLoggingHandler, OpenAIAPIProvider
+from effectful.handlers.llm.providers import (
+    LiteLLMProvider,
+    RetryLLMHandler,
+)
+from effectful.handlers.llm.sampling import KAheadSampler
 from effectful.ops.semantics import handler
+from effectful.ops.types import NotHandled
 
-
-@dataclass(frozen=True)
-class Step:
-    start: int
-    end: int
+type Step = tuple[int, int]
 
 
 @dataclass(frozen=True)
@@ -40,25 +39,8 @@ def new(cls, size: int) -> "GameState":
         towers = tuple(tuple(tower) for tower in towers)
         return cls(size, towers)
 
-    def visualise_text(self):
-        max_disk = self.size
-        width = max_disk * 2 + 3
-        for i, tower in enumerate(self.towers):
-            print(f"\nTower {i}:")
-            for disk in reversed(tower):
-                disk_width = (disk + 1) * 2 - 1
-                padding = (max_disk - disk_width) // 2
-                print(" " * padding + "=" * disk_width + " " * padding)
-            print("=" * width)
-        print()
-
-    def visualise_image(self):
+    def visualise_image(self) -> Image:
         "Uses python graphics libraries to visualise the state of the hanoi game."
-        try:
-            from PIL import Image, ImageDraw
-        except ImportError:
-            return None
-        # Pillow-based visualization
         tower_width = 150
         disk_height = 30
         base_height = 20
@@ -116,16 +98,13 @@ def visualise_image(self):
 
     def visualise(self):
         img = self.visualise_image()
-        if img:
-            img.show()
-        else:
-            self.visualise_text()
+        img.show()
 
     def apply(self, step: Step) -> Optional["GameState"]:
         """
         Given a tower `start` and a target tower `end` moves the topmost disk to the end tower.
         """
-        start, end = step.start, step.end
+        start, end = step
 
         if not (0 <= start < len(self.towers) and 0 <= end < len(self.towers)):
             return None
@@ -149,25 +128,6 @@ def apply(self, step: Step) -> Optional["GameState"]:
         )
         return new_state
 
-    def steps_to_complete(self) -> int:
-        """Compute the number of steps to complete the towers of hanoi from a given configuration if using the optimal algorithm."""
-        # Count disks on each tower
-        total_moves = 0
-
-        # For each tower that's not the destination, we need to move all its disks
-        for tower_idx, tower in enumerate(self.towers):
-            if tower_idx == self.size - 1:
-                continue
-
-            # Number of disks on this tower
-            n_disks = len(tower)
-
-            if n_disks > 0:
-                # Moving n disks from one peg to another requires 2^n - 1 moves
-                total_moves += (2**n_disks) - 1
-
-        return total_moves
-
     def is_done(self) -> bool:
         return all(len(tower) == 0 for tower in self.towers[:-1]) and all(
             self.towers[-1][i] > self.towers[-1][i + 1]
@@ -184,24 +144,62 @@ def valid_steps(self) -> list[Step]:
                     continue
                 # if tower_i's disk is smaller than tower_j's topmost, then it is valid to move from tower i to j
                 if len(tower_j) == 0 or tower_i[-1] < tower_j[-1]:
-                    steps.append(Step(i, j))
+                    steps.append((i, j))
         return steps
 
 
-class MicroAgent:
-    """Micro agent (based on MAKERS paper) responsible for predicting a single next step."""
+class Step(ABC):
+    @property
+    @abstractmethod
+    def start(self) -> int:
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def end(self) -> int:
+        raise NotImplementedError
+
+
+def build_validated_model(game_state: GameState) -> type[Step]:
+    valid_steps = game_state.valid_steps()
+
+    @pydantic.dataclasses.dataclass(frozen=True)
+    class StepModel:
+        start: int
+        end: int
+        explanation: str = ""
+        model_config = ConfigDict(extra="forbid")
+
+        @pydantic.field_validator("start", "end", mode="before")
+        def validate_indices(cls, v, info):
+            if isinstance(v, int):
+                if not (0 <= v < len(game_state.towers)):
+                    raise ValueError(f"{info.field_name} {v} out of range")
+            else:
+                raise TypeError("start/end must both be int")
+            return v
 
-    game_state: GameState
+        @pydantic.model_validator(mode="after")
+        def validate_step(self):
+            if (self.start, self.end) not in valid_steps:
+                raise ValueError("step is not in {self.valid_steps}")
+            return self
 
-    def __init__(self, state: GameState):
-        self.game_state = state
+        def __hash__(self):
+            return hash((self.start, self.end))
+
+    return StepModel
+
+
+def predict_next_step(game_state: GameState) -> Step:
+    ValidStep = build_validated_model(game_state)
 
     @Template.define
-    def predict_next_step(self) -> str:
+    def predict_next_step_inner(game_state) -> ValidStep:
         """
         Given the state of the game of towers of Hanoi as follows:
 
-        {self.game_state}
+        {game_state}
 
         Predict the next step to complete the game (moving all disks to the rightmost tower).
 
@@ -211,96 +209,24 @@ def predict_next_step(self) -> str:
 
         where start and end are zero-based indices for the towers to move. Be concise and avoid wordy answers.
         """
-        pass
-
-    def parse_response(self, response: str) -> Step | None:
-        "Parse the predicted step from an LLM response."
-        pattern = r"<step>\s*(\d+)\s*,\s*(\d+)\s*</step>"
-        m = re.search(pattern, response)
-        if not m:
-            return None
-        return Step(int(m.group(1)), int(m.group(2)))
-
-    def has_no_red_flags(self, response: str) -> Step | None:
-        """Returns the underlying step if the provided step has no red flags."""
-        if len(response) > 450.0:  # based on a sample
-            return None
-
-        step = self.parse_response(response)
-        if not step:
-            return None
-        if not (
-            0 <= step.start < len(self.game_state.towers)
-            and 0 <= step.end < len(self.game_state.towers)
-        ):
-            return None
-        if step not in self.game_state.valid_steps():
-            return None
-        return step
-
-    def get_vote(self):  # algorithm 3
-        while True:
-            resp = self.predict_next_step()
-            if step := self.has_no_red_flags(resp):
-                return step
-
-
-class FirstToAheadMoveSelector:
-    k: int
-    game_state: GameState
-    agents: list[MicroAgent]
-    votes: Counter[Step]
-
-    def __init__(self, state: GameState, no_agents=6, k=3):
-        self.k = k
-        self.game_state = state
-        self.agents = [MicroAgent(self.game_state) for _ in range(no_agents)]
-        self.votes = Counter()
-
-    def do_voting(self) -> Step:  # algorithm 2
-        # run n in parallel repeatedly until k come out in top
-        while True:
-            # submit a batch of votes
-            for vote in futures.as_completed(
-                [Executor.submit(agent.get_vote) for agent in self.agents]
-            ):
-                vote = vote.result()
-                self.votes[vote] += 1
-                max_other_votes = max(
-                    (self.votes[o_vote] for o_vote in self.votes if o_vote != vote),
-                    default=0,
-                )
-                if self.votes[vote] >= max_other_votes + self.k:
-                    return vote
-
-
-def calculate_average_sample_size():
-    """Function I used to calculate the number 450. in the above code."""
-    sizes = []
-    samples = []
+        raise NotHandled
 
-    with handler(OpenAIAPIProvider(OpenAI())):
-        for _ in range(10):
-            s = GameState.new(random.randint(3, 6))
-            for i in range(100):
-                step = random.choice(s.valid_steps())
-                s = s.apply(step) or s
-            resp = MicroAgent(s).predict_next_step()
-            samples.append(resp)
-            sizes.append(len(resp))
-    return sum(sizes) / len(sizes)
+    s = predict_next_step_inner(game_state)
+    return (s.start, s.end)
 
 
 def solve_hanoi(state: GameState):
     log = []
 
     for i in itertools.count():
         print(f"step {i} - {state}")
-        step = FirstToAheadMoveSelector(state).do_voting()
+        with handler(KAheadSampler()), handler(RetryLLMHandler()):
+            step = predict_next_step(state)
         # track the step at each point
-        log.append((state, step))
+        if new_state := state.apply(step):
+            log.append((state, step))
 
-        state = state.apply(step)
+        state = new_state or state
         state.visualise()
         if state.is_done():
             break
@@ -313,8 +239,6 @@ def solve_hanoi(state: GameState):
 )
 
 with (
-    handler(ThreadPoolFuturesInterpretation()),
-    handler(OpenAIAPIProvider(OpenAI())),
-    handler(LLMLoggingHandler()),
+    handler(LiteLLMProvider(model_name="gpt-4o-mini")),
 ):
     solve_hanoi(state=GameState.new(3))