common tapedata.sqlite for the experiment

Author: ollmer

src/agentlab/analyze/tapes.py

@@ -7,6 +7,7 @@
 import numpy as np
 import yaml
 from tapeagents.core import Step, StepMetadata
+from tapeagents.observe import retrieve_all_llm_calls
 from tapeagents.renderers.camera_ready_renderer import CameraReadyRenderer
 from tapeagents.tape_browser import TapeBrowser
 
@@ -88,7 +89,17 @@ def get_steps(self, tape: dict) -> list:
         return tape["steps"]
 
     def load_llm_calls(self):
-        pass
+        sqlite_path = self.exp_path / "tapedata.sqlite"
+        if sqlite_path.exists():
+            try:
+                self.llm_calls = {
+                    call.prompt.id: call for call in retrieve_all_llm_calls(str(sqlite_path))
+                }
+                logger.info(f"Loaded {len(self.llm_calls)} LLM calls from {sqlite_path}")
+            except Exception as e:
+                logger.warning(f"Failed to load LLM calls from {sqlite_path}: {e}")
+        else:
+            logger.warning(f"{sqlite_path} not found")
 
     def get_context(self, tape: Tape) -> list:
         return []
@@ -209,7 +220,8 @@ def load_tapes(self, exp_dir: str) -> list[Tape]:
                     tapes.append(tape)
             except Exception as e:
                 logger.warning(f"Failed to load {json_file}: {e}")
-        logger.info(f"Loaded {len(tapes)} tapes from {exp_dir}")
+        logger.info(f"Loaded {len(tapes)} tapes from {fpath}")
+        self.exp_path = fpath
         return sorted(
             tapes,
             key=lambda x: f"{x.metadata.task.get('Level', '')}{x.metadata.task.get('number', 0):03d}",

src/agentlab/benchmarks/gaia.py

@@ -10,6 +10,7 @@
 import datasets
 import hydra
 import podman
+import tapeagents.config
 from omegaconf import DictConfig
 from pdf2image import convert_from_path
 from pydantic import ConfigDict, Field
@@ -79,16 +80,16 @@ def __init__(
         self.task_seed = task_seed
         self.env_config = env_config
 
-    def make_env(self, exp_dir: str | Path, action_mapping=None) -> GaiaGym:
-        exp_dir = str(exp_dir)
-        logger.info(f"Init gaia env with directory {exp_dir}")
-        os.environ["TAPEAGENTS_SQLITE_DB"] = os.path.join(exp_dir, "tapedata.sqlite")
-        init_code_sandbox(exp_dir)
+    def make_env(self, exp_dir: Path, action_mapping=None) -> GaiaGym:
+        tapeagents.config.DB_DEFAULT_FILENAME = str(exp_dir.parent / "tapedata.sqlite")
+        exp_dir_str = str(exp_dir)
+        logger.info(f"Init gaia env with directory {exp_dir_str}")
+        init_code_sandbox(exp_dir_str)
         for i in range(len(self.env_config.tools)):
             if hasattr(self.env_config.tools[i], "exp_path"):
-                self.env_config.tools[i].exp_path = exp_dir
+                self.env_config.tools[i].exp_path = exp_dir_str
         tools = hydra.utils.instantiate(self.env_config.tools)
-        env = GaiaGym(tools=tools, task=self.task, exp_dir=exp_dir)
+        env = GaiaGym(tools=tools, task=self.task, exp_dir=exp_dir_str)
         return env