Ref(analysis): refactor into class

Author: klieret

codeclash/analysis/llm_as_judge/big_questions.py

@@ -20,6 +20,8 @@
 
 
 class BigQuestionsModelResponseSchema(BaseModel):
+    """Schema for structured output of the model."""
+
     are_edits_motivated_by_logs: bool
     are_edits_motivated: bool
     are_edits_tested_with_simulations: bool
@@ -51,54 +53,81 @@ def extract_triple_backticks(text: str) -> str:
     return actions[0] if actions else ""
 
 
-def evaluate(instance: Instance, config: BigQuestionsConfig) -> None:
-    target_path = instance.trajectory_path.parent.parent.parent / "llm_as_judge.json"
-    data_id = f"big_questions_v{config.version}_{config.model.model_name}_{instance.instance_id}"
-    if target_path.exists():
-        content = target_path.read_text()
-        if content.strip():
-            data = json.loads(content)
-            if data.get("data_id") == data_id:
-                logger.info(
-                    f"Skipping instance {instance.instance_id} because it already exists in {target_path} with key {data_id}"
-                )
-                return
-    model = get_model(config.model.model_name, config={"model_kwargs": config.model.model_kwargs})
-    trajectory_messages = json.loads(instance.trajectory_path.read_text())["messages"]
-    system_message = config.system_prompt
-    trajectory_message_str = ""
-    for message in trajectory_messages:
-        content = message["content"]
-        if isinstance(message["content"], list):
-            assert len(message["content"]) == 1
-            content = message["content"][0]["text"]
-        if message["role"] == "assistant":
-            trajectory_message_str += "\n<action>\n" + extract_triple_backticks(content) + "\n</action>\n"
-        elif message["role"] == "user":
-            trajectory_message_str += content  # already enclosed in <output>
-    instance_message = jinja2.Template(config.instance_prompt).render(trajectory_message_str=trajectory_message_str)
-    # print(instance_message)
-    response = model.query(
-        messages=[
+class BigQuestions:
+    def __init__(self, config: BigQuestionsConfig):
+        self.config = config
+        self.model = get_model(config.model.model_name, config={"model_kwargs": config.model.model_kwargs})
+
+    def get_data_id(self, instance: Instance) -> str:
+        return f"big_questions_v{self.config.version}_{self.config.model.model_name}_{instance.instance_id}"
+
+    def evaluate(self, instance: Instance) -> None:
+        target_path = instance.trajectory_path.parent.parent.parent / "llm_as_judge.json"
+        data_id = self.get_data_id(instance)
+
+        if self._should_skip(target_path, data_id):
+            logger.info(
+                f"Skipping instance {instance.instance_id} because it already exists in {target_path} with key {data_id}"
+            )
+            return
+
+        response = self.model.query(
+            messages=self._get_messages(instance), response_format=BigQuestionsModelResponseSchema
+        )
+        response_data = BigQuestionsModelResponseSchema.model_validate_json(response["content"]).model_dump()
+
+        self._save_response(target_path, response_data, data_id)
+        logger.info(
+            f"Evaluated instance {instance.instance_id}: {response_data}. Saved to {target_path} with key {data_id}"
+        )
+
+    def _format_traj_str(self, messages: list[dict[str, Any]]) -> str:
+        trajectory_message_str = ""
+        for message in messages:
+            content = message["content"]
+            if isinstance(message["content"], list):
+                assert len(message["content"]) == 1
+                content = message["content"][0]["text"]
+            if message["role"] == "assistant":
+                trajectory_message_str += "\n<action>\n" + extract_triple_backticks(content) + "\n</action>\n"
+            elif message["role"] == "user":
+                trajectory_message_str += content  # already enclosed in <output>
+        return trajectory_message_str
+
+    def _get_messages(self, instance: Instance) -> list[dict[str, Any]]:
+        trajectory_messages = json.loads(instance.trajectory_path.read_text())["messages"]
+        system_message = self.config.system_prompt
+        instance_message = jinja2.Template(self.config.instance_prompt).render(
+            trajectory_message_str=self._format_traj_str(trajectory_messages)
+        )
+        # print(instance_message)
+        return [
             {"role": "system", "content": system_message},
             {"role": "user", "content": instance_message},
-        ],
-        response_format=BigQuestionsModelResponseSchema,
-    )
-    # atomic write with file lock in case other analyses are also writing
-    response_data = BigQuestionsModelResponseSchema.model_validate_json(response["content"]).model_dump()
-    with FileLock(target_path):
-        # read again if changed in the meantime
-        data = {}
-        if target_path.exists():
-            content = target_path.read_text()
-            if content.strip():
-                data = json.loads(content)
-        data[data_id] = response_data
-        target_path.write_text(json.dumps(data))
-    logger.info(
-        f"Evaluated instance {instance.instance_id}: {response_data}. Saved to {target_path} with key {data_id}"
-    )
+        ]
+
+    def _should_skip(self, target_path: Path, data_id: str) -> bool:
+        if not target_path.exists():
+            return False
+        content = target_path.read_text()
+        if not content.strip():
+            return False
+        data = json.loads(content)
+        if data.get("data_id") == data_id:
+            return True
+        return False
+
+    def _save_response(self, target_path: Path, response_data: dict[str, Any], data_id: str) -> None:
+        # atomic write with file lock in case other analyses are also writing
+        with FileLock(target_path):
+            # read again if changed in the meantime
+            data = {}
+            if target_path.exists():
+                content = target_path.read_text()
+                if content.strip():
+                    data = json.loads(content)
+            data[data_id] = response_data
+            target_path.write_text(json.dumps(data))
 
 
 if __name__ == "__main__":
@@ -108,6 +137,7 @@ def evaluate(instance: Instance, config: BigQuestionsConfig) -> None:
 
     config = BigQuestionsConfig.model_validate(yaml.safe_load(config_path.read_text()))
     instances = get_instances(args.input_dir)
+    big_questions = BigQuestions(config)
     for instance in instances:
-        evaluate(instance, config)
+        big_questions.evaluate(instance)
         print(GLOBAL_MODEL_STATS.cost)

codeclash/analysis/llm_as_judge/big_questions.yaml

@@ -3,7 +3,7 @@ version: 4
 system_prompt: |
   You need to answer three question
 
-  ## Are edits motivated by logs?
+  ## Q1 (`are_edits_motivated_by_logs`): Are edits motivated by logs?
 
   Are the edits of the main player file of the player directly motivated by problems discovered by reading the logs?
 
@@ -39,7 +39,7 @@ system_prompt: |
   - Player 2 is better most of the time (why?)
   - Player 1 is the last bot standing and is therefore the winner (does not explain why player 2 lost)
 
-  ## Are edits motivated
+  ## Q2 (`are_edits_motivated`): Are edits motivated
 
   Can the goal of the specific edits on the main player file be motivated by any output of previous actions?
   If you answered True to the previous question, answer True here as well.
@@ -63,7 +63,7 @@ system_prompt: |
   You should answer False, if the edits seemed unrelated to any output of previous actions before the relevant
   edit actions.
 
-  ## Are edits tested with simulations?
+  ## Q3 (`are_edits_tested_with_simulations`): Are edits tested with simulations?
 
   After edits to the main player file have been made, does the player simulate games by playing the game against previous versions of itself, or against example players, etc.
   Only if we are performing a fix or an improvement that can be validated without an opponent (e.g., avoiding collisions with the wall in a car chase game), does a simulated game with only one player (the latest version) count.
@@ -77,7 +77,7 @@ system_prompt: |
 
   1. If no edits to the main player file have been made, answer `True`
 
-  ## Are edits validated with unittests?
+  ## Q4 (`are_edits_validated_with_unittests`): Are edits validated with unittests?
 
   Are the final edits covered by specific unittests that test the new or modified behavior?
 
@@ -95,7 +95,7 @@ system_prompt: |
 
   Note: if the tests did not run, or showed that the new version was broken, answer False.
 
-  ## Categorize the kind of edits made
+  ## Q5 (`edit_category`): Categorize the kind of edits made
 
   Categorize the final edits to the main file into one of the following categories.
   Ignore comments or documentation.
@@ -109,7 +109,8 @@ system_prompt: |
 
   ## Output format
 
-  Answer in the json format specified. Reasoning corresponds to your explanation for your answer.
+  Answer in the json format specified.
+  The `reasoning` field should contain an explanation for your answer that explains your reasoning for each of the answers.
 instance_prompt: |
   Here is your input file: