From 83e5eb133a8d48a84c67b46d237d5c95e5caca2f Mon Sep 17 00:00:00 2001
From: luxiangtao <luxiangtao@fuzhi.ai>
Date: Sun, 24 Mar 2024 15:01:45 +0800
Subject: [PATCH] updata di example and Change the location of the .ipynb
 locations

---
 examples/di/data_visualization.py | 15 +++++++++++----
 metagpt/utils/recovery_util.py    |  2 +-
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/examples/di/data_visualization.py b/examples/di/data_visualization.py
index 1a21ab7cb..62e7e7c2f 100644
--- a/examples/di/data_visualization.py
+++ b/examples/di/data_visualization.py
@@ -1,14 +1,21 @@
 import asyncio
-
+import os
+import sys
+import metagpt.utils
+from metagpt.logs import logger
+from metagpt.utils.recovery_util import save_history
 from metagpt.roles.di.data_interpreter import DataInterpreter
 
-
 async def main(requirement: str = ""):
+
     di = DataInterpreter()
-    await di.run(requirement)
+    rsp = await di.run(requirement)
+    logger.info(rsp)
+    save_history(role=di)
 
 
 if __name__ == "__main__":
-    requirement = "Run data analysis on sklearn Iris dataset, include a plot"
 
+    requirement = "Run data analysis on sklearn Iris dataset, include a plot"
     asyncio.run(main(requirement))
+
diff --git a/metagpt/utils/recovery_util.py b/metagpt/utils/recovery_util.py
index d0b197e69..2089ae018 100644
--- a/metagpt/utils/recovery_util.py
+++ b/metagpt/utils/recovery_util.py
@@ -54,5 +54,5 @@ def save_history(role: Role, save_dir: str = ""):
     with open(save_path / "plan.json", "w", encoding="utf-8") as plan_file:
         json.dump(plan, plan_file, indent=4, ensure_ascii=False)
 
-    save_code_file(name=Path(record_time) / "history_nb", code_context=role.execute_code.nb, file_format="ipynb")
+    save_code_file(name=Path(record_time), code_context=role.execute_code.nb, file_format="ipynb")
     return save_path