diff --git a/taskingai/_version.py b/taskingai/_version.py
index 2e42793..8be8d96 100644
--- a/taskingai/_version.py
+++ b/taskingai/_version.py
@@ -1,2 +1,2 @@
 __title__ = "taskingai"
-__version__ = "0.2.0"
+__version__ = "0.2.1"
diff --git a/taskingai/assistant/message.py b/taskingai/assistant/message.py
index 30f31de..a301749 100644
--- a/taskingai/assistant/message.py
+++ b/taskingai/assistant/message.py
@@ -246,7 +246,7 @@ def generate_message(
     """
 
     body = MessageGenerateRequest(
-        system_prompt_variables=system_prompt_variables,
+        system_prompt_variables=system_prompt_variables or {},
         stream=stream,
     )
 
@@ -281,7 +281,7 @@ async def a_generate_message(
     """
 
     body = MessageGenerateRequest(
-        system_prompt_variables=system_prompt_variables,
+        system_prompt_variables=system_prompt_variables or {},
         stream=stream,
     )
 
diff --git a/test/run_test.sh b/test/run_test.sh
index 0c1a3fc..a7fc76b 100644
--- a/test/run_test.sh
+++ b/test/run_test.sh
@@ -4,9 +4,9 @@ export PYTHONPATH="${PYTHONPATH}:${parent_dir}"
 
 echo "Starting tests..."
 
-pytest  ./test/testcase/test_sync
+pytest  ./test/testcase/test_sync  --reruns 2 --reruns-delay 1
 sleep 5
-pytest  ./test/testcase/test_async
+pytest  ./test/testcase/test_async  --reruns 2 --reruns-delay 1
 
 echo "Tests completed."
 
diff --git a/test/testcase/test_async/test_async_retrieval.py b/test/testcase/test_async/test_async_retrieval.py
index 5dea9a2..9c68e47 100644
--- a/test/testcase/test_async/test_async_retrieval.py
+++ b/test/testcase/test_async/test_async_retrieval.py
@@ -257,7 +257,7 @@ async def test_a_query_chunks(self):
 
         query_text = "Machine learning"
         top_k = 1
-        res = await a_query_chunks(collection_id=self.collection_id, query_text=query_text, top_k=top_k)
+        res = await a_query_chunks(collection_id=self.collection_id, query_text=query_text, top_k=top_k, max_tokens=20000)
         pytest.assume(len(res) == top_k)
         for chunk in res:
             chunk_dict = vars(chunk)
diff --git a/test/testcase/test_sync/test_sync_retrieval.py b/test/testcase/test_sync/test_sync_retrieval.py
index ec50689..b295e5e 100644
--- a/test/testcase/test_sync/test_sync_retrieval.py
+++ b/test/testcase/test_sync/test_sync_retrieval.py
@@ -248,7 +248,7 @@ def test_query_chunks(self, collection_id):
 
         query_text = "Machine learning"
         top_k = 1
-        res = query_chunks(collection_id=collection_id, query_text=query_text, top_k=top_k)
+        res = query_chunks(collection_id=collection_id, query_text=query_text, top_k=top_k, max_tokens=20000)
         pytest.assume(len(res) == top_k)
         for chunk in res:
             chunk_dict = vars(chunk)
diff --git a/test_requirements.txt b/test_requirements.txt
index 33c955e..6aba4a4 100644
--- a/test_requirements.txt
+++ b/test_requirements.txt
@@ -15,4 +15,5 @@ pytest-tornasync>=0.6.0
 pytest-trio==0.8.0
 pytest-twisted==1.14.0
 Twisted==24.3.0
-python-dotenv==1.0.0
\ No newline at end of file
+python-dotenv==1.0.0
+pytest-rerunfailures==13.0
\ No newline at end of file