diff --git a/core b/core
new file mode 100644
index 0000000..1c3171f
Binary files /dev/null and b/core differ
diff --git a/src/interface.jl b/src/interface.jl
index cf0ca1b..3133bb1 100644
--- a/src/interface.jl
+++ b/src/interface.jl
@@ -25,7 +25,7 @@ julia>
 
 # Signature
 """
-function jsoncorrection(config::Dict, jsonstring::String, context::Dict) 
+function jsoncorrection(config::Dict, jsonstring::String, context) 
 
   initialstate = Dict{Symbol, Any}(
     :reward=> 0,
@@ -46,13 +46,13 @@ function jsoncorrection(config::Dict, jsonstring::String, context::Dict)
                     reflector=reflector,
                     context=context,
                   )
-  result, _ = LLMMCTS.runMCTS(initialstate, transition, transitionargs;
+  _, result = LLMMCTS.runMCTS(initialstate, transition, transitionargs;
                           totalsample=1, maxdepth=3, maxiterations=1, explorationweight=1.0)
 
   if result[:response] !== nothing
-    return (response=result[:response], select=nothing, reward=0, isterminal=false)
+    return response=result[:code]
   else
-    return (response=result[:errorexplain], select=nothing, reward=0, isterminal=false)
+    return response=result[:errorexplain]
   end
 end
 
@@ -84,7 +84,7 @@ function transition(state::T1, args::NamedTuple
   decisionMaker::Function = args[:decisionMaker]
   evaluator::Function = args[:evaluator]
   reflector::Function = args[:reflector]
-  context::Union{AbstractDict, Nothing} = args[:context]
+  context = args[:context]
 
   # only for the 1st transition
   if state[:errormsg] === nothing
@@ -139,8 +139,8 @@ function decisionMaker(state, config, context)
   At each round of conversation, the user will give you:
   Context: ....
   Query: The user's original incorrect JSON string
-  Code from the last round: Your revised JSON string
-  Execution error: Your JSON string's parsing error
+  Code from the last round: Your previous revised JSON string
+  Execution error: Code from the last round error
 
   You should then respond to the user with:
   - Why: Why couldn't the JSON string be loaded? Are there any steps missing in your plan?  What does the execution error imply?
@@ -157,14 +157,14 @@ function decisionMaker(state, config, context)
   usermsg =
   if state[:code] !== nothing
     """
-    Context: $(JSON3.write(context[:expectedJsonExample]))
+    Context: $context
     Query: $(state[:question])
     Code from the last round: $(state[:code])
     Execution error: $(state[:errormsg])
     """
   else
     """
-    Context: $(JSON3.write(context[:expectedJsonExample]))
+    Context: $context
     Query: $(state[:question])
     Code from the last round: None
     Execution error: $(state[:errormsg])
@@ -210,18 +210,19 @@ function decisionMaker(state, config, context)
       )
     )
 
-  response = GeneralUtils.sendReceiveMqttMsg(outgoingMsg)
-  _responseJsonStr = response[:response][:text]
-  code =
-  if length(split(_responseJsonStr, "Code:")) == 2
-    split(_responseJsonStr, "Code:")
-  elseif length(split(_responseJsonStr, "Code\":")) == 2
-    split(_responseJsonStr, "Code\":")
-  else
-    error("failed to get Code part")
-  end
+  for i in 1:5
+    response = GeneralUtils.sendReceiveMqttMsg(outgoingMsg)
+    _responseJsonStr = response[:response][:text]
 
-  return code
+    if length(split(_responseJsonStr, "Code:")) == 2
+      return split(_responseJsonStr, "Code:")
+    elseif length(split(_responseJsonStr, "Code\":")) == 2
+      return split(_responseJsonStr, "Code\":")
+    else
+      println("Trying to get Code part")
+    end
+  end
+  error("Failed to get Code part.")
 end
 
 
diff --git a/test/runtest.jl b/test/runtest.jl
index 493ea2e..887d990 100644
--- a/test/runtest.jl
+++ b/test/runtest.jl
@@ -3,19 +3,17 @@ using FormatCorrector
 
 incorrectjson = "    \"explain\": \"The user is asking about wines that are suitable for pairing with lamb. I need to explore the database to find relevant information.\",\n    \"plan\":\n        1) List all tables in the database using the function `listalltables` to get an idea of what data is available.\n    \"action\": {\"name\": \"listalltables\", \"input\": null},\n    \"expectation\": The list of all tables in the database,\n    \"observation\": null\n}"
   
-context = Dict(
-  :expectedJsonExample=> 
-    """
-    Here is an expected JSON format:
-    {
-      "explain": ...,
-      "plan": ...,
-      "action": {"name": ..., "input": ...},
-      "expectation": ...,
-      "observation": ...
-    }
-    """
-)
+context =
+  """
+  Here is an expected JSON format:
+  {
+    "explain": ...,
+    "plan": ...,
+    "action": {"name": ..., "input": ...},
+    "expectation": ...,
+    "observation": ...
+  }
+  """
 
 config = Dict(
   :mqttServerInfo => Dict(
@@ -34,12 +32,7 @@ config = Dict(
 
 result = FormatCorrector.jsoncorrection(config, incorrectjson, context)
 
-
-
-
-
-
-
+println("done")