diff --git a/src/llmfunction.jl b/src/llmfunction.jl
index 0104d3d..cab0647 100644
--- a/src/llmfunction.jl
+++ b/src/llmfunction.jl
@@ -358,7 +358,7 @@ julia> result = winestock(agent, input)
 
 # TODO
   - [] update docs
-  - [WORKING] implement the function
+  - [x] implement the function
 
 # Signature
 """
@@ -370,7 +370,7 @@ function winestock(a::T1, input::T2
   # replace because SQLLLM didn't know what food_paired means
   result = replace(wineattributes, "food_pairing" => "food_to_be_paired_with_wine")
 
-  result = SQLLLM.query(Dict(:text=> wineattributes), a.executeSQL, a.text2textInstructLLM)
+  result = SQLLLM.query(result, a.executeSQL, a.text2textInstructLLM)
   return result
 end
 
diff --git a/test/runtest.jl b/test/runtest.jl
index 992ce66..018db2b 100644
--- a/test/runtest.jl
+++ b/test/runtest.jl
@@ -1,5 +1,5 @@
 using Revise      # remove when this package is completed
-using YiemAgent, GeneralUtils, JSON3, MQTTClient, Dates, UUIDs, LibPQ
+using YiemAgent, GeneralUtils, JSON3, MQTTClient, Dates, UUIDs, LibPQ, DataStructures, SQLLLM
 using Base.Threads
 
 # ---------------------------------------------- 100 --------------------------------------------- #
@@ -49,7 +49,7 @@ function text2textInstructLLM(prompt::String)
     :payload=> Dict(
       :text=> prompt,
       :kwargs=> Dict(
-        :max_tokens=> 1024,
+        :max_tokens=> 2048,
         :stop=> ["<|eot_id|>"],
         :temperature=> 0.2,
         )
@@ -70,6 +70,21 @@ end
     id="testingSessionID",  # agent instance id
   )
 
+function main()
+  userinput = "Hello, I would like a get a bottle of wine."
+  for i in 1:10
+    response = YiemAgent.conversation(a, Dict(:text=> userinput))
+    println("")
+    println("--> assistant response: \n", response)
+    println("")
+    println("--> user input:")
+    userinput = readline()
+  end
+end
+
+main()
+
+
 # response = YiemAgent.conversation(a, Dict(:text=> "newtopic",) )
 
 # response = YiemAgent.conversation(a, Dict(:text=> "Hello, I would like a get a bottle of wine."))
@@ -79,37 +94,23 @@ end
 
 
 
-# function main()
-#   userinput = "Hello, I would like a get a bottle of wine."
-#   for i in 1:10
-#     response = YiemAgent.conversation(a, Dict(:text=> userinput))
-#     println("")
-#     println("--> assistant response: \n", response)
-#     println("")
-#     println("--> user input:")
-#     userinput = readline()
-#   end
-# end
-
-# main()
 
 
 
 
-# input = "query=\"off dry, medium tannin, French Rosé\""
-input = "wine_type: rose, country: France, sweetness: 2, tannin: 3"
-YiemAgent.winestock(a, input)
+# # input = "query=\"off dry, medium tannin, French Rosé\""
+# input = "Search the database for wine type: white, country: France, sweetness level: 1"
+# YiemAgent.winestock(a, input)
 
 
 
 
-"""
-I'm having a graduation party this evening. I'll pay at most 30 bucks.
-I have no idea. The party will be formal. What type of wine people usually get for this occasion?
-What about sparkling Rose?
-
-"""
+# """
+# I'm having a graduation party this evening. I'll pay at most 30 bucks.
+# I have no idea. The party will be formal. What type of wine people usually get for this occasion?
+# What about sparkling Rose?
 
+# """