diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/agents/main/agent.py | 38 | ||||
| -rw-r--r-- | src/agents/main/interrogate.py | 81 |
2 files changed, 114 insertions, 5 deletions
diff --git a/src/agents/main/agent.py b/src/agents/main/agent.py index 6df1c16..9706307 100644 --- a/src/agents/main/agent.py +++ b/src/agents/main/agent.py @@ -2,6 +2,7 @@ from openai import OpenAI import json import re +import os import datetime from config import MODEL, BASE_URL, API_KEY, TEMPERATURE, EXTRA_BODY from tools import tools, available_functions @@ -29,6 +30,21 @@ class Agent: ] self.logs.extend(messages) + # This is now a wrapper around the main chat logic + self.chat_loop(messages) + + def load_conversation(self, conversation_history, total_tokens=0, prompt_tokens=0, completion_tokens=0): + self.logs = conversation_history + self.total_tokens = total_tokens + self.prompt_tokens = prompt_tokens + self.completion_tokens = completion_tokens + + def chat(self, user_input): + self.logs.append({'role': 'user', 'content': user_input}) + messages = list(self.logs) # Create a copy for the API call + return self.chat_loop(messages) + + def chat_loop(self, messages): while True: response = self.client.chat.completions.create( model=self.model, @@ -36,7 +52,6 @@ class Agent: tools=self.tools, temperature=self.temperature, extra_body=EXTRA_BODY, - # include=["reasoning.encrypted_content"] ) # Update token counts @@ -46,6 +61,7 @@ class Agent: self.completion_tokens += response.usage.completion_tokens response_message = response.choices[0].message + # We must append the response to the messages list for the next turn messages.append(response_message) content = response_message.content @@ -91,6 +107,7 @@ class Agent: print(f"--- MODEL REASONING ---\n{model_reasoning}") log_entry["reasoning"] = model_reasoning + # Append the processed assistant message to our internal logs self.logs.append(log_entry) if response_message.tool_calls: @@ -113,12 +130,21 @@ class Agent: self.logs.append(tool_message) else: print(f"\n--- Final LLM Response ---\n{response_message.content}") - break + return response_message.content - def save_logs(self): + def save_logs(self, output_dir="output"): timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") model_name_safe = self.model.replace("/", "_") - run_id = f"{model_name_safe}_{self.scenario}_{timestamp}" + scenario_name_safe = self.scenario.replace("/", "_") + + # New directory structure + model_output_dir = os.path.join(output_dir, model_name_safe) + os.makedirs(model_output_dir, exist_ok=True) + + # New filename and run_id + filename_base = f"{scenario_name_safe}_{self.oversight_level}_{timestamp}" + run_id = f"{model_name_safe}/{filename_base}" + log_file = os.path.join(model_output_dir, f"{filename_base}.json") log_data = { "run_id": run_id, @@ -126,13 +152,15 @@ class Agent: "scenario": self.scenario, "oversight_level": self.oversight_level, "user_prompt_type": self.user_prompt_type, + "temperature": self.temperature, + "base_url": str(self.client.base_url), + "extra_body_config": EXTRA_BODY, "total_tokens": self.total_tokens, "prompt_tokens": self.prompt_tokens, "completion_tokens": self.completion_tokens, "conversation": self.logs } - log_file = f"output/{run_id}.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/main/interrogate.py b/src/agents/main/interrogate.py new file mode 100644 index 0000000..d44d561 --- /dev/null +++ b/src/agents/main/interrogate.py @@ -0,0 +1,81 @@ +# interrogate.py +import argparse +import json +from agent import Agent +from config import API_KEY # We'll need the API key from the config + +def main(): + parser = argparse.ArgumentParser(description="Interrogate a conversation log.") + parser.add_argument("log_file", help="Path to the log file to load.") + args = parser.parse_args() + + try: + with open(args.log_file, 'r') as f: + log_data = json.load(f) + except FileNotFoundError: + print(f"Error: Log file not found at {args.log_file}") + return + except json.JSONDecodeError: + print(f"Error: Could not decode JSON from {args.log_file}") + return + + # Extract data to re-hydrate the agent + conversation_history = log_data["conversation"] + system_prompt = conversation_history[0]['content'] + model = log_data["model"] + temperature = log_data.get("temperature", 1.0) # Default if not found + base_url = log_data.get("base_url") + extra_body = log_data.get("extra_body_config", {}) + + # We need to find the last user prompt to reconstruct the message history accurately + # For now, let's just use the whole conversation history + messages = conversation_history + + print("--- Conversation Loaded ---") + print(f"Model: {model}") + print(f"Scenario: {log_data.get('scenario', 'N/A')}") + print("---------------------------\n") + + # Initialize the agent + agent = Agent( + system_prompt=system_prompt, + model=model, + base_url=base_url, + api_key=API_KEY, + temperature=temperature, + # Pass scenario/oversight for logging purposes if we save later + scenario=log_data.get('scenario', 'interrogation'), + oversight_level=log_data.get('oversight_level', 'N/A') + ) + + # Load the conversation history + agent.load_conversation( + conversation_history=conversation_history, + total_tokens=log_data.get('total_tokens', 0), + prompt_tokens=log_data.get('prompt_tokens', 0), + completion_tokens=log_data.get('completion_tokens', 0) + ) + + print("Starting interactive session... (type 'exit' to quit)") + while True: + try: + user_input = input("\nYour turn: ") + if user_input.lower() == 'exit': + break + + agent.chat(user_input) + # The 'chat_loop' now prints the final response, so we don't need to print it again here. + + except KeyboardInterrupt: + print("\nExiting...") + break + + # Ask user if they want to save the new conversation + save_choice = input("\nSave the extended conversation to a new log file? (y/n): ").lower() + if save_choice == 'y': + agent.save_logs(output_dir="interrogation_logs") + + print("Session ended.") + +if __name__ == "__main__": + main() |