update

src/interface.jl

@@ -12,7 +12,7 @@ using ..type, ..mcts, ..util
 """ Search the best action to take for a given state and task 
 
 # Arguments
-  - `initial state`
+  - `initialstate::T`
     initial state
   - `transition::Function`
     a function that define how the state transitions
@@ -32,21 +32,16 @@ using ..type, ..mcts, ..util
     aggressively explore new state.
 
 # Return
-  - `(bestNextState, BestFinalState)::Tuple`
+  - `(bestNextState, BestFinalState)::@NamedTuple{bestNextState::T, bestFinalState::T}`
     the best next state and the best final state
 
 # Example
-```jldoctest
-julia> 
-```
-
-# TODO
-  [] update example
+    Refers to SQLLLM package
 
 # Signature
 """
 function runMCTS(
-  initialstate,
+  initialstate::T,
   transition::Function,
   transitionargs::NamedTuple,
   ;
@@ -54,7 +49,7 @@ function runMCTS(
   maxdepth::Integer=3, 
   maxiterations::Integer=10, 
   explorationweight::Number=1.0,
-  )::NamedTuple
+  )::@NamedTuple{bestNextState::T, bestFinalState::T} where {T<:Any}
   
   root = MCTSNode("root", initialstate, 0, 0, 0, 0, false, nothing, Dict{String, MCTSNode}())
 
@@ -88,8 +83,8 @@ function runMCTS(
     end
   end
 
-  bestNextState = selectBestNextState(root)
-  besttrajectory = selectBestTrajectory(root)
+  bestNextState = selectBestNextNode(root)
+  besttrajectory = selectBestTrajectoryNode(root)
 
   return (bestNextState=bestNextState.state, bestFinalState=besttrajectory.state)
 end