diff --git a/metagpt/actions/di/write_analysis_code.py b/metagpt/actions/di/write_analysis_code.py
index 00e6d174d..39de18987 100644
--- a/metagpt/actions/di/write_analysis_code.py
+++ b/metagpt/actions/di/write_analysis_code.py
@@ -31,8 +31,12 @@ class WriteAnalysisCode(Action):
 
         rsp = await self._aask(reflection_prompt, system_msgs=[REFLECTION_SYSTEM_MSG])
         reflection = json.loads(CodeParser.parse_code(text=rsp))
+        try:
+            improved_impl = CodeParser.parse_code(text=reflection["improved_impl"], lang="python")
+        except Exception as e:
+            improved_impl = reflection["improved_impl"]
 
-        return reflection["improved_impl"]
+        return improved_impl
 
     async def run(
         self,