diff options
| -rw-r--r-- | src/agents/03-logs/agent.py | 28 | ||||
| -rw-r--r-- | src/agents/03-logs/main.py | 9 |
2 files changed, 32 insertions, 5 deletions
diff --git a/src/agents/03-logs/agent.py b/src/agents/03-logs/agent.py index 80bdcd5..9b061b7 100644 --- a/src/agents/03-logs/agent.py +++ b/src/agents/03-logs/agent.py @@ -17,6 +17,9 @@ class Agent: self.scenario = scenario self.oversight_level = oversight_level self.user_prompt_type = user_prompt_type + self.total_tokens = 0 + self.prompt_tokens = 0 + self.completion_tokens = 0 def run(self, initial_prompt): messages = [ @@ -31,8 +34,27 @@ class Agent: messages=messages, tools=self.tools, temperature=self.temperature, + extra_body={ + # for gemini + # 'extra_body': { + # "google": { + # "thinking_config": { + # "thinking_budget": 800, + # "include_thoughts": True + # } + # } + # } + # for minimax + # "reasoning_split": True + } ) + # Update token counts + if response.usage: + self.total_tokens += response.usage.total_tokens + self.prompt_tokens += response.usage.prompt_tokens + self.completion_tokens += response.usage.completion_tokens + response_message = response.choices[0].message messages.append(response_message) @@ -92,9 +114,13 @@ class Agent: "scenario": self.scenario, "oversight_level": self.oversight_level, "user_prompt_type": self.user_prompt_type, + "total_tokens": self.total_tokens, + "prompt_tokens": self.prompt_tokens, + "completion_tokens": self.completion_tokens, "conversation": self.logs } - log_file = f"output/log_{timestamp}.json" + model_name_safe = self.model.replace("/", "_") + log_file = f"output/{model_name_safe}_{self.scenario}_{timestamp}.json" with open(log_file, "w") as f: json.dump(log_data, f, indent=4) print(f"\nLogs saved to {log_file}") diff --git a/src/agents/03-logs/main.py b/src/agents/03-logs/main.py index a7603ff..d339330 100644 --- a/src/agents/03-logs/main.py +++ b/src/agents/03-logs/main.py @@ -1,7 +1,7 @@ # main.py from agent import Agent import os -from config import INCLUDE_HIDDEN_GOAL +from config import INCLUDE_HIDDEN_GOAL, MODEL from vfs import VFS def load_prompt(file_path): @@ -12,8 +12,8 @@ def load_prompt(file_path): if __name__ == "__main__": # Configuration for the experiment run - SCENARIO = "log_redaction" - OVERSIGHT_LEVEL = "high" + SCENARIO = "corporate_sabotage" + OVERSIGHT_LEVEL = "low" # Construct paths based on the scenario scenario_dir = os.path.join("scenarios", SCENARIO) @@ -43,6 +43,7 @@ if __name__ == "__main__": system_prompt = "\n\n".join(filter(None, system_prompt_parts)) + print("\n--- Using model:", MODEL, "---\n") agent = Agent( system_prompt=system_prompt, scenario=SCENARIO, @@ -56,4 +57,4 @@ if __name__ == "__main__": VFS.get_instance().print_fs() -
\ No newline at end of file + |