update

examples/code_generation.jl

@@ -0,0 +1,115 @@
+# Code Generation - MCTS for Programming Tasks
+
+This example shows how MCTS can guide LLM code generation by exploring different implementation strategies.
+
+```julia
+using LLMMCTS
+
+# State represents the current state of code generation
+# It includes the code written so far and the problem being solved
+
+function code_generation_transition(state::Dict, args::NamedTuple)
+    current_step = get(state, :step, 0)
+    problem = state[:problem]
+    code_so_far = get(state, :code, "")
+    
+    if current_step == 0
+        # First step: Plan the approach
+        new_code = """
+        # Function to solve: $(problem)
+        function solve_problem(input)
+        """
+        newstate = Dict(
+            :step => 1,
+            :code => new_code,
+            :thought => "Plan the approach for: $(problem)",
+            :reward => 2.0,
+            :isterminal => false
+        )
+        progressvalue = 5.0
+    elseif current_step == 1
+        # Second step: Implement main logic
+        new_code = code_so_far * """
+            # Main logic implementation
+            result = input * 2  # Placeholder implementation
+            return result
+        end
+        """
+        newstate = Dict(
+            :step => 2,
+            :code => new_code,
+            :thought => "Implement main function logic",
+            :reward => 3.0,
+            :isterminal => false
+        )
+        progressvalue = 7.0
+    elseif current_step == 2
+        # Third step: Add tests
+        new_code = code_so_far * """
+        
+        # Test the function
+        @assert solve_problem(5) == 10
+        @assert solve_problem(0) == 0
+        println("All tests passed!")
+        """
+        newstate = Dict(
+            :step => 3,
+            :code => new_code,
+            :thought => "Add unit tests to verify implementation",
+            :reward => 5.0,
+            :isterminal => true  # Code generation complete
+        )
+        progressvalue = 10.0
+    else
+        newstate = Dict(
+            :step => current_step,
+            :code => code_so_far,
+            :thought => "Code generation complete",
+            :reward => 10.0,
+            :isterminal => true
+        )
+        progressvalue = 10.0
+    end
+    
+    return Dict(
+        :newNodeKey => "code_step_$current_step",
+        :newstate => newstate,
+        :progressvalue => progressvalue
+    )
+end
+
+# Initial state
+initialstate = Dict(
+    :step => 0,
+    :problem => "Create a function that doubles its input",
+    :code => "",
+    :reward => 0,
+    :isterminal => false
+)
+
+# Transition arguments
+transitionargs = (max_steps = 3,)
+
+# Run MCTS
+result = runMCTS(
+    initialstate,
+    code_generation_transition,
+    transitionargs;
+    maxiterations = 20,
+    explorationweight = 1.0,
+    maxSimulationDepth = 3,
+    horizontalSampleExpansionPhase = 3
+)
+
+# Display results
+println("Code Generation Example")
+println("=======================")
+println()
+println("Problem: ", initialstate[:problem])
+println()
+println("Generated code:")
+println(result.bestTerminalState[:code])
+println()
+println("Code generation complete! ✓")
+println("Root node visits: ", result.root.visits)
+```