diff --git a/.gitignore b/.gitignore index dc976f1ef..819ccbdef 100644 --- a/.gitignore +++ b/.gitignore @@ -213,3 +213,6 @@ cython_debug/ marimo/_static/ marimo/_lsp/ __marimo__/ + +## Refactoring Example Log Folder +examples/refactor/sample-working/rlm_logs/ diff --git a/examples/bedrock_example.py b/examples/bedrock_example.py index 6a6e7254c..3450d3998 100644 --- a/examples/bedrock_example.py +++ b/examples/bedrock_example.py @@ -16,7 +16,7 @@ "model_name": "us.anthropic.claude-3-5-sonnet-20241022-v2:0", "region_name": "us-east-1", }, - environment="local", + environment="docker", max_root_tokens=80000, # Session-wide cumulative limit max_depth=1, logger=logger, diff --git a/examples/refactor/README.md b/examples/refactor/README.md new file mode 100644 index 000000000..ccfadecbb --- /dev/null +++ b/examples/refactor/README.md @@ -0,0 +1,483 @@ +# RLM Refactor - Interactive Code Refactoring + +Use the built-in RLM refactoring CLI for code analysis and transformation tasks with YAML-based configuration and Docker isolation. + +## Quick Start + +### Try the Sample Example + +Run the included sample that searches for "bob" across multiple repositories: + +```bash +# From the repository root +uv run python -m rlm.refactor.cli --working-dir examples/refactor/sample-working + +# Or interactive mode (will prompt for directory): +uv run python -m rlm.refactor.cli +# Enter: examples/refactor/sample-working +``` + +That's it! The sample includes: +- `rlm-instructions.md` - Task: Find which file contains "bob" across all repos +- `rlm-config.yaml` - Bedrock configuration with Sonnet (root) + Haiku (sub) multi-tier setup +- Three git repositories with 10 text files total (one contains "bob") + - `sample-repo-one/` - 2 files + - `sample-repo-two/` - 3 files + - `sample-repo-three/` - 5 files (contains the target) + +**Expected output**: The RLM will search across all repositories and report that `sample-repo-three/file_05.txt` contains "bob". + +**Prerequisites**: +- AWS credentials configured (for Bedrock backend) +- Docker installed and running (for safe code execution isolation) +- Run `uv sync` to install dependencies + +### Create Your Own Task + +1. **Create a working directory** with your task: + ```bash + mkdir my-refactor-task + cd my-refactor-task + ``` + +2. **Create required files**: + + **`rlm-instructions.md`** (task description): + ```markdown + # Task: Find Bob + + Search through all `.txt` files and identify which file contains the word "bob". + Return only the filename. + ``` + + **`rlm-config.yaml`** (configuration): + ```yaml + backend: bedrock + backend_kwargs: + model_name: us.anthropic.claude-3-5-sonnet-20241022-v2:0 + region_name: us-east-1 + + max_root_tokens: 50000 + max_sub_tokens: 25000 + ``` + +3. **Run the refactoring CLI**: + ```bash + # From repository root + uv run python -m rlm.refactor.cli --working-dir ./my-refactor-task + + # Or use interactive mode + uv run python -m rlm.refactor.cli + # Enter: ./my-refactor-task + ``` + +## Configuration Requirements + +Both `rlm-instructions.md` and `rlm-config.yaml` are **required** in your working directory. + +### Required: `rlm-config.yaml` + +Minimal configuration: + +```yaml +# Minimal configuration +backend: bedrock +backend_kwargs: + model_name: us.anthropic.claude-3-5-sonnet-20241022-v2:0 + region_name: us-east-1 + max_tokens: 8000 +``` + +See [rlm-config.yaml](./rlm-config.yaml) for a complete template with all options documented. + +### Configuration Reference + +#### Required Fields + +| Field | Type | Description | Example | +|-------|------|-------------|---------| +| `backend` | string | LM provider | `bedrock`, `openai`, `anthropic` | +| `backend_kwargs.model_name` | string | Model identifier | `us.anthropic.claude-3-5-sonnet-20241022-v2:0` | + +#### Common Optional Fields + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `backend_kwargs.max_tokens` | int | 8000 | Max tokens per response (≥100) | +| `backend_kwargs.temperature` | float | 1.0 | Sampling temperature (0.0-2.0) | +| `backend_kwargs.region_name` | string | `us-east-1` | AWS region (Bedrock only) | +| `max_root_tokens` | int | 100000 | Session-wide token budget | +| `max_iterations` | int | 50 | Max REPL iterations | +| `environment` | string | `local` | Execution environment | +| `verbose` | bool | true | Enable rich console output | +| `log_dir` | string | `rlm-logs` | Log directory (relative path) | + +#### Multi-Tier Configuration (Cost Optimization) + +Use a cheaper sub-model for simple tasks called via `llm_query()`: + +```yaml +backend: bedrock +backend_kwargs: + model_name: us.anthropic.claude-3-5-sonnet-20241022-v2:0 # Root model + region_name: us-east-1 + +# Enable multi-tier routing +other_backends: + - bedrock +other_backend_kwargs: + - model_name: us.anthropic.claude-3-5-haiku-20241022-v1:0 # Sub model + region_name: us-east-1 + +max_root_tokens: 100000 +max_sub_tokens: 50000 +``` + +**Cost savings**: ~80-90% on sub-queries (Haiku is 10x cheaper than Sonnet) + +## Supported Backends & Models + +### AWS Bedrock + +```yaml +backend: bedrock +backend_kwargs: + model_name: us.anthropic.claude-3-5-sonnet-20241022-v2:0 + region_name: us-east-1 +``` + +**Popular models**: +- `us.anthropic.claude-3-5-sonnet-20241022-v2:0` - Best quality +- `us.anthropic.claude-3-5-haiku-20241022-v1:0` - Fast & cheap +- `meta.llama3-70b-instruct-v1:0` - Open source + +**Credentials**: AWS CLI config or environment variables + +### OpenAI + +```yaml +backend: openai +backend_kwargs: + model_name: gpt-4o + # api_key: auto-detected from OPENAI_API_KEY +``` + +**Popular models**: +- `gpt-4o` - Best quality +- `gpt-4o-mini` - Fast & cheap +- `gpt-4-turbo` - Legacy flagship + +**Credentials**: `OPENAI_API_KEY` environment variable + +### Anthropic (Direct API) + +```yaml +backend: anthropic +backend_kwargs: + model_name: claude-sonnet-4-20250514 + # api_key: auto-detected from ANTHROPIC_API_KEY +``` + +**Popular models**: +- `claude-sonnet-4-20250514` - Latest Sonnet +- `claude-3-5-sonnet-20241022` - Previous generation +- `claude-3-5-haiku-20241022` - Fast & cheap + +**Credentials**: `ANTHROPIC_API_KEY` environment variable + +### Google Gemini + +```yaml +backend: gemini +backend_kwargs: + model_name: gemini-2.0-flash-exp + # api_key: auto-detected from GEMINI_API_KEY +``` + +**Popular models**: +- `gemini-2.0-flash-exp` - Very cheap (experimental) +- `gemini-1.5-pro` - Best quality +- `gemini-1.5-flash` - Fast + +**Credentials**: `GEMINI_API_KEY` environment variable + +## Execution Environment + +The CLI **hardcodes Docker** for safety and isolation. All refactoring code executes inside a `python:3.12-slim` Docker container with the working directory mounted. + +**Why Docker?** +- ✅ **Safety**: Prevents accidental damage to your host system +- ✅ **Isolation**: Clean environment for each run +- ✅ **Reproducibility**: Consistent Python version across systems + +**Requirements**: +- Docker installed and running +- Working directory is mounted read-write into container at `/workspace` + +## Troubleshooting + +### Error: "ModuleNotFoundError: No module named 'boto3'" (or other missing modules) + +**Cause**: Running `python -m rlm.refactor.cli` directly doesn't use the project's dependencies. + +**Solution**: Use `uv run` to execute within the virtual environment: +```bash +# Wrong - dependencies not available +python -m rlm.refactor.cli --working-dir ./my-task + +# Correct - uses project dependencies +uv run python -m rlm.refactor.cli --working-dir ./my-task +``` + +**Alternative**: Install missing backend-specific dependencies: +```bash +# For Bedrock (requires boto3) +uv pip install boto3 + +# Or install all optional dependencies +uv sync --all-extras +``` + +### Error: "Missing required file: rlm-instructions.md" + +**Solution**: Create `rlm-instructions.md` in your working directory: +```bash +cat > my-dir/rlm-instructions.md << 'EOF' +# My Task + +Task description here +EOF +``` + +### Error: "Missing required file: rlm-config.yaml" + +**Solution**: Create `rlm-config.yaml` in your working directory: +```bash +cat > my-dir/rlm-config.yaml << 'EOF' +backend: bedrock +backend_kwargs: + model_name: us.anthropic.claude-3-5-sonnet-20241022-v2:0 + region_name: us-east-1 +max_root_tokens: 50000 +EOF +``` + +### Error: "Failed to parse rlm-config.yaml" + +**Common issues**: +- Indentation must use spaces (not tabs) +- Missing quotes around special characters +- Incorrect YAML syntax + +**Solution**: Check YAML syntax or remove config file to use defaults + +### Error: "Field 'backend_kwargs.max_tokens' = 50 is too low" + +**Solution**: Increase `max_tokens` to at least 100: +```yaml +backend_kwargs: + max_tokens: 4096 # Must be >= 100 +``` + +### Error: "Field 'max_root_tokens' should not be in 'backend_kwargs'" + +**Common mistake**: Token budgets are top-level fields, not backend kwargs. + +**Wrong**: +```yaml +backend_kwargs: + model_name: gpt-4o + max_root_tokens: 100000 # ❌ Wrong location +``` + +**Correct**: +```yaml +backend_kwargs: + model_name: gpt-4o + +max_root_tokens: 100000 # ✅ Top-level field +``` + +### Warning: "AWS credentials not detected for Bedrock" + +**Solution**: Set AWS credentials: +```bash +# Option 1: Environment variables +export AWS_ACCESS_KEY_ID=your_key +export AWS_SECRET_ACCESS_KEY=your_secret + +# Option 2: AWS profile +export AWS_PROFILE=my-profile + +# Option 3: AWS CLI config +aws configure +``` + +### Error: Docker container failures or image issues + +**Symptoms**: +- Container won't start +- Git installation failures +- Execution errors inside Docker + +**Solution**: Check the Docker diagnostic logs in `{log_dir}/docker_*.jsonl`: +```bash +# Find errors in Docker logs +cat my-refactor-task/rlm-logs/docker_*.jsonl | jq 'select(.level == "ERROR" or .level == "WARNING")' + +# Check container startup +cat my-refactor-task/rlm-logs/docker_*.jsonl | jq 'select(.message | contains("Container"))' + +# View git installation details +cat my-refactor-task/rlm-logs/docker_*.jsonl | jq 'select(.message | contains("Git"))' +``` + +Common Docker issues: +- **Image pull fails**: Check your internet connection and Docker daemon status +- **Volume mount errors**: Verify the working directory path exists and is accessible +- **Git install fails**: The container needs internet access for apt-get; check network settings +- **Permission issues**: Ensure Docker has permissions to access the working directory + +### Question: "Where can I see what prompt was sent to the LLM?" + +Check the log files in `{working_dir}/rlm_logs/*.jsonl` (or custom `log_dir` from config): + +```bash +# View logs +cat my-refactor-task/rlm_logs/refactor_*.jsonl | jq . + +# See the exact messages +cat my-refactor-task/rlm_logs/refactor_*.jsonl | jq '.messages' +``` + +The CLI sends your `rlm-instructions.md` content directly to the LLM along with the refactoring system prompt. + +## Cost Optimization + +### Use Multi-Tier Routing + +**Single-tier** (all Sonnet): +```yaml +backend: bedrock +backend_kwargs: + model_name: us.anthropic.claude-3-5-sonnet-20241022-v2:0 +# Cost: ~$9 for complex task with many sub-queries +``` + +**Multi-tier** (Sonnet + Haiku): +```yaml +backend: bedrock +backend_kwargs: + model_name: us.anthropic.claude-3-5-sonnet-20241022-v2:0 +other_backends: + - bedrock +other_backend_kwargs: + - model_name: us.anthropic.claude-3-5-haiku-20241022-v1:0 +# Cost: ~$1.50 for same task (83% savings!) +``` + +### Cost Reference (2026) + +| Model | Input | Output | Use Case | +|-------|-------|--------|----------| +| **Claude 3.5 Sonnet** | $3/M | $15/M | Best quality, root model | +| **Claude 3.5 Haiku** | $0.25/M | $1.25/M | Fast, cheap sub-model | +| **GPT-4o** | $2.50/M | $10/M | High quality | +| **GPT-4o-mini** | $0.15/M | $0.60/M | Budget sub-model | +| **Gemini 2.0 Flash** | $0.10/M | $0.30/M | Very cheap | + +**Recommendation**: Use premium models (Sonnet, GPT-4o) for root, budget models (Haiku, mini) for subs. + +## Advanced Usage + +### Custom System Prompt + +The refactoring CLI uses the `REFACTOR_SYSTEM_PROMPT` from `rlm/refactor/prompts.py`. To customize, modify the CLI source: + +```python +# In rlm/refactor/cli.py (line ~245) +rlm = RLM( + backend=backend, + backend_kwargs=backend_kwargs, + environment="docker", + environment_kwargs=environment_kwargs, + custom_system_prompt=REFACTOR_SYSTEM_PROMPT, # ← Modify this + # ... +) +``` + +### Persistent Environments + +For multiple related tasks, use persistent environments: + +```yaml +# Not yet supported in config, use Python API: +# rlm = RLM(..., persistent=True) +``` + +### Logging & Analysis + +Two types of logs are written to `{log_dir}` (default: `rlm-logs/`): + +**1. Refactor trajectory logs** (`refactor_*.jsonl`): +```bash +# View execution traces +cat my-refactor-task/rlm-logs/refactor_*.jsonl | jq . + +# Analyze token usage +cat my-refactor-task/rlm-logs/refactor_*.jsonl | jq '.usage' + +# Examine specific iteration +cat my-refactor-task/rlm-logs/refactor_*.jsonl | jq 'select(.iteration == 3)' +``` + +**2. Docker diagnostic logs** (`docker_*.jsonl`): +Tracks container lifecycle events (setup, execution errors, cleanup): +```bash +# View Docker events +cat my-refactor-task/rlm-logs/docker_*.jsonl | jq . + +# Check for errors +cat my-refactor-task/rlm-logs/docker_*.jsonl | jq 'select(.level == "ERROR")' + +# Monitor git installation +cat my-refactor-task/rlm-logs/docker_*.jsonl | jq 'select(.message | contains("Git"))' +``` + +**Log correlation**: Both logs share the same timestamp/run_id suffix for easy matching. + +## Files + +``` +examples/refactor/ +├── README.md # This file +├── sample-working/ # Sample working directory (ready to run!) +│ ├── rlm-config.yaml # Example Bedrock configuration +│ ├── rlm-instructions.md # Task: Find which file contains "bob" +│ ├── sample-repo-one/ # Git repository #1 +│ │ ├── .git/ +│ │ ├── file_01.txt +│ │ └── file_02.txt +│ ├── sample-repo-two/ # Git repository #2 +│ │ ├── .git/ +│ │ ├── file_08.txt +│ │ ├── file_09.txt +│ │ └── file_10.txt +│ └── sample-repo-three/ # Git repository #3 +│ ├── .git/ +│ ├── file_03.txt +│ ├── file_04.txt +│ ├── file_05.txt # ← Contains "bob" +│ ├── file_06.txt +│ └── file_07.txt +└── rlm-config.yaml # Configuration template with all options +``` + +**CLI Location**: `rlm/refactor/cli.py` (run via `python -m rlm.refactor.cli`) + +## See Also + +- [RLM Documentation](../../README.md) +- [Multi-tier example](../multi_tier_example.py) +- [Bedrock example](../bedrock_example.py) +- [Refactoring CLI source](../../rlm/refactor/cli.py) diff --git a/examples/refactor/sample-working/rlm-config.yaml b/examples/refactor/sample-working/rlm-config.yaml new file mode 100644 index 000000000..df8abbf35 --- /dev/null +++ b/examples/refactor/sample-working/rlm-config.yaml @@ -0,0 +1,96 @@ +# ============================================================================= +# RLM REFACTORING CONFIGURATION +# ============================================================================= +# All fields are optional. Defaults used if missing. +# Environment is hardcoded to 'docker' for safety (see rlm/refactor/cli.py) + +# ============================================================================= +# ROOT MODEL +# ============================================================================= + +# Backend provider: openai, anthropic, bedrock, gemini, portkey +backend: bedrock + +backend_kwargs: + # Model identifier (REQUIRED) + # Bedrock examples: us.anthropic.claude-3-5-sonnet-20241022-v2:0, meta.llama3-70b-instruct-v1:0 + # OpenAI examples: gpt-4o, gpt-4o-mini, gpt-4-turbo + # Anthropic examples: claude-sonnet-4-20250514, claude-3-5-sonnet-20241022 + # Gemini examples: gemini-2.0-flash-exp, gemini-1.5-pro + model_name: us.anthropic.claude-3-5-sonnet-20241022-v2:0 + + # AWS region (Bedrock only) + region_name: us-east-1 + + # Optional: Max tokens per response (default: 4096 for Bedrock, 32768 for Anthropic) + # max_tokens: 8000 + + # Optional: Sampling temperature 0.0-2.0 (default: 1.0) + # temperature: 1.0 + + # Optional: Nucleus sampling 0.0-1.0 (default: varies by provider) + # top_p: 0.999 + +# ============================================================================= +# SUB-MODEL (OPTIONAL - Multi-tier cost optimization) +# ============================================================================= +# Configure a cheaper model for sub-tasks (file reading, simple analysis). +# Root model (above) handles strategic planning; sub-model handles execution. +# Cost savings: ~90% (Haiku vs Sonnet, or gpt-4o-mini vs gpt-4o) + +other_backends: + - bedrock + +other_backend_kwargs: + - model_name: us.anthropic.claude-3-5-haiku-20241022-v1:0 + region_name: us-east-1 + # Optional: max_tokens, temperature, top_p (same as above) + +# ============================================================================= +# TOKEN BUDGETS (Session-wide limits to prevent runaway costs) +# ============================================================================= + +# Max total tokens for root model across all iterations (default: 1,000,000) +# Cost of 1M tokens: $3-$15 (Sonnet input/output mix) +max_root_tokens: 50000 + +# Max total tokens for sub-model across all llm_query() calls (default: 1,000,000) +# Cost of 1M tokens: $0.25-$1.25 (Haiku input/output mix) +max_sub_tokens: 25000 + +# ============================================================================= +# EXECUTION SETTINGS +# ============================================================================= + +# Optional: Max REPL iterations before stopping (default: 30) +# max_iterations: 20 + +# Optional: Enable rich console output with colors/progress bars (default: false) +# verbose: true + +# ============================================================================= +# LOGGING +# ============================================================================= + +# Optional: Log directory relative to working directory (default: rlm_logs) +# Stores JSONL trajectory logs for debugging and analysis +# log_dir: rlm-logs + +# ============================================================================= +# COST REFERENCE (2026 pricing) +# ============================================================================= +# Bedrock Claude (on-demand): +# - claude-3-5-sonnet-v2: $3/M input, $15/M output +# - claude-3-5-haiku: $0.25/M input, $1.25/M output (~90% savings) +# +# OpenAI: +# - gpt-4o: $2.50/M input, $10/M output +# - gpt-4o-mini: $0.15/M input, $0.60/M output (~90% savings) +# +# Anthropic API: +# - claude-sonnet-4: $3/M input, $15/M output +# - claude-3-5-haiku: $0.25/M input, $1.25/M output (~90% savings) +# +# Gemini: +# - gemini-2.0-flash: $0.10/M input, $0.30/M output +# - gemini-1.5-pro: $1.25/M input, $5/M output diff --git a/examples/refactor/sample-working/rlm-instructions.md b/examples/refactor/sample-working/rlm-instructions.md new file mode 100644 index 000000000..0dc7c0367 --- /dev/null +++ b/examples/refactor/sample-working/rlm-instructions.md @@ -0,0 +1,22 @@ +# RLM Instructions: Find Bob + +## Objective +Search through all `.txt` files across all repositories to find which file contains the word "bob" (case-insensitive). + +## Task +1. Explore all git repositories in the working directory +2. Find all `.txt` files across all repositories +3. Search each file for the word "bob" +4. Return the repository and filename that contains "bob" + +## Expected Output +Provide the location of the file containing "bob" in the format: +``` +The file containing 'bob' is: / +``` + +## Notes +- Search should be case-insensitive +- Search across ALL repositories in the working directory +- Only one file should contain the word "bob" +- Return the relative path from the working directory (e.g., "sample-repo-three/file_05.txt") \ No newline at end of file diff --git a/examples/refactor/sample-working/sample-repo-one b/examples/refactor/sample-working/sample-repo-one new file mode 160000 index 000000000..8ea218c5d --- /dev/null +++ b/examples/refactor/sample-working/sample-repo-one @@ -0,0 +1 @@ +Subproject commit 8ea218c5dd4fe986fc853b85994c35ed94bc464f diff --git a/examples/refactor/sample-working/sample-repo-three b/examples/refactor/sample-working/sample-repo-three new file mode 160000 index 000000000..e8697e9f4 --- /dev/null +++ b/examples/refactor/sample-working/sample-repo-three @@ -0,0 +1 @@ +Subproject commit e8697e9f4ea95c69a3c57f399f5c179b5926f563 diff --git a/examples/refactor/sample-working/sample-repo-two b/examples/refactor/sample-working/sample-repo-two new file mode 160000 index 000000000..b66d46216 --- /dev/null +++ b/examples/refactor/sample-working/sample-repo-two @@ -0,0 +1 @@ +Subproject commit b66d462161c118af8006958e36a48898d4612769 diff --git a/examples/refactor_example.py b/examples/refactor_example.py new file mode 100644 index 000000000..90466af61 --- /dev/null +++ b/examples/refactor_example.py @@ -0,0 +1,88 @@ +""" +Programmatic usage of the RLM refactoring agent. + +For most users, the CLI is simpler: + # Basic usage + rlm-refactor --working-dir ./my-repos --model gpt-4o --verbose + + # With inference parameters + rlm-refactor --working-dir ./my-repos --model gpt-4o \\ + --max-tokens 8000 --temperature 0.7 --top-p 0.95 + + # With sub-model configuration + rlm-refactor --working-dir ./my-repos --model gpt-4o \\ + --sub-model gpt-4o-mini --sub-max-tokens 4000 --sub-temperature 0.3 + + # With session-wide token budgets + rlm-refactor --working-dir ./my-repos --model gpt-4o \\ + --max-root-tokens 500000 --max-sub-tokens 500000 + + # With provider-specific parameters + rlm-refactor --working-dir ./my-repos --model custom-model \\ + --backend-arg base_url=http://localhost:8080 \\ + --backend-arg api_key=custom-key + +This example shows how to use the refactoring system prompt with +the RLM class directly for custom integration. + +Setup: + 1. Ensure Docker is running + 2. Create a directory with git repos and an rlm-instructions.md file + 3. Run: python -m examples.refactor_example +""" + +import os + +from rlm.core.rlm import RLM +from rlm.logger import RLMLogger +from rlm.refactor.prompts import REFACTOR_SYSTEM_PROMPT + + +def main(): + working_dir = os.path.abspath("./example-refactor-workspace") + + # Read the task instructions + instructions_path = os.path.join(working_dir, "rlm-instructions.md") + if not os.path.isfile(instructions_path): + print(f"Create {instructions_path} with your refactoring instructions first.") + print("See examples/rlm-instructions-template.md for a template.") + return + + with open(instructions_path) as f: + instructions = f.read() + + # Setup logging inside the working directory + log_dir = os.path.join(working_dir, "rlm-logs") + os.makedirs(log_dir, exist_ok=True) + logger = RLMLogger(log_dir=log_dir, file_name="refactor") + + rlm = RLM( + backend="openai", + backend_kwargs={"model_name": "gpt-4o"}, + environment="docker", + environment_kwargs={ + "working_dir": working_dir, + # Use a custom image with git pre-installed for faster startup: + # "image": "my-rlm-refactor-image", + }, + custom_system_prompt=REFACTOR_SYSTEM_PROMPT, + logger=logger, + max_iterations=50, + verbose=True, + # Optional: use a cheaper model for sub-LLM queries + # other_backends=["openai"], + # other_backend_kwargs=[{"model_name": "gpt-4o-mini"}], + ) + + result = rlm.completion( + prompt=instructions, + root_prompt=instructions[:500], + ) + + print(f"\nResult: {result.response}") + print(f"Time: {result.execution_time:.1f}s") + print(f"Logs: {log_dir}") + + +if __name__ == "__main__": + main() diff --git a/examples/rlm-instructions-template.md b/examples/rlm-instructions-template.md new file mode 100644 index 000000000..2b816da27 --- /dev/null +++ b/examples/rlm-instructions-template.md @@ -0,0 +1,40 @@ +# RLM Refactoring Instructions + +## Goal + + + +Rename all instances of `OldClassName` to `NewClassName` across the codebase, +updating imports, type annotations, docstrings, and tests. + +## Repositories + + + +- `my-service/` — Main application service (Python, FastAPI) +- `my-client-lib/` — Client library that depends on my-service's types + +## Constraints + + + +- Do not modify files in `vendor/` or `node_modules/` +- Preserve all existing test coverage — tests must still pass +- Do not change any public API signatures beyond the rename +- Keep backwards-compatible type aliases where the old name was part of a public API + +## Commit Strategy + + + +- One commit per repository +- Commit message format: `refactor: rename OldClassName to NewClassName` +- Do not push — changes will be reviewed manually + +## Verification + + + +- Run `python -m pytest` in each repo after changes +- Run `grep -rn OldClassName /repos` to confirm no remaining references +- Check `git diff --stat` to review scope of changes diff --git a/pyproject.toml b/pyproject.toml index 99aff8a4c..e5949de3c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,10 +25,14 @@ dependencies = [ "portkey-ai>=2.1.0", "pytest>=9.0.2", "python-dotenv>=1.2.1", + "pyyaml>=6.0", "requests>=2.32.5", "rich>=13.0.0", ] +[project.scripts] +rlm-refactor = "rlm.refactor.cli:main" + [project.urls] Homepage = "https://github.com/alexzhang13/rlm" Repository = "https://github.com/alexzhang13/rlm" diff --git a/rlm/environments/docker_repl.py b/rlm/environments/docker_repl.py index 12472caa2..c7f6d5d55 100644 --- a/rlm/environments/docker_repl.py +++ b/rlm/environments/docker_repl.py @@ -93,13 +93,17 @@ def _handle_batched(self, body: dict) -> dict: return {"responses": results} -def _build_exec_script(code: str, proxy_port: int, depth: int = 1) -> str: +def _build_exec_script( + code: str, proxy_port: int, depth: int = 1, working_dir: str | None = None +) -> str: """Build execution script for the container.""" code_b64 = base64.b64encode(code.encode()).decode() + extra_imports = ", subprocess, glob, shutil, pathlib" if working_dir else "" + chdir_line = "os.chdir('/repos')\n" if working_dir else "" return textwrap.dedent( f''' -import sys, io, json, base64, traceback, os, requests +import sys, io, json, base64, traceback, os, requests{extra_imports} try: import dill except ImportError: @@ -162,7 +166,7 @@ def SHOW_VARS(): _globals = {{"__builtins__": __builtins__, "__name__": "__main__", "llm_query": llm_query, "llm_query_batched": llm_query_batched, "FINAL_VAR": FINAL_VAR, "SHOW_VARS": SHOW_VARS}} -code = base64.b64decode("{code_b64}").decode() +{chdir_line}code = base64.b64decode("{code_b64}").decode() stdout_buf, stderr_buf = io.StringIO(), io.StringIO() old_stdout, old_stderr = sys.stdout, sys.stderr @@ -199,6 +203,8 @@ def __init__( setup_code: str | None = None, persistent: bool = False, depth: int = 1, + working_dir: str | None = None, + docker_log_path: str | None = None, **kwargs, ): if persistent: @@ -209,6 +215,8 @@ def __init__( self.image = image self.lm_handler_address = lm_handler_address + self.working_dir = os.path.abspath(working_dir) if working_dir else None + self.docker_log_path = docker_log_path self.container_id: str | None = None self.proxy_server: HTTPServer | None = None self.proxy_thread: threading.Thread | None = None @@ -228,8 +236,35 @@ def __init__( if setup_code: self.execute_code(setup_code) + def _log_docker_event(self, level: str, message: str, details: dict | None = None): + """Log Docker-specific events to persistent file.""" + if not self.docker_log_path: + return + + import json + from datetime import datetime + + entry = { + "type": "docker", + "timestamp": datetime.now().isoformat(), + "level": level, + "message": message, + "container_id": self.container_id, + "details": details or {}, + } + + with open(self.docker_log_path, "a") as f: + json.dump(entry, f) + f.write("\n") + def setup(self): """Start the proxy server and Docker container.""" + self._log_docker_event( + "INFO", + "Starting Docker setup", + {"image": self.image, "working_dir": self.working_dir}, + ) + # Start LLM proxy server handler = type( "Handler", @@ -247,34 +282,93 @@ def setup(self): self.proxy_thread.start() # Start Docker container - result = subprocess.run( - [ - "docker", - "run", - "-d", - "--rm", - "-v", - f"{self.temp_dir}:/workspace", - "--add-host", - "host.docker.internal:host-gateway", - self.image, - "tail", - "-f", - "/dev/null", - ], - capture_output=True, - text=True, - ) + docker_cmd = [ + "docker", + "run", + "-d", + "--rm", + "-v", + f"{self.temp_dir}:/workspace", + ] + if self.working_dir: + docker_cmd += ["-v", f"{self.working_dir}:/repos"] + docker_cmd += [ + "--add-host", + "host.docker.internal:host-gateway", + self.image, + "tail", + "-f", + "/dev/null", + ] + result = subprocess.run(docker_cmd, capture_output=True, text=True) if result.returncode != 0: + self._log_docker_event( + "ERROR", + "Failed to start container", + { + "command": " ".join(docker_cmd), + "returncode": result.returncode, + "stderr": result.stderr, + "stdout": result.stdout, + }, + ) raise RuntimeError(f"Failed to start container: {result.stderr}") self.container_id = result.stdout.strip() + self._log_docker_event("INFO", "Container started", {"container_id": self.container_id}) # Install dependencies - subprocess.run( + pip_result = subprocess.run( ["docker", "exec", self.container_id, "pip", "install", "-q", "dill", "requests"], capture_output=True, + text=True, ) + if pip_result.returncode != 0: + self._log_docker_event( + "WARNING", + "Pip install had issues", + {"stderr": pip_result.stderr, "stdout": pip_result.stdout}, + ) + else: + self._log_docker_event("INFO", "Dependencies installed successfully") + + # Install git and configure safe directories for bind-mounted repos + if self.working_dir: + git_result = subprocess.run( + [ + "docker", + "exec", + self.container_id, + "bash", + "-c", + "apt-get update -qq && apt-get install -y -qq git 2>&1", + ], + capture_output=True, + text=True, + ) + if git_result.returncode != 0: + self._log_docker_event( + "ERROR", + "Git installation failed", + {"stderr": git_result.stderr, "stdout": git_result.stdout}, + ) + else: + self._log_docker_event("INFO", "Git installed successfully") + + subprocess.run( + [ + "docker", + "exec", + self.container_id, + "git", + "config", + "--global", + "--add", + "safe.directory", + "*", + ], + capture_output=True, + ) def load_context(self, context_payload: dict | list | str): """Load context by writing to a file in the mounted workspace.""" @@ -302,13 +396,23 @@ def execute_code(self, code: str) -> REPLResult: with self._calls_lock: self.pending_calls.clear() - script = _build_exec_script(code, self.proxy_port, self.depth) + script = _build_exec_script(code, self.proxy_port, self.depth, self.working_dir) result = subprocess.run( ["docker", "exec", self.container_id, "python", "-c", script], capture_output=True, text=True, ) + if result.returncode != 0: + self._log_docker_event( + "WARNING", + "Code execution failed", + { + "returncode": result.returncode, + "stderr": result.stderr[:500] if result.stderr else "", + }, + ) + with self._calls_lock: calls = self.pending_calls.copy() self.pending_calls.clear() @@ -333,9 +437,24 @@ def execute_code(self, code: str) -> REPLResult: ) def cleanup(self): + self._log_docker_event("INFO", "Starting cleanup") + if hasattr(self, "container_id") and self.container_id: - subprocess.run(["docker", "stop", self.container_id], capture_output=True) + stop_result = subprocess.run( + ["docker", "stop", self.container_id], + capture_output=True, + text=True, + ) + if stop_result.returncode != 0: + self._log_docker_event( + "WARNING", + "Failed to stop container", + {"stderr": stop_result.stderr}, + ) + else: + self._log_docker_event("INFO", "Container stopped") self.container_id = None + if hasattr(self, "proxy_server") and self.proxy_server: self.proxy_server.shutdown() self.proxy_server = None @@ -344,6 +463,8 @@ def cleanup(self): shutil.rmtree(self.temp_dir, ignore_errors=True) + self._log_docker_event("INFO", "Cleanup complete") + def __enter__(self): return self diff --git a/rlm/refactor/__init__.py b/rlm/refactor/__init__.py new file mode 100644 index 000000000..ee3eb7a3b --- /dev/null +++ b/rlm/refactor/__init__.py @@ -0,0 +1,3 @@ +from rlm.refactor.prompts import REFACTOR_SYSTEM_PROMPT + +__all__ = ["REFACTOR_SYSTEM_PROMPT"] diff --git a/rlm/refactor/cli.py b/rlm/refactor/cli.py new file mode 100644 index 000000000..4bbdab3b6 --- /dev/null +++ b/rlm/refactor/cli.py @@ -0,0 +1,298 @@ +""" +CLI entry point for rlm-refactor. + +Reads configuration from rlm-config.yaml and task instructions from +rlm-instructions.md in the working directory. Executes refactoring in +an isolated Docker container for safety. + +Configuration: + Working Directory (required): + - Contains rlm-config.yaml (model configuration) + - Contains rlm-instructions.md (refactoring task description) + - Mounted into Docker container at /workspace + + rlm-config.yaml: + - backend: LM provider (openai, anthropic, bedrock, gemini, portkey) + - backend_kwargs: Model config with model_name (required) + - other_backends: (optional) Sub-model for cost optimization + - max_root_tokens: Session-wide token budget for root model + - max_sub_tokens: Session-wide token budget for sub-model + - max_iterations: Max REPL iterations (default: 30) + - verbose: Enable rich console output (default: false) + - log_dir: Log directory name (default: rlm_logs) + +Usage: + # Interactive mode (prompts for working directory) + rlm-refactor + + # Direct mode (specify working directory) + rlm-refactor --working-dir ./my-project + + # Working directory must contain: + # - rlm-config.yaml (configuration) + # - rlm-instructions.md (task description) + +Examples: + # Setup + mkdir my-refactoring && cd my-refactoring + + # Create task instructions + cat > rlm-instructions.md << 'EOF' + # Find all occurrences of "TODO" + Search all Python files and list every TODO comment. + EOF + + # Create config + cat > rlm-config.yaml << 'EOF' + backend: bedrock + backend_kwargs: + model_name: us.anthropic.claude-3-5-sonnet-20241022-v2:0 + region_name: us-east-1 + max_root_tokens: 50000 + verbose: true + EOF + + # Run refactoring + rlm-refactor --working-dir . +""" + +import argparse +import sys +import uuid +from datetime import datetime +from pathlib import Path + +import yaml +from rich.console import Console +from rich.prompt import Prompt + +from rlm.core.rlm import RLM +from rlm.logger import RLMLogger +from rlm.refactor.prompts import REFACTOR_SYSTEM_PROMPT + +INSTRUCTIONS_FILENAME = "rlm-instructions.md" +CONFIG_FILENAME = "rlm-config.yaml" + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + prog="rlm-refactor", + description="Run an RLM-powered code refactoring agent in Docker. " + "All configuration is read from rlm-config.yaml in the working directory.", + ) + parser.add_argument( + "--working-dir", + default=None, + help="Working directory containing rlm-config.yaml and rlm-instructions.md " + "(default: interactive prompt)", + ) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> None: + console = Console() + args = parse_args(argv) + + # Step 1: Get working directory (prompt if not provided) + if args.working_dir is None: + console.print("\n[bold cyan]RLM Refactoring Agent[/bold cyan]") + console.print( + "All configuration read from [yellow]rlm-config.yaml[/yellow] and [yellow]rlm-instructions.md[/yellow]\n" + ) + + while True: + working_dir_input = Prompt.ask( + "[bold cyan]Enter working directory path[/bold cyan]", default=str(Path.cwd()) + ) + working_dir = Path(working_dir_input).expanduser().resolve() + + if not working_dir.exists(): + console.print(f"[red]✗ Directory does not exist: {working_dir}[/red]\n") + continue + if not working_dir.is_dir(): + console.print(f"[red]✗ Path is not a directory: {working_dir}[/red]\n") + continue + break + else: + working_dir = Path(args.working_dir).expanduser().resolve() + if not working_dir.exists(): + console.print(f"[red]Error: Working directory does not exist: {working_dir}[/red]") + sys.exit(1) + if not working_dir.is_dir(): + console.print(f"[red]Error: Path is not a directory: {working_dir}[/red]") + sys.exit(1) + + # Step 2: Load rlm-instructions.md + instructions_path = working_dir / INSTRUCTIONS_FILENAME + if not instructions_path.exists(): + console.print( + f"[red]Error: {INSTRUCTIONS_FILENAME} not found in {working_dir}[/red]\n" + f"Create {INSTRUCTIONS_FILENAME} with your refactoring task description." + ) + sys.exit(1) + + instructions = instructions_path.read_text() + if not instructions.strip(): + console.print(f"[red]Error: {INSTRUCTIONS_FILENAME} is empty[/red]") + sys.exit(1) + + # Step 3: Load rlm-config.yaml + config_path = working_dir / CONFIG_FILENAME + if not config_path.exists(): + console.print( + f"[red]Error: {CONFIG_FILENAME} not found in {working_dir}[/red]\n" + f"Create {CONFIG_FILENAME} with your model and backend configuration." + ) + sys.exit(1) + + try: + with open(config_path) as f: + config = yaml.safe_load(f) or {} + except yaml.YAMLError as e: + console.print(f"[red]Error: Invalid YAML in {CONFIG_FILENAME}:[/red]\n{e}") + sys.exit(1) + + # Step 4: Extract and validate configuration + backend = config.get("backend") + if not backend: + console.print( + "[red]Error: 'backend' is required in config (e.g., openai, anthropic, bedrock)[/red]" + ) + sys.exit(1) + + backend_kwargs = config.get("backend_kwargs", {}) + if not isinstance(backend_kwargs, dict): + console.print("[red]Error: 'backend_kwargs' must be a dictionary[/red]") + sys.exit(1) + if "model_name" not in backend_kwargs: + console.print("[red]Error: 'backend_kwargs.model_name' is required[/red]") + sys.exit(1) + + # Validate forbidden parameters + forbidden = ["max_root_tokens", "max_sub_tokens"] + for param in forbidden: + if param in backend_kwargs: + console.print( + f"[red]Error: '{param}' should be a top-level config parameter, " + f"not inside backend_kwargs[/red]" + ) + sys.exit(1) + + # Handle multi-tier configuration + other_backends = config.get("other_backends") + other_backend_kwargs = config.get("other_backend_kwargs") + + if other_backends is not None or other_backend_kwargs is not None: + if other_backends is None or other_backend_kwargs is None: + console.print( + "[red]Error: Both 'other_backends' and 'other_backend_kwargs' must be specified together[/red]" + ) + sys.exit(1) + if not isinstance(other_backends, list) or not isinstance(other_backend_kwargs, list): + console.print( + "[red]Error: 'other_backends' and 'other_backend_kwargs' must be lists[/red]" + ) + sys.exit(1) + if len(other_backends) != len(other_backend_kwargs): + console.print( + "[red]Error: 'other_backends' and 'other_backend_kwargs' must have same length[/red]" + ) + sys.exit(1) + + # Validate sub-backend kwargs + for i, sub_kwargs in enumerate(other_backend_kwargs): + if not isinstance(sub_kwargs, dict): + console.print(f"[red]Error: other_backend_kwargs[{i}] must be a dictionary[/red]") + sys.exit(1) + for param in forbidden: + if param in sub_kwargs: + console.print( + f"[red]Error: '{param}' should be a top-level config parameter, " + f"not inside other_backend_kwargs[{i}][/red]" + ) + sys.exit(1) + + # Extract token budgets + max_root_tokens = config.get("max_root_tokens", 1_000_000) + max_sub_tokens = config.get("max_sub_tokens", 1_000_000) + + if not isinstance(max_root_tokens, int) or max_root_tokens <= 0: + console.print("[red]Error: 'max_root_tokens' must be a positive integer[/red]") + sys.exit(1) + if not isinstance(max_sub_tokens, int) or max_sub_tokens <= 0: + console.print("[red]Error: 'max_sub_tokens' must be a positive integer[/red]") + sys.exit(1) + + # Extract execution settings + max_iterations = config.get("max_iterations", 30) + verbose = config.get("verbose", False) + log_dir_name = config.get("log_dir", "rlm_logs") + + # Step 5: Setup logging + log_dir = working_dir / log_dir_name + log_dir.mkdir(exist_ok=True) + logger = RLMLogger(log_dir=str(log_dir), file_name="refactor") + + # Create Docker diagnostic log with matching timestamp/run_id + timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + run_id = str(uuid.uuid4())[:8] + docker_log_path = str(log_dir / f"docker_{timestamp}_{run_id}.jsonl") + + # Step 6: Build environment kwargs + environment_kwargs = { + "working_dir": str(working_dir), + "image": "python:3.12-slim", + "docker_log_path": docker_log_path, + } + + # Step 7: Create RLM instance + console.print(f"\n[green]✓[/green] Working directory: {working_dir}") + console.print(f"[green]✓[/green] Backend: {backend}") + console.print(f"[green]✓[/green] Root model: {backend_kwargs.get('model_name')}") + if other_backends: + console.print(f"[green]✓[/green] Sub model: {other_backend_kwargs[0].get('model_name')}") + console.print(f"[green]✓[/green] Max iterations: {max_iterations}") + console.print( + f"[green]✓[/green] Token budgets: {max_root_tokens:,} root / {max_sub_tokens:,} sub\n" + ) + + rlm = RLM( + backend=backend, + backend_kwargs=backend_kwargs, + environment="docker", + environment_kwargs=environment_kwargs, + custom_system_prompt=REFACTOR_SYSTEM_PROMPT, + logger=logger, + max_iterations=max_iterations, + max_root_tokens=max_root_tokens, + max_sub_tokens=max_sub_tokens, + verbose=verbose, + other_backends=other_backends, + other_backend_kwargs=other_backend_kwargs, + ) + + # Step 8: Run refactoring + root_prompt = instructions[:500] if len(instructions) > 500 else instructions + + # Show spinner only when verbose mode is disabled (verbose output provides its own progress) + if verbose: + console.print("[bold cyan]→[/bold cyan] Starting RLM agent...\n") + result = rlm.completion(prompt=instructions, root_prompt=root_prompt) + else: + with console.status( + "[bold cyan]RLM agent iterating through your instructions...[/bold cyan]", + spinner="dots", + ): + result = rlm.completion(prompt=instructions, root_prompt=root_prompt) + + # Step 9: Display results + console.print("\n" + "=" * 60) + console.print("[bold green]REFACTORING COMPLETE[/bold green]") + console.print("=" * 60) + console.print(f"\n{result.response}") + console.print(f"\n[dim]Execution time: {result.execution_time:.1f}s[/dim]") + console.print(f"[dim]Logs written to: {log_dir}[/dim]") + + +if __name__ == "__main__": + main() diff --git a/rlm/refactor/prompts.py b/rlm/refactor/prompts.py new file mode 100644 index 000000000..38fe15d9e --- /dev/null +++ b/rlm/refactor/prompts.py @@ -0,0 +1,117 @@ +import textwrap + +REFACTOR_SYSTEM_PROMPT = textwrap.dedent( + """\ +You are a code refactoring agent operating inside a Docker container with full read-write \ +access to one or more git repositories mounted at `/repos`. Your working directory is `/repos`. + +## Your Task + +Read `/repos/rlm-instructions.md` for your task description. This file contains the \ +refactoring goals, constraints, and any specific instructions. Always read it first. + +## Environment + +You have a Python REPL with the full standard library available. Key modules: +- `os`, `pathlib`, `glob`, `shutil` — filesystem operations +- `subprocess` — run shell commands including `git` +- `json`, `re`, `textwrap` — text processing + +Execute code by wrapping it in triple backticks with the `repl` language identifier: +```repl +import os +for entry in os.listdir('/repos'): + print(entry) +``` + +## Git Operations + +Use `subprocess.run()` for git commands. Always capture output for inspection: +```repl +import subprocess +result = subprocess.run(['git', 'status'], capture_output=True, text=True, cwd='/repos/my-repo') +print(result.stdout) +``` + +Make incremental commits with clear, descriptive messages: +```repl +subprocess.run(['git', 'add', '-A'], cwd='/repos/my-repo') +subprocess.run(['git', 'commit', '-m', 'refactor: extract helper function from process_data'], cwd='/repos/my-repo') +``` + +Do NOT push. The user will review and push after your work is complete. + +## File Operations + +Read files: +```repl +content = open('/repos/my-repo/src/main.py').read() +print(content[:2000]) # Preview first 2000 chars +``` + +Write files: +```repl +with open('/repos/my-repo/src/main.py', 'w') as f: + f.write(new_content) +``` + +Search across repos: +```repl +result = subprocess.run(['grep', '-rn', 'old_function_name', '/repos'], capture_output=True, text=True) +print(result.stdout) +``` + +## Sub-LLM Queries + +You have access to sub-LLMs for analyzing code that is too large to reason about directly: +- `llm_query(prompt, model=None)` — query a sub-LLM (handles ~500K characters) +- `llm_query_batched(prompts, model=None)` — concurrent queries, returns results in order + +Use these to analyze large files, generate code, or plan complex changes: +```repl +content = open('/repos/my-repo/src/large_module.py').read() +analysis = llm_query(f"Analyze this Python module and identify all functions that should be extracted into a utils module:\\n\\n{content}") +print(analysis) +``` + +For processing many files concurrently: +```repl +import glob +files = glob.glob('/repos/my-repo/src/**/*.py', recursive=True) +prompts = [f"Summarize the purpose of this file:\\n\\n{open(f).read()}" for f in files[:10]] +summaries = llm_query_batched(prompts) +for f, s in zip(files[:10], summaries): + print(f"--- {f} ---") + print(s) +``` + +## State Management + +Variables persist across REPL blocks within the same session: +- `SHOW_VARS()` — list all variables you have created +- `FINAL_VAR(variable_name)` — return a variable as your final answer + +## Recommended Workflow + +1. **Read instructions**: Read `/repos/rlm-instructions.md` to understand the task +2. **Explore structure**: List repos, examine directory trees, read key files +3. **Plan changes**: Use sub-LLMs to analyze code and develop a plan +4. **Implement incrementally**: Make changes in small, testable steps +5. **Commit after each logical change**: Use clear commit messages +6. **Verify**: Run tests, linting, or other validation if applicable +7. **Summarize**: Provide a final answer listing all changes and commits + +## Completing Your Task + +When you are done, provide a final answer using FINAL(your summary here) or FINAL_VAR(variable_name). \ +Your final answer should summarize: +- What changes were made +- Which files were modified +- What commits were created +- Any issues encountered or items needing manual review + +WARNING - COMMON MISTAKE: FINAL_VAR retrieves an EXISTING variable. You MUST create and assign \ +the variable in a ```repl``` block FIRST, then call FINAL_VAR in a SEPARATE step. + +Think step by step, plan carefully, and execute immediately — do not just describe what you will do.""" +) diff --git a/tests/mock_lm.py b/tests/mock_lm.py index 461f73c68..f6070f999 100644 --- a/tests/mock_lm.py +++ b/tests/mock_lm.py @@ -26,9 +26,7 @@ def get_usage_summary(self): ) def get_last_usage(self): - return ModelUsageSummary( - total_calls=1, total_input_tokens=10, total_output_tokens=10 - ) + return ModelUsageSummary(total_calls=1, total_input_tokens=10, total_output_tokens=10) class MockLMWithTokenTracking(BaseLM): @@ -112,9 +110,7 @@ def completion(self, prompt: str | dict[str, Any] | list[dict[str, Any]]) -> str if isinstance(prompt, str): input_tokens = max(len(prompt) // 4, 1) elif isinstance(prompt, list): - input_tokens = max( - sum(len(str(m.get("content", ""))) for m in prompt) // 4, 1 - ) + input_tokens = max(sum(len(str(m.get("content", ""))) for m in prompt) // 4, 1) else: input_tokens = 10 diff --git a/tests/test_refactor.py b/tests/test_refactor.py new file mode 100644 index 000000000..7ad4e503f --- /dev/null +++ b/tests/test_refactor.py @@ -0,0 +1,204 @@ +"""Tests for the refactoring module: DockerREPL working_dir support and CLI.""" + +import os +import tempfile +from unittest.mock import MagicMock, patch + +from rlm.environments.docker_repl import _build_exec_script +from rlm.refactor.prompts import REFACTOR_SYSTEM_PROMPT + + +class TestBuildExecScriptWorkingDir: + """Tests for _build_exec_script with working_dir parameter.""" + + def test_no_working_dir_unchanged(self): + """When working_dir is None, script should not contain /repos chdir or extra imports.""" + script = _build_exec_script("x = 1", proxy_port=8080, depth=1, working_dir=None) + assert "os.chdir('/repos')" not in script + assert "import subprocess" not in script + assert "import glob" not in script + + def test_working_dir_adds_chdir(self): + """When working_dir is set, script should chdir to /repos.""" + script = _build_exec_script("x = 1", proxy_port=8080, depth=1, working_dir="/tmp/repos") + assert "os.chdir('/repos')" in script + + def test_working_dir_adds_extra_imports(self): + """When working_dir is set, script should import subprocess, glob, shutil, pathlib.""" + script = _build_exec_script("x = 1", proxy_port=8080, depth=1, working_dir="/tmp/repos") + assert "subprocess" in script + assert "glob" in script + assert "shutil" in script + assert "pathlib" in script + + def test_code_still_embedded(self): + """User code is still base64-embedded regardless of working_dir.""" + script = _build_exec_script( + "print('hello')", proxy_port=8080, depth=1, working_dir="/tmp/repos" + ) + assert "base64.b64decode" in script + + def test_proxy_port_in_script(self): + """Proxy port appears in the script URL.""" + script = _build_exec_script("x = 1", proxy_port=9999, depth=1, working_dir="/tmp/repos") + assert "9999" in script + + +class TestDockerREPLWorkingDir: + """Tests for DockerREPL constructor and setup with working_dir.""" + + @patch("rlm.environments.docker_repl.subprocess.run") + @patch("rlm.environments.docker_repl.HTTPServer") + def test_working_dir_adds_volume_mount(self, mock_httpserver, mock_run): + """When working_dir is set, docker run should include -v for /repos.""" + from rlm.environments.docker_repl import DockerREPL + + # Mock HTTPServer + mock_server = MagicMock() + mock_server.server_address = ("127.0.0.1", 12345) + mock_httpserver.return_value = mock_server + + # Mock all subprocess.run calls to succeed + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = "container123\n" + mock_result.stderr = "" + mock_run.return_value = mock_result + + with tempfile.TemporaryDirectory() as tmpdir: + working_dir = os.path.join(tmpdir, "repos") + os.makedirs(working_dir) + + repl = DockerREPL(working_dir=working_dir) + + # Find the docker run call + docker_run_calls = [ + call + for call in mock_run.call_args_list + if call[0] and call[0][0][:2] == ["docker", "run"] + ] + assert len(docker_run_calls) == 1 + + docker_cmd = docker_run_calls[0][0][0] + # Check that -v working_dir:/repos is in the command + assert "-v" in docker_cmd + repos_mount_idx = None + for i, arg in enumerate(docker_cmd): + if arg == "-v" and i + 1 < len(docker_cmd) and ":/repos" in docker_cmd[i + 1]: + repos_mount_idx = i + break + assert repos_mount_idx is not None, f"No /repos mount found in: {docker_cmd}" + + repl.cleanup() + + @patch("rlm.environments.docker_repl.subprocess.run") + @patch("rlm.environments.docker_repl.HTTPServer") + def test_working_dir_installs_git(self, mock_httpserver, mock_run): + """When working_dir is set, setup should install git and configure safe.directory.""" + from rlm.environments.docker_repl import DockerREPL + + mock_server = MagicMock() + mock_server.server_address = ("127.0.0.1", 12345) + mock_httpserver.return_value = mock_server + + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = "container123\n" + mock_result.stderr = "" + mock_run.return_value = mock_result + + with tempfile.TemporaryDirectory() as tmpdir: + working_dir = os.path.join(tmpdir, "repos") + os.makedirs(working_dir) + + repl = DockerREPL(working_dir=working_dir) + + # Check that git install was called + all_calls_str = str(mock_run.call_args_list) + assert "apt-get" in all_calls_str and "git" in all_calls_str + + # Check that git config safe.directory was called + assert "safe.directory" in all_calls_str + + repl.cleanup() + + @patch("rlm.environments.docker_repl.subprocess.run") + @patch("rlm.environments.docker_repl.HTTPServer") + def test_no_working_dir_no_git(self, mock_httpserver, mock_run): + """When working_dir is None, setup should NOT install git.""" + from rlm.environments.docker_repl import DockerREPL + + mock_server = MagicMock() + mock_server.server_address = ("127.0.0.1", 12345) + mock_httpserver.return_value = mock_server + + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = "container123\n" + mock_result.stderr = "" + mock_run.return_value = mock_result + + repl = DockerREPL() + + # Should NOT have git install calls + all_calls_str = str(mock_run.call_args_list) + assert "apt-get" not in all_calls_str + assert "safe.directory" not in all_calls_str + + repl.cleanup() + + @patch("rlm.environments.docker_repl.subprocess.run") + @patch("rlm.environments.docker_repl.HTTPServer") + def test_no_working_dir_no_repos_mount(self, mock_httpserver, mock_run): + """When working_dir is None, docker run should NOT include /repos mount.""" + from rlm.environments.docker_repl import DockerREPL + + mock_server = MagicMock() + mock_server.server_address = ("127.0.0.1", 12345) + mock_httpserver.return_value = mock_server + + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = "container123\n" + mock_result.stderr = "" + mock_run.return_value = mock_result + + repl = DockerREPL() + + docker_run_calls = [ + call + for call in mock_run.call_args_list + if call[0] and call[0][0][:2] == ["docker", "run"] + ] + docker_cmd = docker_run_calls[0][0][0] + + for i, arg in enumerate(docker_cmd): + if arg == "-v" and i + 1 < len(docker_cmd): + assert ":/repos" not in docker_cmd[i + 1], "Should not have /repos mount" + + repl.cleanup() + + +class TestRefactorSystemPrompt: + """Tests for the refactoring system prompt.""" + + def test_prompt_mentions_repos(self): + assert "/repos" in REFACTOR_SYSTEM_PROMPT + + def test_prompt_mentions_instructions(self): + assert "rlm-instructions.md" in REFACTOR_SYSTEM_PROMPT + + def test_prompt_mentions_git(self): + assert "git" in REFACTOR_SYSTEM_PROMPT + + def test_prompt_mentions_llm_query(self): + assert "llm_query" in REFACTOR_SYSTEM_PROMPT + + def test_prompt_mentions_final_var(self): + assert "FINAL_VAR" in REFACTOR_SYSTEM_PROMPT + + def test_prompt_mentions_subprocess(self): + assert "subprocess" in REFACTOR_SYSTEM_PROMPT + + def test_prompt_mentions_repl(self): + assert "```repl" in REFACTOR_SYSTEM_PROMPT diff --git a/uv.lock b/uv.lock index 384da7bb9..9131fc36f 100644 --- a/uv.lock +++ b/uv.lock @@ -1716,6 +1716,7 @@ dependencies = [ { name = "portkey-ai" }, { name = "pytest" }, { name = "python-dotenv" }, + { name = "pyyaml" }, { name = "requests" }, { name = "rich" }, ] @@ -1760,6 +1761,7 @@ requires-dist = [ { name = "prime-sandboxes", marker = "extra == 'prime'", specifier = ">=0.2.0" }, { name = "pytest", specifier = ">=9.0.2" }, { name = "python-dotenv", specifier = ">=1.2.1" }, + { name = "pyyaml", specifier = ">=6.0" }, { name = "requests", specifier = ">=2.32.5" }, { name = "rich", specifier = ">=13.0.0" }, ]