phony command, joblib stuff, took think out of prompt

Author: TLSDC

pyproject.toml

@@ -57,3 +57,4 @@ exclude = '''
 [project.scripts]
 agentlab-assistant = "agentlab.ui_assistant:main"
 agentlab-xray = "agentlab.analyze.agent_xray:main"
+agentlab-analyze = "agentlab.analyze.error_analysis.pipeline:main"

src/agentlab/analyze/error_analysis/pipeline.py

@@ -23,6 +23,11 @@ def __call__(self, *args, **kwds):
         return "analysis"
 
 
+def analyze(exp_result, episode_summarizer, save_analysis_func):
+    error_analysis = episode_summarizer(exp_result)
+    save_analysis_func(exp_result, error_analysis)
+
+
 @dataclass
 class ErrorAnalysisPipeline:
     exp_dir: Path
@@ -36,12 +41,21 @@ def filter_exp_results(self) -> Generator[ExpResult, None, None]:
             if self.filter is None or self.filter in str(exp_result.exp_dir):
                 yield exp_result
 
-    def run_analysis(self):
+    def run_analysis(self, parallel=False, jobs=-1):
         filtered_results = self.filter_exp_results()
 
-        for exp_result in filtered_results:
-            error_analysis = self.episode_summarizer(exp_result)
-            self.save_analysis(exp_result, error_analysis)
+        if parallel:
+            import joblib
+
+            joblib.Parallel(n_jobs=jobs, backend="threading")(
+                joblib.delayed(analyze)(exp_result, self.episode_summarizer, self.save_analysis)
+                for exp_result in filtered_results
+            )
+
+        else:
+            for exp_result in filtered_results:
+                error_analysis = self.episode_summarizer(exp_result)
+                self.save_analysis(exp_result, error_analysis)
 
     def save_analysis(self, exp_result: ExpResult, error_analysis: dict, exists_ok=True):
         """Save the analysis to json"""
@@ -56,28 +70,37 @@ def save_analysis(self, exp_result: ExpResult, error_analysis: dict, exists_ok=T
 HTML_FORMATTER = lambda x: x.get("pruned_html", "No HTML available")
 
 
-if __name__ == "__main__":
+def main():
     import argparse
 
     parser = argparse.ArgumentParser()
     parser.add_argument("-e", "--exp_dir", type=str)
     parser.add_argument("-f", "--filter", type=str, default=None)
+    parser.add_argument("-p", "--parallel", action="store_true")
+    parser.add_argument("-j", "--jobs", type=int, default=-1)
 
     args = parser.parse_args()
+
+    assert args.exp_dir is not None, "Please provide an exp_dir, e.g., -e /path/to/exp_dir"
+
     exp_dir = Path(args.exp_dir)
     filter = args.filter
+    parallel = args.parallel
+    jobs = args.jobs
 
     from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT
 
     llm = CHAT_MODEL_ARGS_DICT["azure/gpt-4o-2024-08-06"].make_model()
 
-    step_summarizer = ChangeSummarizer(llm, lambda x: x)
-    episode_summarizer = EpisodeSummarizer()
-
     pipeline = ErrorAnalysisPipeline(
         exp_dir=exp_dir,
         filter=filter,
         episode_summarizer=EpisodeErrorSummarizer(ChangeSummarizer(llm, AXTREE_FORMATTER), llm),
     )
 
-    pipeline.run_analysis()
+    pipeline.run_analysis(parallel=parallel, jobs=jobs)
+
+
+if __name__ == "__main__":
+
+    main()

src/agentlab/analyze/error_analysis/summarizer.py

@@ -6,8 +6,8 @@
     CHANGE_SUMMARIZER_PROMPT,
     ERROR_CLASSIFICATION_PROMPT,
 )
-from agentlab.analyze.inspect_results import summarize
 from agentlab.llm.llm_utils import json_parser, parse_html_tags
+from agentlab.llm.tracking import set_tracker
 
 
 def _diff(past_obs, current_obs):
@@ -94,14 +94,20 @@ def __call__(self, exp_results: ExpResult) -> EpisodeAnalysis:
         # if exp_results.steps_info[-1].reward == 1:
         #     return {"analysis": "Success", "summaries": {}}
 
-        summaries = self.make_change_summaries(exp_results)
+        with set_tracker("summary") as summaries_tracker:
+            summaries = self.make_change_summaries(exp_results)
         prompt = self.make_prompt(exp_results, summaries)
-        raw_analysis = self.llm(prompt)["content"]
+
+        with set_tracker("analysis") as analysis_tracker:
+            raw_analysis = self.llm(prompt)["content"]
         analysis = self.parse(raw_analysis)
-        return {
+        res = {
             "analysis": analysis,
             "summaries": {i: a for i, a in enumerate(summaries)},
         }
+        res.update(analysis_tracker.stats)
+        res.update(summaries_tracker.stats)
+        return res
 
     def make_change_summaries(self, exp_result: ExpResult) -> list[str]:
         summaries = []  # type: list[str]
@@ -136,16 +142,15 @@ def format_summary(summary):
 
         txt_summaries = "\n".join([format_summary(summary) for summary in summaries])
 
-        thoughts = [step.agent_info.think for step in exp_results.steps_info[:-1]]
         actions = [step.action for step in exp_results.steps_info[:-1]]
         action_errors = "\n".join(
             [step.obs["last_action_error"] for step in exp_results.steps_info[1:]]
         )
 
         txt_actions = "\n".join(
             [
-                f"Thoughts: {thought}\nAction: {action}\nAction Error: {action_error}"
-                for action, thought, action_error in zip(actions, thoughts, action_errors)
+                f"Action: {action}\nAction Error: {action_error}"
+                for action, action_error in zip(actions, action_errors)
             ]
         )
         return ERROR_CLASSIFICATION_PROMPT.format(