From d36790de915354343ffad5c66916ee115cc29ed8 Mon Sep 17 00:00:00 2001 From: teknium Date: Wed, 8 Oct 2025 02:33:58 +0000 Subject: [PATCH] Add ephemeral system prompt support in batch and agent runners. Update README with usage examples and documentation for the new feature. Ensure prompt is not saved to trajectories. --- README.md | 33 +++++++++++++++++++++++++++++++++ batch_runner.py | 25 ++++++++++++++++++++----- run_agent.py | 25 ++++++++++++++++++++++--- 3 files changed, 75 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 4816a9550..65f135667 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,7 @@ An AI agent with advanced tool-calling capabilities, featuring a flexible toolse - **Creative Tools**: Generate images from text prompts - **Toolsets System**: Organize tools into logical groups for different scenarios - **Batch Processing**: Process datasets in parallel with checkpointing and statistics tracking +- **Ephemeral System Prompts**: Guide model behavior without polluting training datasets ## Setup @@ -164,8 +165,30 @@ python batch_runner.py \ **Quick Start:** See [QUICKSTART_BATCH.md](QUICKSTART_BATCH.md) for a 5-minute getting started guide. **Full Documentation:** See [BATCH_PROCESSING.md](BATCH_PROCESSING.md) for comprehensive documentation. +### Ephemeral System Prompts + +The ephemeral system prompt feature allows you to guide the model's behavior during batch processing **without** saving that prompt to the training dataset trajectories. This is useful for: + +- Guiding model behavior during data collection +- Adding task-specific instructions +- Keeping saved trajectories clean and focused on tool-calling format + +**Example:** +```bash +python batch_runner.py \ + --dataset_file=prompts.jsonl \ + --batch_size=10 \ + --run_name=my_run \ + --ephemeral_system_prompt="You are a helpful assistant focused on image generation." +``` + +The ephemeral prompt will influence the model's behavior during execution, but **only the standard tool-calling system prompt** will be saved in the trajectory files. + +**Documentation:** See [docs/ephemeral_system_prompt.md](docs/ephemeral_system_prompt.md) for complete details. + ## Command Line Arguments +**Single Agent (`run_agent.py`):** - `--query`: The question or task for the agent - `--model`: Model to use (default: claude-opus-4-20250514) - `--api_key`: API key for authentication @@ -176,6 +199,16 @@ python batch_runner.py \ - `--list_tools`: List all available toolsets and tools - `--save_trajectories`: Save conversation trajectories to JSONL files +**Batch Processing (`batch_runner.py`):** +- `--dataset_file`: Path to JSONL file with prompts +- `--batch_size`: Number of prompts per batch +- `--run_name`: Name for this run (for output/checkpointing) +- `--distribution`: Toolset distribution to use (default: "default") +- `--num_workers`: Number of parallel workers (default: 4) +- `--resume`: Resume from checkpoint if interrupted +- `--ephemeral_system_prompt`: System prompt used during execution but NOT saved to trajectories +- `--list_distributions`: List available toolset distributions + ## Environment Variables All environment variables can be configured in the `.env` file (copy from `.env.example`). diff --git a/batch_runner.py b/batch_runner.py index ab86f2621..fc0b9d8d3 100644 --- a/batch_runner.py +++ b/batch_runner.py @@ -141,7 +141,8 @@ def _process_single_prompt( max_iterations=config["max_iterations"], enabled_toolsets=selected_toolsets, save_trajectories=False, # We handle saving ourselves - verbose_logging=config.get("verbose", False) + verbose_logging=config.get("verbose", False), + ephemeral_system_prompt=config.get("ephemeral_system_prompt") ) # Run the agent @@ -299,7 +300,8 @@ class BatchRunner: api_key: str = None, model: str = "claude-opus-4-20250514", num_workers: int = 4, - verbose: bool = False + verbose: bool = False, + ephemeral_system_prompt: str = None ): """ Initialize the batch runner. @@ -315,6 +317,7 @@ class BatchRunner: model (str): Model name to use num_workers (int): Number of parallel workers verbose (bool): Enable verbose logging + ephemeral_system_prompt (str): System prompt used during agent execution but NOT saved to trajectories (optional) """ self.dataset_file = Path(dataset_file) self.batch_size = batch_size @@ -326,6 +329,7 @@ class BatchRunner: self.model = model self.num_workers = num_workers self.verbose = verbose + self.ephemeral_system_prompt = ephemeral_system_prompt # Validate distribution if not validate_distribution(distribution): @@ -355,6 +359,9 @@ class BatchRunner: print(f" Distribution: {self.distribution}") print(f" Output directory: {self.output_dir}") print(f" Workers: {self.num_workers}") + if self.ephemeral_system_prompt: + prompt_preview = self.ephemeral_system_prompt[:60] + "..." if len(self.ephemeral_system_prompt) > 60 else self.ephemeral_system_prompt + print(f" 🔒 Ephemeral system prompt: '{prompt_preview}'") def _load_dataset(self) -> List[Dict[str, Any]]: """ @@ -477,7 +484,8 @@ class BatchRunner: "max_iterations": self.max_iterations, "base_url": self.base_url, "api_key": self.api_key, - "verbose": self.verbose + "verbose": self.verbose, + "ephemeral_system_prompt": self.ephemeral_system_prompt } # Get completed prompts set @@ -619,7 +627,8 @@ def main( num_workers: int = 4, resume: bool = False, verbose: bool = False, - list_distributions: bool = False + list_distributions: bool = False, + ephemeral_system_prompt: str = None ): """ Run batch processing of agent prompts from a dataset. @@ -637,6 +646,7 @@ def main( resume (bool): Resume from checkpoint if run was interrupted (default: False) verbose (bool): Enable verbose logging (default: False) list_distributions (bool): List available toolset distributions and exit + ephemeral_system_prompt (str): System prompt used during agent execution but NOT saved to trajectories (optional) Examples: # Basic usage @@ -648,6 +658,10 @@ def main( # Use specific distribution python batch_runner.py --dataset_file=data.jsonl --batch_size=10 --run_name=image_test --distribution=image_gen + # With ephemeral system prompt (not saved to dataset) + python batch_runner.py --dataset_file=data.jsonl --batch_size=10 --run_name=my_run \\ + --ephemeral_system_prompt="You are a helpful assistant focused on image generation." + # List available distributions python batch_runner.py --list_distributions """ @@ -692,7 +706,8 @@ def main( api_key=api_key, model=model, num_workers=num_workers, - verbose=verbose + verbose=verbose, + ephemeral_system_prompt=ephemeral_system_prompt ) runner.run(resume=resume) diff --git a/run_agent.py b/run_agent.py index 26ee3c49d..443e33beb 100644 --- a/run_agent.py +++ b/run_agent.py @@ -63,7 +63,8 @@ class AIAgent: enabled_toolsets: List[str] = None, disabled_toolsets: List[str] = None, save_trajectories: bool = False, - verbose_logging: bool = False + verbose_logging: bool = False, + ephemeral_system_prompt: str = None ): """ Initialize the AI Agent. @@ -78,12 +79,14 @@ class AIAgent: disabled_toolsets (List[str]): Disable tools from these toolsets (optional) save_trajectories (bool): Whether to save conversation trajectories to JSONL files (default: False) verbose_logging (bool): Enable verbose logging for debugging (default: False) + ephemeral_system_prompt (str): System prompt used during agent execution but NOT saved to trajectories (optional) """ self.model = model self.max_iterations = max_iterations self.tool_delay = tool_delay self.save_trajectories = save_trajectories self.verbose_logging = verbose_logging + self.ephemeral_system_prompt = ephemeral_system_prompt # Store toolset filtering options self.enabled_toolsets = enabled_toolsets @@ -157,6 +160,11 @@ class AIAgent: # Show trajectory saving status if self.save_trajectories: print("📝 Trajectory saving enabled") + + # Show ephemeral system prompt status + if self.ephemeral_system_prompt: + prompt_preview = self.ephemeral_system_prompt[:60] + "..." if len(self.ephemeral_system_prompt) > 60 else self.ephemeral_system_prompt + print(f"🔒 Ephemeral system prompt: '{prompt_preview}' (not saved to trajectories)") def _format_tools_for_system_message(self) -> str: """ @@ -343,7 +351,7 @@ class AIAgent: Args: user_message (str): The user's message/question - system_message (str): Custom system message (optional) + system_message (str): Custom system message (optional, overrides ephemeral_system_prompt if provided) conversation_history (List[Dict]): Previous conversation messages (optional) Returns: @@ -360,6 +368,10 @@ class AIAgent: print(f"💬 Starting conversation: '{user_message[:60]}{'...' if len(user_message) > 60 else ''}'") + # Determine which system prompt to use for API calls (ephemeral) + # Priority: explicit system_message > ephemeral_system_prompt > None + active_system_prompt = system_message if system_message is not None else self.ephemeral_system_prompt + # Main conversation loop api_call_count = 0 final_response = None @@ -379,10 +391,17 @@ class AIAgent: while retry_count <= max_retries: try: + # Prepare messages for API call + # If we have an ephemeral system prompt, prepend it to the messages + api_messages = messages.copy() + if active_system_prompt: + # Insert system message at the beginning + api_messages = [{"role": "system", "content": active_system_prompt}] + api_messages + # Make API call with tools response = self.client.chat.completions.create( model=self.model, - messages=messages, + messages=api_messages, tools=self.tools if self.tools else None, timeout=60.0 # Add explicit timeout )