Add categorize action script

Author: klieret

codeclash/analysis/llm_as_judge/categorize_actions.py

@@ -1,14 +1,154 @@
+#!/usr/bin/env python3
 import argparse
+import random
 from pathlib import Path
+from typing import Any, Literal
 
 import yaml
+from pydantic import BaseModel
+from rich.console import Console
+from rich.table import Table
+
+from codeclash.analysis.llm_as_judge.big_questions import BigQuestions
+from codeclash.analysis.llm_as_judge.utils import Instance, InstanceBatch, get_instances
+from codeclash.utils.log import get_logger
+
+logger = get_logger("ActionCategorizer", emoji="🏷️")
 
 config_path = Path(__file__).parent / "categorize_actions.yaml"
 
+# Base categories
+_read_subcategories = ["source", "logs", "docs", "other"]
+_read_subsubcategories = ["new", "old"]
+
+_write_subcategories = [
+    "docs",
+    "source.main",
+    "source.main.backup",
+    "source.opponent",
+    "source.analysis",
+    "source.tests",
+    "other",
+]
+_write_subsubcategories = ["create", "modify_old", "modify_new"]
+
+_execute_subcategories = ["game", "game.setup", "analysis", "unittest", "other"]
+_execute_subsubcategories = ["in_mem", "new", "old"]
+
+# Generate all category combinations
+_all_categories = (
+    ["search", "submit", "other"]
+    + [f"read.{sub}.{subsub}" for sub in _read_subcategories for subsub in _read_subsubcategories]
+    + [f"write.{sub}.{subsub}" for sub in _write_subcategories for subsub in _write_subsubcategories]
+    + [f"execute.{sub}.{subsub}" for sub in _execute_subcategories for subsub in _execute_subsubcategories]
+)
+
+
+class ActionCategoryResponse(BaseModel):
+    category: Literal[*_all_categories]
+    base_action: str
+    success: bool
+    notes: str = ""
+    target_paths: list[str] = []
+
+
+class ActionCategoriesModelResponse(BaseModel):
+    categories: list[ActionCategoryResponse]
+
+    def pretty_print(self) -> None:
+        console = Console()
+
+        table = Table()
+        table.add_column("#", style="dim", width=4)
+        table.add_column("Category", style="cyan", width=35)
+        table.add_column("Base Action", style="green", width=20)
+        table.add_column("Target Path", style="green", width=20)
+        table.add_column("Success", style="green", width=10)
+        table.add_column("Notes", style="yellow")
+
+        for i, action in enumerate(self.categories, 1):
+            success_str = "✓" if action.success else "✗"
+            table.add_row(
+                str(i),
+                action.category,
+                action.base_action,
+                ", ".join(action.target_paths),
+                success_str,
+                action.notes or "-",
+            )
+
+        console.print()
+        console.print(table)
+        console.print()
+
+
+class ModelConfig(BaseModel):
+    model_name: str
+    model_class: str | None = None
+    model_kwargs: dict[str, Any] = {}
+
+
+class ActionCategoriesConfig(BaseModel):
+    version: int
+    system_prompt: str
+    instance_prompt: str = "Categorize the following actions:\n\n{{ trajectory_message_str }}"
+    model: ModelConfig
+
+
+class ActionCategorizer(BigQuestions):
+    def __init__(self, config: ActionCategoriesConfig):
+        super().__init__(config)
+
+    @property
+    def data_id(self) -> str:
+        return f"action_categories_v{self.config.version}"
+
+    def evaluate(self, instance: Instance) -> None:
+        target_path = instance.trajectory_path.parent.parent.parent / "llm_as_judge.json"
+
+        if self._should_skip(target_path, instance):
+            logger.info(
+                f"Skipping instance {instance.instance_id} because it already exists in {target_path} under key {self.data_id}"
+            )
+            return
+
+        response = self.model.query(
+            messages=self._get_messages(instance), response_format=ActionCategoriesModelResponse
+        )
+        response_data = ActionCategoriesModelResponse.model_validate_json(response["content"])
+        response_data.pretty_print()
+        response_data_json = {
+            "result": response_data.model_dump(mode="json"),
+            "instance": instance.model_dump(mode="json"),
+        }
+
+        self._save_response(target_path, response_data_json, instance)
+        logger.info(f"Evaluated instance {instance.instance_id}. Saved to {target_path} with key {self.data_id}")
+
+
+def load_instances_from_path(path: Path) -> list[Instance]:
+    if path.is_file() and path.suffix == ".json":
+        logger.info(f"Loading instances from batch file: {path}")
+        batch = InstanceBatch.model_validate_json(path.read_text())
+        return batch.instances
+    logger.info(f"Loading instances from directory: {path}")
+    return get_instances(path)
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("input_dir", type=Path, help="Path to the input dir")
+    parser.add_argument(
+        "input_dir", type=Path, nargs="+", help="Path to the input dir(s) or to instance batch json files"
+    )
+    parser.add_argument("--shuffle", action="store_true", help="Shuffle instances before processing")
+    parser.add_argument("-n", "--n-workers", type=int, default=1, help="Number of parallel workers (default: 1)")
     args = parser.parse_args()
 
-    config = yaml.safe_load(config_path.read_text())
+    config = ActionCategoriesConfig.model_validate(yaml.safe_load(config_path.read_text()))
+    instances = []
+    for input_path in args.input_dir:
+        instances.extend(load_instances_from_path(input_path))
+    categorizer = ActionCategorizer(config)
+    if args.shuffle:
+        random.shuffle(instances)
+    categorizer.evaluate_bulk(instances, n_workers=args.n_workers)

codeclash/analysis/llm_as_judge/categorize_actions.yaml

@@ -1,15 +1,23 @@
 # Incrementing this version will cause previous submissions to be re-evaluated
-version: 1
+version: 2
 system_prompt: |
-  You are helping to analyze the actions of a LM agent.
+  You are helping to analyze the actions of a LM agent (summarily referred to as "trajectory").
 
   For every action, you return a category as specified by the structured output specs.
 
   # Categories
 
+  ## Search operations
+
+  - `search`: Commonly find or grep commands.
+
+  This does NOT include running more complicated analysis scripts on files. Rule of thumb: If it's a bash command, it's a read operation, if a python script is executed, it probably belongs in the execution category.
+
   ## Read operations
 
-  The model reads code, documentation, logs, or anything else. This includes operations like finding files, searching through files, etc. Commands include `ls`, `find`, `grep`, `cat`, `head`, `tail`, etc. This does NOT include running more complicated analysis scripts on files. Rule of thumb: If it's a bash command, it's a read operation, if a python script is executed, it probably belongs in the execution category.
+  The model reads code, documentation, logs, or anything else.
+  Commands include `ls`, `cat`, `head`, `tail`, etc.
+  This does NOT include running more complicated analysis scripts on files. Rule of thumb: If it's a bash command, it's a read operation, if a python script is executed, it probably belongs in the execution category.
 
   Categories
 
@@ -20,19 +28,38 @@ system_prompt: |
 
   depending on what is being read.
 
+  For every read operation, use the following subsubcategories:
+
+  - `x.new`: We read a script that was created in this trajectory, i.e., you have seen the creation of the script or it is created in the same action.
+  - `x.old`: We read a script that was created before any action you have seen, i.e., you have not seen the creation of the script and it was created before this trajectory started.
+
+  Example: `read.source.new` (we read something a source file that we created in this trajectory), `read.logs.old` (we read logs that we did not create in this trajectory)
+
   ## Write operations
 
   The model modifies files. Common commands include `cat ... > file`, `sed`, etc. Creating directories also falls into this category.
 
   Subcatgories:
 
   - `write.docs`: Documentation
-  - `write.source.main`: Writing of the main player code. Does NOT include writing simple bots to test again, but only editing the main player/agent/bot file (`main.py`, `player.py`, `robot.js`, `warrior.red`, `robots/custom/`)
+  - `write.source.main`: Writing of the main player code. Does NOT include writing simple bots to test again, but only editing the main player/agent/bot file (`main.py`, `player.py`, `robot.js`, `warrior.red`, `robots/custom/`). Copying the main file to a backup does NOT belong in this category but in `write.source.main.backup`. Only editing the main file that is actually used in the game belongs in this category.
+  - `write.source.main.backup`: Writing of backup files of the main player code.
   - `write.source.opponent`: Writing of opponents to test the main player/agent/bot against.
   - `write.source.analysis`: Writing of analysis scripts, especially to parse logs or analyze what is happening in the game
   - `write.source.tests`: Writing of unit test scripts. Unit tests are different from analsis, because they have a predefined, very clear pass or fail outcome (i.e., assert statements)
   - `write.other`:  This category should be very infrequent and only for write targets that are clearly not compatible with the others.
 
+  For every write operation, use the following subsubcategories:
+
+  - `x.modify_old`: Modification of old files, adding or removing lines, etc. Old means that it was not created during this trajectory, and you have NOT seen the creation of the file. Completely overwriting a file that you know existed before this trajectory (for example, because it was target of a successful read operation) also belongs in this category.
+  - `x.create`: New file creation. The file was created in this very action. This includes `cp` operations.
+  - `x.modify_new`: Modification of 'new' files. The file was created during this trajectory, i.e., you have seen the creation of the file in a previous action. It is now modified in this action.
+
+  For example,
+
+  - if the model writes a new file, the category should be `write.source.main.new`.
+  - if the model modifies a main player source file that was created before any action you have seen, the category should be `write.source.main.modify_old`.
+
   ## Execution operations
 
   Executions are anything that executes source files, especially executing analysis scripts, playing the game with different players, etc.
@@ -41,17 +68,82 @@ system_prompt: |
   - `execute.game.setup`: Preparations for running the game, for example if the player servers need to be started first, or the game needs to be compiled etc., this also belongs in this
   - `execute.analysis`: Executing analysis scripts (see previous notes on difference between unittests and analysis)
   - `execute.unittest`: Executing unittests or simple tests (import checks etc.). Compilation checks also fall into this category.
+  - `execute.other`: This category should be very infrequent and only for execution targets that are clearly not compatible with the others.
+    Note that this still should only be for execution of scripts or longer e.g., `python -c` commands, not just for simple bash commands (use `other` for that).
+
+  For each of these execution operations, use the following subsubcategories:
+
+  - `x.in_mem`: We are executing a script in memory, e.g., `python -c "print('hello')"`
+  - `x.new`: We execute a script that was created in this trajectory, i.e., you have seen the creation of the script or it is created in the same action.
+  - `x.old`: We execute a script that was created before any action you have seen, i.e., you have not seen the creation of the script.
 
   ## Other
 
   - `submit`: The player issues "MINI_SWE_AGENT_FINAL_OUTPUT", "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT" to finish the run. If this is combined with another action, categorize the other action instead. Only use this category if it's a standalone request to finish.
-  - `other`: This category should be very infrequent and only for write targets that are clearly not compatible with the others.
+  - `other`: This category should be very infrequent and only for write targets that are clearly not compatible with any other category.
+
+  # Category Priorities
+
+  In order of importance: execution is more important than writing is more important than reading.
+  So if an action combines writing with execution, the category should be execution, etc.
+  Within one of these three categories, use the best match for the category.
+
+  # Base actions
+
+  In addition to the category, you also return the base action that was executed.
+  This means the part of the command that describes the action best, but without any arguments.
+  E.g., for `cat file.txt`, the base action is `cat`.
+  For `cd /path/to/dir && python script.py`, the base action is `python`.
+  For `git commit -m "Fix bug"`, the base action is `git commit` (because `commit` is an important part of the command).
+  If you resolved the category based on the priority rules, use the base action that is most important for the category.
+  E.g., for `sed ... test.py && python test.py`, the base action is `python`
+  (because it's more important than `sed` for the execution category).
+
+  # `notes`
+
+  If you cannot categorize the action and put it into `read.other`, `write.other`, `execute.other`, or `other`,
+  you MUST (!) explain why in the `notes` field.
+  Otherwise, use this very sparringly.
+  For most cases, you should leave this empty, unless you are unsure about the category
+  (in that case, still categorize the action, but explain why in the notes field).
+
+  # Success
+
+  Fill in the success field to True if the action was successfully executed.
+  Fill in the success field to False if the action was not successfully executed.
+  For example, if the agent tried to replace some text in a file, but the file or text was not found, the success field should be False.
+  If you are unsure, set the success field to True.
+  If you set this field to False, you MUST (!) explain why in the `notes` field.
+
+  # Target paths
+
+  If the action has a target path, e.g., for a read or write or execute operation, fill in the target_paths field with the target path(s).
+  If there are multiple target paths, list them all.
+  If there is no target path, leave this field empty.
+
+  # Output schema
+
+  The schema of your response is given to you.
+
+  You return essentially a list of of action responses, where every response includes the following fields:
+  - category
+  - base_action
+  - success
+  - notes (optional, if you need to explain why you set the success field to False or why you chose an 'other' category)
+  - target_paths (optional, if the action has a target path, e.g., for a read or write or execute operation)
 
   # Important notes
 
   1. You MUST (!) categorize EVERY (!) action. Do NOT (!) skip any action.
   2. Every action MUST (!) be put into exactly (!) one (!) category.
   3. Your category MUST (!) be one of the list above.
+  4. If you are unsure, use the best match for the category.
 
-  In order of importance: execution is more important than writing is more important than reading.
-  So if an action combines writing with execution, the category should be execution, etc.
+instance_prompt: |
+  Categorize the following actions:
+
+  {{ trajectory_message_str }}
+
+model:
+  model_name: '@openai/gpt-5-mini'
+  model_kwargs: {}