From 74757bd2a6450d6f9e38b9016b3b546b41803df3 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:10:59 -0700 Subject: [PATCH 1/4] refactor(trajectory_compressor): dedupe sync/async compress twins into plan/assemble helpers, split main into mode helpers, table-drive from_yaml and _detect_provider --- trajectory_compressor.py | 1617 +++++++++++++------------------------- 1 file changed, 566 insertions(+), 1051 deletions(-) diff --git a/trajectory_compressor.py b/trajectory_compressor.py index 58c0ca3245..163204b046 100644 --- a/trajectory_compressor.py +++ b/trajectory_compressor.py @@ -1,43 +1,27 @@ #!/usr/bin/env python3 -""" -Trajectory Compressor +"""Trajectory Compressor — post-process agent trajectories into a token budget. -Post-processes completed agent trajectories to compress them within a target -token budget while preserving training signal quality. - -Compression Strategy: -1. Protect first turns (system, human, first gpt, first tool) -2. Protect last N turns (final actions and conclusions) -3. Compress MIDDLE turns only, starting from 2nd tool response -4. Compress only as much as needed to fit under target -5. Replace compressed region with a single human summary message -6. Keep remaining tool calls intact (model continues working after summary) +Strategy: protect the head (system, human, first gpt, first tool) and the last N +turns; from the middle, summarize only as many turns as needed (never splitting a +/ pair) and replace them with one human summary turn. Usage: - # Compress a directory of JSONL files - python trajectory_compressor.py --input=data/my_run - - # Compress a single JSONL file - python trajectory_compressor.py --input=data/trajectories.jsonl - - # Compress 15% sample of a file + python trajectory_compressor.py --input=data/my_run # directory python trajectory_compressor.py --input=data/trajectories.jsonl --sample_percent=15 - - # Compress with custom output and token target - python trajectory_compressor.py --input=data/trajectories.jsonl --output=compressed.jsonl --target_max_tokens=16000 - - # Compress 10% sample from a directory - python trajectory_compressor.py --input=data/my_run --sample_percent=10 + python trajectory_compressor.py --input=data/trajectories.jsonl --output=out.jsonl --target_max_tokens=16000 """ import json import os +import random +import shutil +import tempfile import time import yaml import logging import asyncio from pathlib import Path -from typing import List, Dict, Any, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple from dataclasses import dataclass, field from datetime import datetime @@ -57,12 +41,10 @@ load_hermes_dotenv(hermes_home=_hermes_home, project_env=_project_env) def _response_finish_reason(response: Any) -> str: - """Return ``choices[0].finish_reason`` from a dict- or object-shaped response. + """Lowercased ``choices[0].finish_reason`` of a dict/object response, ``""`` if absent. - Local copy of ``agent.context_compressor._response_finish_reason`` — - trajectory_compressor is a standalone CLI tool and deliberately avoids - importing the (heavy) context compressor module. Returns the lowercased - finish reason, or ``""`` when absent/unreadable. + Local copy of ``agent.context_compressor._response_finish_reason``: this + standalone CLI deliberately avoids importing the heavy context compressor. """ try: if isinstance(response, dict): @@ -83,13 +65,14 @@ def _response_finish_reason(response: Any) -> str: def _effective_temperature_for_model( model: str, - requested_temperature: float, + requested_temperature: Optional[float], base_url: Optional[str] = None, ) -> Optional[float]: """Apply fixed model temperature contracts to direct client calls. Returns ``None`` when the model manages temperature server-side (Kimi); callers must omit the ``temperature`` kwarg entirely in that case. + Shared with ``mini_swe_runner`` (which passes ``requested_temperature=None``). """ try: from agent.auxiliary_client import _fixed_temperature_for_model, OMIT_TEMPERATURE @@ -104,24 +87,91 @@ def _effective_temperature_for_model( return requested_temperature +def _load_jsonl( + path: Path, + on_error: Optional[Callable[[int, json.JSONDecodeError], None]] = None, + start: int = 0, +) -> List[Tuple[int, Any]]: + """Return ``(line_num, entry)`` for each non-blank line; bad lines go to ``on_error``.""" + entries = [] + with open(path, 'r', encoding='utf-8') as f: + for line_num, line in enumerate(f, start): + line = line.strip() + if not line: + continue + try: + entries.append((line_num, json.loads(line))) + except json.JSONDecodeError as e: + if on_error is not None: + on_error(line_num, e) + return entries + + +def _write_jsonl(path: Path, entries) -> None: + with open(path, 'w', encoding='utf-8') as f: + for entry in entries: + f.write(json.dumps(entry, ensure_ascii=False) + '\n') + + +# YAML section -> ((yaml_key, CompressionConfig attribute), ...) +_YAML_SECTIONS: Dict[str, Tuple[Tuple[str, str], ...]] = { + "tokenizer": (("name", "tokenizer_name"), ("trust_remote_code", "trust_remote_code")), + "compression": ( + ("target_max_tokens", "target_max_tokens"), + ("summary_target_tokens", "summary_target_tokens"), + ), + "protected_turns": ( + ("first_system", "protect_first_system"), + ("first_human", "protect_first_human"), + ("first_gpt", "protect_first_gpt"), + ("first_tool", "protect_first_tool"), + ("last_n_turns", "protect_last_n_turns"), + ), + "summarization": ( + ("model", "summarization_model"), + ("base_url", "base_url"), + ("api_key_env", "api_key_env"), + ("temperature", "temperature"), + ("max_retries", "max_retries"), + ("retry_delay", "retry_delay"), + ), + "output": ( + ("add_summary_notice", "add_summary_notice"), + ("summary_notice_text", "summary_notice_text"), + ("output_suffix", "output_suffix"), + ), + "processing": ( + ("num_workers", "num_workers"), + ("max_concurrent_requests", "max_concurrent_requests"), + ("skip_under_target", "skip_under_target"), + ("save_over_limit", "save_over_limit"), + ), + "metrics": ( + ("enabled", "metrics_enabled"), + ("per_trajectory", "metrics_per_trajectory"), + ("output_file", "metrics_output_file"), + ), +} + + @dataclass class CompressionConfig: """Configuration for trajectory compression.""" # Tokenizer tokenizer_name: str = "moonshotai/Kimi-K2-Thinking" trust_remote_code: bool = True - + # Compression targets target_max_tokens: int = 15250 summary_target_tokens: int = 750 - + # Protected turns protect_first_system: bool = True protect_first_human: bool = True protect_first_gpt: bool = True protect_first_tool: bool = True protect_last_n_turns: int = 4 - + # Summarization (OpenRouter) summarization_model: str = "google/gemini-3-flash-preview" base_url: str = OPENROUTER_BASE_URL @@ -129,78 +179,39 @@ class CompressionConfig: temperature: float = 0.3 max_retries: int = 3 retry_delay: int = 2 - + # Output add_summary_notice: bool = True summary_notice_text: str = "\n\nSome of your previous tool responses may be summarized to preserve context." output_suffix: str = "_compressed" - + # Processing num_workers: int = 4 max_concurrent_requests: int = 50 # Max concurrent API calls for summarization skip_under_target: bool = True save_over_limit: bool = True per_trajectory_timeout: int = 300 # Timeout per trajectory in seconds (default: 5 min) - + # Metrics metrics_enabled: bool = True metrics_per_trajectory: bool = True metrics_output_file: str = "compression_metrics.json" - + @classmethod def from_yaml(cls, yaml_path: str) -> "CompressionConfig": - """Load configuration from YAML file.""" + """Load configuration from YAML file (missing keys keep the defaults).""" with open(yaml_path, 'r', encoding="utf-8") as f: data = yaml.safe_load(f) or {} config = cls() - - # Tokenizer - if 'tokenizer' in data: - config.tokenizer_name = data['tokenizer'].get('name', config.tokenizer_name) - config.trust_remote_code = data['tokenizer'].get('trust_remote_code', config.trust_remote_code) - - # Compression - if 'compression' in data: - config.target_max_tokens = data['compression'].get('target_max_tokens', config.target_max_tokens) - config.summary_target_tokens = data['compression'].get('summary_target_tokens', config.summary_target_tokens) - - # Protected turns - if 'protected_turns' in data: - config.protect_first_system = data['protected_turns'].get('first_system', config.protect_first_system) - config.protect_first_human = data['protected_turns'].get('first_human', config.protect_first_human) - config.protect_first_gpt = data['protected_turns'].get('first_gpt', config.protect_first_gpt) - config.protect_first_tool = data['protected_turns'].get('first_tool', config.protect_first_tool) - config.protect_last_n_turns = data['protected_turns'].get('last_n_turns', config.protect_last_n_turns) - - # Summarization - if 'summarization' in data: - config.summarization_model = data['summarization'].get('model', config.summarization_model) - config.base_url = data['summarization'].get('base_url') or config.base_url - config.api_key_env = data['summarization'].get('api_key_env', config.api_key_env) - config.temperature = data['summarization'].get('temperature', config.temperature) - config.max_retries = data['summarization'].get('max_retries', config.max_retries) - config.retry_delay = data['summarization'].get('retry_delay', config.retry_delay) - - # Output - if 'output' in data: - config.add_summary_notice = data['output'].get('add_summary_notice', config.add_summary_notice) - config.summary_notice_text = data['output'].get('summary_notice_text', config.summary_notice_text) - config.output_suffix = data['output'].get('output_suffix', config.output_suffix) - - # Processing - if 'processing' in data: - config.num_workers = data['processing'].get('num_workers', config.num_workers) - config.max_concurrent_requests = data['processing'].get('max_concurrent_requests', config.max_concurrent_requests) - config.skip_under_target = data['processing'].get('skip_under_target', config.skip_under_target) - config.save_over_limit = data['processing'].get('save_over_limit', config.save_over_limit) - - # Metrics - if 'metrics' in data: - config.metrics_enabled = data['metrics'].get('enabled', config.metrics_enabled) - config.metrics_per_trajectory = data['metrics'].get('per_trajectory', config.metrics_per_trajectory) - config.metrics_output_file = data['metrics'].get('output_file', config.metrics_output_file) - + for section, keys in _YAML_SECTIONS.items(): + if section not in data: + continue + for yaml_key, attr in keys: + value = data[section].get(yaml_key, getattr(config, attr)) + if attr == "base_url": + value = value or config.base_url # ``base_url: null`` keeps the default + setattr(config, attr, value) return config @@ -211,22 +222,22 @@ class TrajectoryMetrics: compressed_tokens: int = 0 tokens_saved: int = 0 compression_ratio: float = 1.0 - + original_turns: int = 0 compressed_turns: int = 0 turns_removed: int = 0 - + turns_compressed_start_idx: int = -1 turns_compressed_end_idx: int = -1 turns_in_compressed_region: int = 0 - + was_compressed: bool = False still_over_limit: bool = False skipped_under_target: bool = False - + summarization_api_calls: int = 0 summarization_errors: int = 0 - + def to_dict(self) -> Dict[str, Any]: return { "original_tokens": self.original_tokens, @@ -249,7 +260,11 @@ class TrajectoryMetrics: } -@dataclass +def _mean(values, default): + return sum(values) / len(values) if values else default + + +@dataclass class AggregateMetrics: """Aggregate metrics across all trajectories.""" total_trajectories: int = 0 @@ -257,27 +272,27 @@ class AggregateMetrics: trajectories_skipped_under_target: int = 0 trajectories_still_over_limit: int = 0 trajectories_failed: int = 0 - + total_tokens_before: int = 0 total_tokens_after: int = 0 total_tokens_saved: int = 0 - + total_turns_before: int = 0 total_turns_after: int = 0 total_turns_removed: int = 0 - + total_summarization_calls: int = 0 total_summarization_errors: int = 0 - + # Distribution stats compression_ratios: List[float] = field(default_factory=list) tokens_saved_list: List[int] = field(default_factory=list) turns_removed_list: List[int] = field(default_factory=list) - + processing_start_time: str = "" processing_end_time: str = "" processing_duration_seconds: float = 0.0 - + def add_trajectory_metrics(self, metrics: TrajectoryMetrics): """Add a trajectory's metrics to the aggregate.""" self.total_trajectories += 1 @@ -289,33 +304,20 @@ class AggregateMetrics: self.total_turns_removed += metrics.turns_removed self.total_summarization_calls += metrics.summarization_api_calls self.total_summarization_errors += metrics.summarization_errors - + if metrics.was_compressed: self.trajectories_compressed += 1 self.compression_ratios.append(metrics.compression_ratio) self.tokens_saved_list.append(metrics.tokens_saved) self.turns_removed_list.append(metrics.turns_removed) - + if metrics.skipped_under_target: self.trajectories_skipped_under_target += 1 - + if metrics.still_over_limit: self.trajectories_still_over_limit += 1 - + def to_dict(self) -> Dict[str, Any]: - avg_compression_ratio = ( - sum(self.compression_ratios) / len(self.compression_ratios) - if self.compression_ratios else 1.0 - ) - avg_tokens_saved = ( - sum(self.tokens_saved_list) / len(self.tokens_saved_list) - if self.tokens_saved_list else 0 - ) - avg_turns_removed = ( - sum(self.turns_removed_list) / len(self.turns_removed_list) - if self.turns_removed_list else 0 - ) - return { "summary": { "total_trajectories": self.total_trajectories, @@ -337,9 +339,9 @@ class AggregateMetrics: "total_removed": self.total_turns_removed, }, "averages": { - "avg_compression_ratio": round(avg_compression_ratio, 4), - "avg_tokens_saved_per_compressed": round(avg_tokens_saved, 1), - "avg_turns_removed_per_compressed": round(avg_turns_removed, 2), + "avg_compression_ratio": round(_mean(self.compression_ratios, 1.0), 4), + "avg_tokens_saved_per_compressed": round(_mean(self.tokens_saved_list, 0), 1), + "avg_turns_removed_per_compressed": round(_mean(self.turns_removed_list, 0), 2), }, "summarization": { "total_api_calls": self.total_summarization_calls, @@ -354,31 +356,59 @@ class AggregateMetrics: } +# Ordered (hostname, provider) table for _detect_provider (codex is matched separately). +_PROVIDER_HOSTS: Tuple[Tuple[str, str], ...] = ( + ("openrouter.ai", "openrouter"), + ("nousresearch.com", "nous"), + ("z.ai", "zai"), + ("moonshot.ai", "kimi-coding"), + ("moonshot.cn", "kimi-coding"), + ("api.kimi.com", "kimi-coding"), + ("arcee.ai", "arcee"), + ("minimaxi.com", "minimax-cn"), + ("minimax.io", "minimax"), +) + +_SUMMARY_FALLBACK = "[CONTEXT SUMMARY]: [Summary generation failed - previous turns contained tool calls and responses that have been compressed to save context space.]" +_STATUS_FMT = "[dim]✅ {compressed} compressed | ⏭️ {skipped} skipped | ⏱️ {timeouts} timeout | 🔄 {api_calls} API calls | ⚡ {in_flight} in-flight[/dim]" + + +@dataclass +class _RunProgress: + """Shared counters + rich progress handles for one directory run.""" + progress: Any + main_task: Any + status_task: Any + lock: asyncio.Lock + semaphore: asyncio.Semaphore + compressed: int = 0 + skipped: int = 0 + api_calls: int = 0 + in_flight: int = 0 + timeouts: int = 0 + + def advance(self, update_status: bool = True) -> None: + self.progress.advance(self.main_task) + if update_status: + self.progress.update(self.status_task, description=_STATUS_FMT.format( + compressed=self.compressed, skipped=self.skipped, timeouts=self.timeouts, + api_calls=self.api_calls, in_flight=self.in_flight)) + + class TrajectoryCompressor: + """Compresses agent trajectories to fit within a target token budget. + + Keeps protected head/tail turns, summarizes only as much of the middle as + needed into one human summary turn, and keeps the remaining middle intact. """ - Compresses agent trajectories to fit within a target token budget. - - Compression strategy: - 1. Keep protected head turns (system, human, first gpt+tool) - 2. Keep protected tail turns (last N turns) - 3. From the compressible middle region, compress only as much as needed - 4. Replace compressed turns with a single human summary message - 5. Keep remaining middle turns intact (model continues with tools) - """ - + def __init__(self, config: CompressionConfig): - """Initialize the compressor.""" self.config = config self.aggregate_metrics = AggregateMetrics() - - # Initialize tokenizer self._init_tokenizer() - - # Initialize OpenRouter client self._init_summarizer() - self.logger = logging.getLogger(__name__) - + def _init_tokenizer(self): """Initialize HuggingFace tokenizer for token counting.""" try: @@ -390,21 +420,13 @@ class TrajectoryCompressor: print(f"✅ Loaded tokenizer: {self.config.tokenizer_name}") except Exception as e: raise RuntimeError(f"Failed to load tokenizer '{self.config.tokenizer_name}': {e}") - + def _init_summarizer(self): - """Initialize LLM routing for summarization (sync and async). - - Uses call_llm/async_call_llm from the centralized provider router - which handles auth, headers, and provider detection internally. - For custom endpoints, falls back to raw client construction. - """ - + """Route summarization through call_llm for known providers, else a raw client.""" provider = self._detect_provider() if provider: - # Store provider for use in _generate_summary calls self._llm_provider = provider self._use_call_llm = True - # Verify the provider is available from agent.auxiliary_client import resolve_provider_client client, _ = resolve_provider_client( provider, model=self.config.summarization_model) @@ -426,10 +448,9 @@ class TrajectoryCompressor: from agent.auxiliary_client import _to_openai_base_url self.client = OpenAI( api_key=api_key, base_url=_to_openai_base_url(self.config.base_url)) - # AsyncOpenAI is created lazily in _get_async_client() so it - # binds to the current event loop — avoids "Event loop is closed" - # when process_directory() is called multiple times (each call - # creates a new loop via asyncio.run()). + # AsyncOpenAI is created lazily in _get_async_client() so it binds to + # the current event loop — each process_directory() call runs its own + # asyncio.run(), and a shared client would hit "Event loop is closed". self.async_client = None self._async_client_api_key = api_key @@ -437,15 +458,9 @@ class TrajectoryCompressor: print(f" Max concurrent requests: {self.config.max_concurrent_requests}") def _get_async_client(self): - """Return an AsyncOpenAI client bound to the current event loop. - - Created lazily so that each ``asyncio.run()`` call in - ``process_directory()`` gets a client tied to its own loop, - avoiding "Event loop is closed" errors on repeated calls. - """ + """Return a fresh AsyncOpenAI client bound to the running event loop.""" from openai import AsyncOpenAI from agent.auxiliary_client import _to_openai_base_url - # Always create a fresh client so it binds to the running loop. self.async_client = AsyncOpenAI( api_key=self._async_client_api_key, base_url=_to_openai_base_url(self.config.base_url), @@ -453,95 +468,54 @@ class TrajectoryCompressor: return self.async_client def _detect_provider(self) -> str: - """Detect the provider name from the configured base_url.""" + """Provider name for the configured base_url, or ``""`` when unknown.""" url = self.config.base_url or "" - if base_url_host_matches(url, "openrouter.ai"): - return "openrouter" - if base_url_host_matches(url, "nousresearch.com"): - return "nous" - if ( - base_url_hostname(url) == "chatgpt.com" - and "/backend-api/codex" in url.lower() - ): + if base_url_hostname(url) == "chatgpt.com" and "/backend-api/codex" in url.lower(): return "codex" - if base_url_host_matches(url, "z.ai"): - return "zai" - if ( - base_url_host_matches(url, "moonshot.ai") - or base_url_host_matches(url, "moonshot.cn") - or base_url_host_matches(url, "api.kimi.com") - ): - return "kimi-coding" - if base_url_host_matches(url, "arcee.ai"): - return "arcee" - if base_url_host_matches(url, "minimaxi.com"): - return "minimax-cn" - if base_url_host_matches(url, "minimax.io"): - return "minimax" - # Unknown base_url — not a known provider + for host, provider in _PROVIDER_HOSTS: + if base_url_host_matches(url, host): + return provider return "" - + def count_tokens(self, text: str) -> int: - """Count tokens in text using the configured tokenizer.""" + """Token count via the configured tokenizer; falls back to len//4.""" if not text: return 0 try: return len(self.tokenizer.encode(text)) except Exception: - # Fallback to character estimate return len(text) // 4 - + def count_trajectory_tokens(self, trajectory: List[Dict[str, str]]) -> int: - """Count total tokens in a trajectory.""" - return sum(self.count_tokens(turn.get("value", "")) for turn in trajectory) - + return sum(self.count_turn_tokens(trajectory)) + def count_turn_tokens(self, trajectory: List[Dict[str, str]]) -> List[int]: - """Count tokens for each turn in a trajectory.""" return [self.count_tokens(turn.get("value", "")) for turn in trajectory] - + def _find_protected_indices(self, trajectory: List[Dict[str, str]]) -> Tuple[set, int, int]: - """ - Find indices of protected turns. - - Returns: - Tuple of (protected_set, compressible_start, compressible_end) - """ + """Return ``(protected_set, compressible_start, compressible_end)``.""" n = len(trajectory) protected = set() - - # Track first occurrences - first_system = first_human = first_gpt = first_tool = None - + + first_seen: Dict[str, int] = {} for i, turn in enumerate(trajectory): - role = turn.get("from", "") - if role == "system" and first_system is None: - first_system = i - elif role == "human" and first_human is None: - first_human = i - elif role == "gpt" and first_gpt is None: - first_gpt = i - elif role == "tool" and first_tool is None: - first_tool = i - - # Protect first turns - if self.config.protect_first_system and first_system is not None: - protected.add(first_system) - if self.config.protect_first_human and first_human is not None: - protected.add(first_human) - if self.config.protect_first_gpt and first_gpt is not None: - protected.add(first_gpt) - if self.config.protect_first_tool and first_tool is not None: - protected.add(first_tool) - - # Protect last N turns - for i in range(max(0, n - self.config.protect_last_n_turns), n): - protected.add(i) - - # Determine compressible region - # Start after the last protected head turn + first_seen.setdefault(turn.get("from", ""), i) + + for role, enabled in ( + ("system", self.config.protect_first_system), + ("human", self.config.protect_first_human), + ("gpt", self.config.protect_first_gpt), + ("tool", self.config.protect_first_tool), + ): + if enabled and role in first_seen: + protected.add(first_seen[role]) + + protected.update(range(max(0, n - self.config.protect_last_n_turns), n)) + + # Compressible region: after the last protected head turn, before the first tail turn. head_protected = [i for i in protected if i < n // 2] tail_protected = [i for i in protected if i >= n // 2] - + compressible_start = max(head_protected) + 1 if head_protected else 0 compressible_end = min(tail_protected) if tail_protected else n @@ -549,14 +523,11 @@ class TrajectoryCompressor: @staticmethod def _is_boundary_clean(trajectory: List[Dict[str, str]], idx: int) -> bool: - """Return True if a region boundary at ``idx`` does not split a turn pair. + """True if a boundary at ``idx`` does not split a gpt /tool pair. - In the from/value trajectory format a ``tool`` turn (carrying - ```` markers) is always emitted immediately after the - ``gpt`` turn whose ```` it answers. A compression boundary - that lands *on* a ``tool`` turn therefore cuts between a tool call and - its response. A boundary is only clean when it sits at the very end of - the trajectory or on a non-``tool`` turn. + A ``tool`` turn always directly follows the ``gpt`` turn it answers, so a + boundary landing *on* a tool turn cuts the pair; only the end of the + trajectory or a non-``tool`` turn is clean. """ return idx >= len(trajectory) or trajectory[idx].get("from") != "tool" @@ -568,13 +539,10 @@ class TrajectoryCompressor: min_idx: int, max_idx: int, ) -> int: - """Move a compression boundary onto the nearest clean turn boundary. + """Move a boundary onto the nearest clean turn boundary within ``[min_idx, max_idx]``. - Moving forward is preferred so that an orphaned ``tool`` turn is folded - into the region that already holds its ``gpt`` turn; if no clean - boundary exists ahead (for example the protected tail itself begins on a - ``tool`` turn) the boundary is moved backward instead. The result is - clamped to ``[min_idx, max_idx]``. + Forward is preferred (folds an orphaned ``tool`` turn into the region that + holds its ``gpt`` turn); backward only when nothing clean lies ahead. """ forward = idx while forward < max_idx and not cls._is_boundary_clean(trajectory, forward): @@ -587,29 +555,15 @@ class TrajectoryCompressor: return backward def _extract_turn_content_for_summary(self, trajectory: List[Dict[str, str]], start: int, end: int) -> str: - """ - Extract content from turns to be summarized. - - Args: - trajectory: Full trajectory - start: Start index (inclusive) - end: End index (exclusive) - - Returns: - Formatted string of turn contents for summarization - """ + """Format turns ``[start, end)`` for the summarization prompt (long values truncated).""" parts = [] for i in range(start, end): turn = trajectory[i] role = turn.get("from", "unknown") value = turn.get("value", "") - - # Truncate very long values for the summary prompt if len(value) > 3000: value = value[:1500] + "\n...[truncated]...\n" + value[-500:] - parts.append(f"[Turn {i} - {role.upper()}]:\n{value}") - return "\n\n".join(parts) @staticmethod @@ -626,19 +580,9 @@ class TrajectoryCompressor: if text.startswith("[CONTEXT SUMMARY]:"): return text return "[CONTEXT SUMMARY]:" if not text else f"[CONTEXT SUMMARY]: {text}" - - def _generate_summary(self, content: str, metrics: TrajectoryMetrics) -> str: - """ - Generate a summary of the compressed turns using OpenRouter. - - Args: - content: The content to summarize - metrics: Metrics object to update - - Returns: - Summary string - """ - prompt = f"""Summarize the following agent conversation turns concisely. This summary will replace these turns in the conversation history. + + def _summary_prompt(self, content: str) -> str: + return f"""Summarize the following agent conversation turns concisely. This summary will replace these turns in the conversation history. Write the summary from a neutral perspective describing what the assistant did and learned. Include: 1. What actions the assistant took (tool calls, searches, file operations) @@ -655,500 +599,255 @@ TURNS TO SUMMARIZE: Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" + def _summary_request(self, prompt: str) -> Tuple[Optional[float], Dict[str, Any]]: + """Return ``(temperature, create-kwargs)``; temperature None means omit it.""" + temperature = _effective_temperature_for_model( + self.config.summarization_model, + self.config.temperature, + self.config.base_url, + ) + kwargs = { + "model": self.config.summarization_model, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": self.config.summary_target_tokens * 2, + } + if not getattr(self, '_use_call_llm', False) and temperature is not None: + kwargs["temperature"] = temperature + return temperature, kwargs + + def _finish_summary(self, response: Any) -> str: + """Extract the summary text; a ``length`` stop is a failure (partial summary).""" + if _response_finish_reason(response) == "length": + # Storing a truncated summary silently corrupts the trajectory's + # memory, so raise and let the retry/backoff loop handle it. + raise RuntimeError( + "trajectory summarization hit the output token cap " + "(finish_reason=length); summary is incomplete" + ) + summary = self._coerce_summary_content(response.choices[0].message.content) + return self._ensure_summary_prefix(summary) + + def _summary_attempt_failed(self, metrics: TrajectoryMetrics, attempt: int, exc: Exception) -> Optional[float]: + """Record a failed attempt; return the backoff delay, or None on the last attempt.""" + metrics.summarization_errors += 1 + self.logger.warning("Summarization attempt %d failed: %s", attempt + 1, exc) + if attempt < self.config.max_retries - 1: + return jittered_backoff(attempt + 1, base_delay=self.config.retry_delay, max_delay=30.0) + return None + + def _generate_summary(self, content: str, metrics: TrajectoryMetrics) -> str: + """Summarize ``content`` with retries; returns a fallback summary after the last failure.""" + prompt = self._summary_prompt(content) for attempt in range(self.config.max_retries): try: metrics.summarization_api_calls += 1 - summary_temperature = _effective_temperature_for_model( - self.config.summarization_model, - self.config.temperature, - self.config.base_url, - ) - + temperature, kwargs = self._summary_request(prompt) if getattr(self, '_use_call_llm', False): from agent.auxiliary_client import call_llm - response = call_llm( - provider=self._llm_provider, - model=self.config.summarization_model, - messages=[{"role": "user", "content": prompt}], - temperature=summary_temperature, - max_tokens=self.config.summary_target_tokens * 2, - ) + response = call_llm(provider=self._llm_provider, temperature=temperature, **kwargs) else: - _create_kwargs = { - "model": self.config.summarization_model, - "messages": [{"role": "user", "content": prompt}], - "max_tokens": self.config.summary_target_tokens * 2, - } - if summary_temperature is not None: - _create_kwargs["temperature"] = summary_temperature - response = self.client.chat.completions.create(**_create_kwargs) - - _fr = _response_finish_reason(response) - if _fr == "length": - # Length stop = partial summary; storing it as the turn - # replacement silently truncates the trajectory's memory. - # Raise so the retry/backoff loop treats it as a failure - # (pi#7048 class). - raise RuntimeError( - "trajectory summarization hit the output token cap " - "(finish_reason=length); summary is incomplete" - ) - summary = self._coerce_summary_content(response.choices[0].message.content) - return self._ensure_summary_prefix(summary) - + response = self.client.chat.completions.create(**kwargs) + return self._finish_summary(response) except Exception as e: - metrics.summarization_errors += 1 - self.logger.warning("Summarization attempt %d failed: %s", attempt + 1, e) - - if attempt < self.config.max_retries - 1: - time.sleep(jittered_backoff(attempt + 1, base_delay=self.config.retry_delay, max_delay=30.0)) - else: - # Fallback: create a basic summary - return "[CONTEXT SUMMARY]: [Summary generation failed - previous turns contained tool calls and responses that have been compressed to save context space.]" - + delay = self._summary_attempt_failed(metrics, attempt, e) + if delay is None: + return _SUMMARY_FALLBACK + time.sleep(delay) + async def _generate_summary_async(self, content: str, metrics: TrajectoryMetrics) -> str: - """ - Generate a summary of the compressed turns using OpenRouter (async version). - - Args: - content: The content to summarize - metrics: Metrics object to update - - Returns: - Summary string - """ - prompt = f"""Summarize the following agent conversation turns concisely. This summary will replace these turns in the conversation history. - -Write the summary from a neutral perspective describing what the assistant did and learned. Include: -1. What actions the assistant took (tool calls, searches, file operations) -2. Key information or results obtained -3. Any important decisions or findings -4. Relevant data, file names, values, or outputs - -Keep the summary factual and informative. Target approximately {self.config.summary_target_tokens} tokens. - ---- -TURNS TO SUMMARIZE: -{content} ---- - -Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" - + """Async twin of ``_generate_summary``.""" + prompt = self._summary_prompt(content) for attempt in range(self.config.max_retries): try: metrics.summarization_api_calls += 1 - summary_temperature = _effective_temperature_for_model( - self.config.summarization_model, - self.config.temperature, - self.config.base_url, - ) - + temperature, kwargs = self._summary_request(prompt) if getattr(self, '_use_call_llm', False): from agent.auxiliary_client import async_call_llm - response = await async_call_llm( - provider=self._llm_provider, - model=self.config.summarization_model, - messages=[{"role": "user", "content": prompt}], - temperature=summary_temperature, - max_tokens=self.config.summary_target_tokens * 2, - ) + response = await async_call_llm(provider=self._llm_provider, temperature=temperature, **kwargs) else: - _create_kwargs = { - "model": self.config.summarization_model, - "messages": [{"role": "user", "content": prompt}], - "max_tokens": self.config.summary_target_tokens * 2, - } - if summary_temperature is not None: - _create_kwargs["temperature"] = summary_temperature - response = await self._get_async_client().chat.completions.create(**_create_kwargs) - - if _response_finish_reason(response) == "length": - # Length stop = partial summary; see sync sibling above - # (pi#7048 class). - raise RuntimeError( - "trajectory summarization hit the output token cap " - "(finish_reason=length); summary is incomplete" - ) - summary = self._coerce_summary_content(response.choices[0].message.content) - return self._ensure_summary_prefix(summary) - + response = await self._get_async_client().chat.completions.create(**kwargs) + return self._finish_summary(response) except Exception as e: - metrics.summarization_errors += 1 - self.logger.warning("Summarization attempt %d failed: %s", attempt + 1, e) - - if attempt < self.config.max_retries - 1: - await asyncio.sleep(jittered_backoff(attempt + 1, base_delay=self.config.retry_delay, max_delay=30.0)) - else: - # Fallback: create a basic summary - return "[CONTEXT SUMMARY]: [Summary generation failed - previous turns contained tool calls and responses that have been compressed to save context space.]" - + delay = self._summary_attempt_failed(metrics, attempt, e) + if delay is None: + return _SUMMARY_FALLBACK + await asyncio.sleep(delay) + + def _plan_compression( + self, trajectory: List[Dict[str, str]], metrics: TrajectoryMetrics + ) -> Optional[Tuple[int, int]]: + """Choose the ``[start, until)`` region to summarize, or None if nothing can be. + + Fills the pre-compression metrics either way. Accumulates turns from the + start of the compressible middle until the savings cover the overage plus + the summary itself, then snaps both boundaries off ``tool`` turns. + """ + cfg = self.config + turn_tokens = self.count_turn_tokens(trajectory) + total_tokens = sum(turn_tokens) + metrics.original_turns = metrics.compressed_turns = len(trajectory) + metrics.original_tokens = metrics.compressed_tokens = total_tokens + + if total_tokens <= cfg.target_max_tokens: + metrics.skipped_under_target = True + return None + metrics.still_over_limit = True + + _, start, end = self._find_protected_indices(trajectory) + # Never *start* on an orphaned whose is in the protected head. + start = self._snap_boundary(trajectory, start, start, end) + if start >= end: + return None + + # Replacing N turns with one summary saves sum(N) - summary_target_tokens. + target_tokens_to_compress = total_tokens - cfg.target_max_tokens + cfg.summary_target_tokens + accumulated = 0 + until = start + for i in range(start, end): + accumulated += turn_tokens[i] + until = i + 1 + if accumulated >= target_tokens_to_compress: + break + if accumulated < target_tokens_to_compress and until < end: + until = end + + # The remainder is kept verbatim, so a tail boundary on a tool turn would orphan a marker. + until = self._snap_boundary(trajectory, until, start, end) + # A region no larger than the summary replacing it cannot shrink the trajectory. + if until <= start or sum(turn_tokens[start:until]) <= cfg.summary_target_tokens: + return None + + metrics.turns_compressed_start_idx = start + metrics.turns_compressed_end_idx = until + metrics.turns_in_compressed_region = until - start + return start, until + + def _assemble_compressed( + self, + trajectory: List[Dict[str, str]], + start: int, + until: int, + summary: str, + metrics: TrajectoryMetrics, + ) -> List[Dict[str, str]]: + """Head (with summary notice on system) + summary human turn + verbatim tail; finalize metrics.""" + compressed = [] + for turn in trajectory[:start]: + turn = turn.copy() + if turn.get("from") == "system" and self.config.add_summary_notice: + turn["value"] = turn["value"] + self.config.summary_notice_text + compressed.append(turn) + compressed.append({"from": "human", "value": summary}) + compressed.extend(turn.copy() for turn in trajectory[until:]) + + metrics.compressed_turns = len(compressed) + metrics.compressed_tokens = self.count_trajectory_tokens(compressed) + metrics.turns_removed = metrics.original_turns - metrics.compressed_turns + metrics.tokens_saved = metrics.original_tokens - metrics.compressed_tokens + metrics.compression_ratio = metrics.compressed_tokens / max(metrics.original_tokens, 1) + metrics.was_compressed = True + metrics.still_over_limit = metrics.compressed_tokens > self.config.target_max_tokens + return compressed + def compress_trajectory( self, trajectory: List[Dict[str, str]] ) -> Tuple[List[Dict[str, str]], TrajectoryMetrics]: - """ - Compress a single trajectory to fit within target token budget. - - Algorithm: - 1. Count total tokens - 2. If under target, skip - 3. Find compressible region (between protected head and tail) - 4. Calculate how many tokens need to be saved - 5. Accumulate turns from start of compressible region until savings met - 6. Replace accumulated turns with single human summary - 7. Keep remaining turns intact - - Args: - trajectory: List of conversation turns - - Returns: - Tuple of (compressed_trajectory, metrics) - """ + """Compress one trajectory into the target budget; returns ``(trajectory, metrics)``.""" metrics = TrajectoryMetrics() - metrics.original_turns = len(trajectory) - - # Count tokens per turn - turn_tokens = self.count_turn_tokens(trajectory) - total_tokens = sum(turn_tokens) - metrics.original_tokens = total_tokens - - # Check if compression needed - if total_tokens <= self.config.target_max_tokens: - metrics.skipped_under_target = True - metrics.compressed_tokens = total_tokens - metrics.compressed_turns = len(trajectory) - metrics.compression_ratio = 1.0 + region = self._plan_compression(trajectory, metrics) + if region is None: return trajectory, metrics - - # Find protected regions - protected, compress_start, compress_end = self._find_protected_indices(trajectory) + start, until = region + summary = self._generate_summary( + self._extract_turn_content_for_summary(trajectory, start, until), metrics) + return self._assemble_compressed(trajectory, start, until, summary, metrics), metrics - # Snap the head boundary so the compressible region never *starts* on an - # orphaned whose lives in the protected head. - compress_start = self._snap_boundary(trajectory, compress_start, compress_start, compress_end) - - # Check if there's anything to compress - if compress_start >= compress_end: - # Nothing to compress, return as-is - metrics.compressed_tokens = total_tokens - metrics.compressed_turns = len(trajectory) - metrics.still_over_limit = total_tokens > self.config.target_max_tokens - return trajectory, metrics - - # Calculate how much we need to save - tokens_to_save = total_tokens - self.config.target_max_tokens - - # We'll replace N turns with 1 summary turn - # Net savings = (sum of N turns' tokens) - summary_target_tokens - # We need: net_savings >= tokens_to_save - # So: sum of turns >= tokens_to_save + summary_target_tokens - target_tokens_to_compress = tokens_to_save + self.config.summary_target_tokens - - # Accumulate turns from compress_start until we have enough savings - accumulated_tokens = 0 - compress_until = compress_start - - for i in range(compress_start, compress_end): - accumulated_tokens += turn_tokens[i] - compress_until = i + 1 # Exclusive end - - # Check if we have enough savings - if accumulated_tokens >= target_tokens_to_compress: - break - - # If we still don't have enough savings, compress the entire compressible region - if accumulated_tokens < target_tokens_to_compress and compress_until < compress_end: - compress_until = compress_end - accumulated_tokens = sum(turn_tokens[compress_start:compress_end]) - - # Snap the tail boundary so we never cut between a and its - # : the summary replaces [compress_start, compress_until) - # and the remainder is kept verbatim, so a boundary on a tool turn would - # leave an orphaned marker and corrupt the training trajectory. - compress_until = self._snap_boundary(trajectory, compress_until, compress_start, compress_end) - if compress_until <= compress_start: - # Snapping collapsed the region; nothing can be safely compressed. - metrics.compressed_tokens = total_tokens - metrics.compressed_turns = len(trajectory) - metrics.still_over_limit = total_tokens > self.config.target_max_tokens - return trajectory, metrics - - # If the region we can safely compress is no larger than the summary - # that would replace it, compression cannot reduce the token count -- - # it would grow the trajectory and still spend a summarization call. - if ( - sum(turn_tokens[compress_start:compress_until]) - <= self.config.summary_target_tokens - ): - metrics.compressed_tokens = total_tokens - metrics.compressed_turns = len(trajectory) - metrics.still_over_limit = total_tokens > self.config.target_max_tokens - return trajectory, metrics - - # Record compression region - metrics.turns_compressed_start_idx = compress_start - metrics.turns_compressed_end_idx = compress_until - metrics.turns_in_compressed_region = compress_until - compress_start - - # Extract content for summary - content_to_summarize = self._extract_turn_content_for_summary( - trajectory, compress_start, compress_until - ) - - # Generate summary - summary = self._generate_summary(content_to_summarize, metrics) - - # Build compressed trajectory - compressed = [] - - # Add head (turns before compression region) - for i in range(compress_start): - turn = trajectory[i].copy() - # Add notice to system message - if turn.get("from") == "system" and self.config.add_summary_notice: - turn["value"] = turn["value"] + self.config.summary_notice_text - compressed.append(turn) - - # Add summary as human message - compressed.append({ - "from": "human", - "value": summary - }) - - # Add tail (turns after compression region) - for i in range(compress_until, len(trajectory)): - compressed.append(trajectory[i].copy()) - - # Calculate final metrics - metrics.compressed_turns = len(compressed) - metrics.compressed_tokens = self.count_trajectory_tokens(compressed) - metrics.turns_removed = metrics.original_turns - metrics.compressed_turns - metrics.tokens_saved = metrics.original_tokens - metrics.compressed_tokens - metrics.compression_ratio = metrics.compressed_tokens / max(metrics.original_tokens, 1) - metrics.was_compressed = True - metrics.still_over_limit = metrics.compressed_tokens > self.config.target_max_tokens - - return compressed, metrics - async def compress_trajectory_async( self, trajectory: List[Dict[str, str]] ) -> Tuple[List[Dict[str, str]], TrajectoryMetrics]: - """ - Compress a single trajectory to fit within target token budget (async version). - - Same algorithm as compress_trajectory but uses async API calls for summarization. - """ + """Async twin of ``compress_trajectory``.""" metrics = TrajectoryMetrics() - metrics.original_turns = len(trajectory) - - # Count tokens per turn - turn_tokens = self.count_turn_tokens(trajectory) - total_tokens = sum(turn_tokens) - metrics.original_tokens = total_tokens - - # Check if compression needed - if total_tokens <= self.config.target_max_tokens: - metrics.skipped_under_target = True - metrics.compressed_tokens = total_tokens - metrics.compressed_turns = len(trajectory) - metrics.compression_ratio = 1.0 + region = self._plan_compression(trajectory, metrics) + if region is None: return trajectory, metrics - - # Find protected regions - protected, compress_start, compress_end = self._find_protected_indices(trajectory) + start, until = region + summary = await self._generate_summary_async( + self._extract_turn_content_for_summary(trajectory, start, until), metrics) + return self._assemble_compressed(trajectory, start, until, summary, metrics), metrics - # Snap the head boundary so the compressible region never *starts* on an - # orphaned whose lives in the protected head. - compress_start = self._snap_boundary(trajectory, compress_start, compress_start, compress_end) - - # Check if there's anything to compress - if compress_start >= compress_end: - metrics.compressed_tokens = total_tokens - metrics.compressed_turns = len(trajectory) - metrics.still_over_limit = total_tokens > self.config.target_max_tokens - return trajectory, metrics - - # Calculate how much we need to save - tokens_to_save = total_tokens - self.config.target_max_tokens - target_tokens_to_compress = tokens_to_save + self.config.summary_target_tokens - - # Accumulate turns from compress_start until we have enough savings - accumulated_tokens = 0 - compress_until = compress_start - - for i in range(compress_start, compress_end): - accumulated_tokens += turn_tokens[i] - compress_until = i + 1 - if accumulated_tokens >= target_tokens_to_compress: - break - - # If we still don't have enough savings, compress the entire compressible region - if accumulated_tokens < target_tokens_to_compress and compress_until < compress_end: - compress_until = compress_end - accumulated_tokens = sum(turn_tokens[compress_start:compress_end]) - - # Snap the tail boundary so we never cut between a and its - # : the summary replaces [compress_start, compress_until) - # and the remainder is kept verbatim, so a boundary on a tool turn would - # leave an orphaned marker and corrupt the training trajectory. - compress_until = self._snap_boundary(trajectory, compress_until, compress_start, compress_end) - if compress_until <= compress_start: - # Snapping collapsed the region; nothing can be safely compressed. - metrics.compressed_tokens = total_tokens - metrics.compressed_turns = len(trajectory) - metrics.still_over_limit = total_tokens > self.config.target_max_tokens - return trajectory, metrics - - # If the region we can safely compress is no larger than the summary - # that would replace it, compression cannot reduce the token count -- - # it would grow the trajectory and still spend a summarization call. - if ( - sum(turn_tokens[compress_start:compress_until]) - <= self.config.summary_target_tokens - ): - metrics.compressed_tokens = total_tokens - metrics.compressed_turns = len(trajectory) - metrics.still_over_limit = total_tokens > self.config.target_max_tokens - return trajectory, metrics - - # Record compression region - metrics.turns_compressed_start_idx = compress_start - metrics.turns_compressed_end_idx = compress_until - metrics.turns_in_compressed_region = compress_until - compress_start - - # Extract content for summary - content_to_summarize = self._extract_turn_content_for_summary( - trajectory, compress_start, compress_until - ) - - # Generate summary (ASYNC) - summary = await self._generate_summary_async(content_to_summarize, metrics) - - # Build compressed trajectory - compressed = [] - - # Add head (turns before compression region) - for i in range(compress_start): - turn = trajectory[i].copy() - if turn.get("from") == "system" and self.config.add_summary_notice: - turn["value"] = turn["value"] + self.config.summary_notice_text - compressed.append(turn) - - # Add summary as human message - compressed.append({ - "from": "human", - "value": summary - }) - - # Add tail (turns after compression region) - for i in range(compress_until, len(trajectory)): - compressed.append(trajectory[i].copy()) - - # Calculate final metrics - metrics.compressed_turns = len(compressed) - metrics.compressed_tokens = self.count_trajectory_tokens(compressed) - metrics.turns_removed = metrics.original_turns - metrics.compressed_turns - metrics.tokens_saved = metrics.original_tokens - metrics.compressed_tokens - metrics.compression_ratio = metrics.compressed_tokens / max(metrics.original_tokens, 1) - metrics.was_compressed = True - metrics.still_over_limit = metrics.compressed_tokens > self.config.target_max_tokens - - return compressed, metrics - async def process_entry_async(self, entry: Dict[str, Any]) -> Tuple[Dict[str, Any], TrajectoryMetrics]: - """ - Process a single JSONL entry (async version). - """ + """Compress one JSONL entry's ``conversations``; attach metrics when compressed.""" if "conversations" not in entry: - metrics = TrajectoryMetrics() - return entry, metrics - - trajectory = entry["conversations"] - compressed_trajectory, metrics = await self.compress_trajectory_async(trajectory) - - # Create new entry with compressed trajectory + return entry, TrajectoryMetrics() + compressed_trajectory, metrics = await self.compress_trajectory_async(entry["conversations"]) result = entry.copy() result["conversations"] = compressed_trajectory - - # Add compression metadata if enabled if self.config.metrics_per_trajectory and metrics.was_compressed: result["compression_metrics"] = metrics.to_dict() - return result, metrics - - def process_entry(self, entry: Dict[str, Any]) -> Tuple[Dict[str, Any], TrajectoryMetrics]: - """ - Process a single JSONL entry. - - Args: - entry: JSONL entry containing 'conversations' field - - Returns: - Tuple of (processed_entry, metrics) - """ - if "conversations" not in entry: - metrics = TrajectoryMetrics() - return entry, metrics - - trajectory = entry["conversations"] - compressed_trajectory, metrics = self.compress_trajectory(trajectory) - - # Create new entry with compressed trajectory - result = entry.copy() - result["conversations"] = compressed_trajectory - - # Add compression metadata if enabled - if self.config.metrics_per_trajectory and metrics.was_compressed: - result["compression_metrics"] = metrics.to_dict() - - return result, metrics - + def process_directory(self, input_dir: Path, output_dir: Path): - """ - Process all JSONL files in a directory using async parallel processing. - - Args: - input_dir: Input directory containing JSONL files - output_dir: Output directory for compressed files - """ - # Run the async version + """Compress every ``*.jsonl`` in ``input_dir`` into ``output_dir`` (async, parallel API calls).""" asyncio.run(self._process_directory_async(input_dir, output_dir)) - + + async def _process_one( + self, run: _RunProgress, file_path: Path, entry_idx: int, entry: Dict + ) -> Optional[Tuple[Dict[str, Any], TrajectoryMetrics]]: + """Process one entry under the semaphore/timeout; None means dropped (timed out).""" + async with run.semaphore: + async with run.lock: + run.in_flight += 1 + try: + processed_entry, metrics = await asyncio.wait_for( + self.process_entry_async(entry), + timeout=self.config.per_trajectory_timeout + ) + async with run.lock: + self.aggregate_metrics.add_trajectory_metrics(metrics) + if metrics.was_compressed: + run.compressed += 1 + run.api_calls += metrics.summarization_api_calls + if metrics.skipped_under_target: + run.skipped += 1 + run.in_flight -= 1 + run.advance() + return processed_entry, metrics + except asyncio.TimeoutError: + self.logger.warning("Timeout processing entry from %s:%s (>%ss)", file_path, entry_idx, self.config.per_trajectory_timeout) + async with run.lock: + self.aggregate_metrics.trajectories_failed += 1 + run.timeouts += 1 + run.in_flight -= 1 + run.advance() + return None + except Exception as e: + self.logger.error("Error processing entry from %s:%s: %s", file_path, entry_idx, e) + async with run.lock: + self.aggregate_metrics.trajectories_failed += 1 + run.in_flight -= 1 + run.advance(update_status=False) + return entry, TrajectoryMetrics() # keep the original on error + async def _process_directory_async(self, input_dir: Path, output_dir: Path): - """ - Async implementation of directory processing with parallel API calls. - """ console = Console() - - # Record start time self.aggregate_metrics.processing_start_time = datetime.now().isoformat() start_time = time.time() - - # Find all JSONL files + jsonl_files = sorted(input_dir.glob("*.jsonl")) - if not jsonl_files: self.logger.warning("No JSONL files found in %s", input_dir) return - - # Load ALL entries from all files + console.print("\n[dim]Loading all entries...[/dim]") all_entries = [] # List of (file_path, entry_idx, entry) - for file_path in jsonl_files: - with open(file_path, 'r', encoding='utf-8') as f: - for line_num, line in enumerate(f): - line = line.strip() - if line: - try: - entry = json.loads(line) - all_entries.append((file_path, line_num, entry)) - except json.JSONDecodeError as e: - self.logger.warning("Skipping invalid JSON at %s:%s: %s", file_path, line_num, e) - + def _warn(line_num, e, file_path=file_path): + self.logger.warning("Skipping invalid JSON at %s:%s: %s", file_path, line_num, e) + all_entries.extend((file_path, idx, entry) for idx, entry in _load_jsonl(file_path, _warn)) total_entries = len(all_entries) - + console.print(f"\n{'='*60}") console.print(f"📂 Input: {input_dir}") console.print(f"📂 Output: {output_dir}") @@ -1158,89 +857,7 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" console.print(f"📝 Summary target tokens: {self.config.summary_target_tokens}") console.print(f"⚡ Max concurrent API calls: {self.config.max_concurrent_requests}") console.print(f"{'='*60}\n") - - # Create semaphore for rate limiting - semaphore = asyncio.Semaphore(self.config.max_concurrent_requests) - - # Tracking for progress display (thread-safe with lock) - progress_lock = asyncio.Lock() - compressed_count = 0 - skipped_count = 0 - api_calls = 0 - in_flight = 0 - - # Results storage: {file_path: {entry_idx: (processed_entry, metrics)}} - results = {f: {} for f in jsonl_files} - - # Track timeouts separately - timeout_count = 0 - - async def process_single(file_path: Path, entry_idx: int, entry: Dict, - progress, main_task, status_task): - """Process a single entry with semaphore rate limiting and timeout.""" - nonlocal compressed_count, skipped_count, api_calls, in_flight, timeout_count - - async with semaphore: - # Track in-flight - async with progress_lock: - in_flight += 1 - - try: - # Apply per-trajectory timeout - processed_entry, metrics = await asyncio.wait_for( - self.process_entry_async(entry), - timeout=self.config.per_trajectory_timeout - ) - results[file_path][entry_idx] = (processed_entry, metrics) - - # Update aggregate metrics (with lock for thread safety) - async with progress_lock: - self.aggregate_metrics.add_trajectory_metrics(metrics) - - # Update counters - if metrics.was_compressed: - compressed_count += 1 - api_calls += metrics.summarization_api_calls - if metrics.skipped_under_target: - skipped_count += 1 - - in_flight -= 1 - - # Update progress - progress.advance(main_task) - progress.update( - status_task, - description=f"[dim]✅ {compressed_count} compressed | ⏭️ {skipped_count} skipped | ⏱️ {timeout_count} timeout | 🔄 {api_calls} API calls | ⚡ {in_flight} in-flight[/dim]" - ) - - except asyncio.TimeoutError: - self.logger.warning("Timeout processing entry from %s:%s (>%ss)", file_path, entry_idx, self.config.per_trajectory_timeout) - - async with progress_lock: - self.aggregate_metrics.trajectories_failed += 1 - timeout_count += 1 - in_flight -= 1 - progress.advance(main_task) - progress.update( - status_task, - description=f"[dim]✅ {compressed_count} compressed | ⏭️ {skipped_count} skipped | ⏱️ {timeout_count} timeout | 🔄 {api_calls} API calls | ⚡ {in_flight} in-flight[/dim]" - ) - - # Skip this entry entirely (don't include in output) - results[file_path][entry_idx] = None - - except Exception as e: - self.logger.error("Error processing entry from %s:%s: %s", file_path, entry_idx, e) - - async with progress_lock: - self.aggregate_metrics.trajectories_failed += 1 - in_flight -= 1 - progress.advance(main_task) - - # Keep original entry on error - results[file_path][entry_idx] = (entry, TrajectoryMetrics()) - - # Create progress bar + with Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), @@ -1253,90 +870,62 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" console=console, refresh_per_second=10 # Higher refresh for async ) as progress: - # Main task for overall progress - main_task = progress.add_task( - f"[cyan]Compressing {total_entries:,} trajectories", - total=total_entries + run = _RunProgress( + progress=progress, + main_task=progress.add_task(f"[cyan]Compressing {total_entries:,} trajectories", total=total_entries), + status_task=progress.add_task("[dim]Starting...[/dim]", total=None), + lock=asyncio.Lock(), + semaphore=asyncio.Semaphore(self.config.max_concurrent_requests), ) - - # Status line task - status_task = progress.add_task( - "[dim]Starting...[/dim]", - total=None - ) - - # Create all tasks - tasks = [ - process_single(file_path, entry_idx, entry, progress, main_task, status_task) + outcomes = await asyncio.gather(*( + self._process_one(run, file_path, entry_idx, entry) for file_path, entry_idx, entry in all_entries - ] - - # Run all tasks concurrently (semaphore limits actual concurrency) - await asyncio.gather(*tasks) - - # Remove status task - progress.remove_task(status_task) - - # Write results to output files (preserving original order) + )) + progress.remove_task(run.status_task) + + # Write results preserving original order; timed-out entries are dropped. console.print("\n[dim]Writing output files...[/dim]") output_dir.mkdir(parents=True, exist_ok=True) - + results = {f: [] for f in jsonl_files} + for (file_path, _, _), outcome in zip(all_entries, outcomes): + if outcome is not None: + results[file_path].append(outcome[0]) for file_path in jsonl_files: - output_path = output_dir / file_path.name - file_results = results[file_path] - - # Sort by original entry index to preserve order, skip None (timed out) entries - sorted_entries = [ - file_results[idx][0] - for idx in sorted(file_results.keys()) - if file_results[idx] is not None - ] - - with open(output_path, 'w', encoding='utf-8') as f: - for entry in sorted_entries: - f.write(json.dumps(entry, ensure_ascii=False) + '\n') - - # Record end time + _write_jsonl(output_dir / file_path.name, results[file_path]) + self.aggregate_metrics.processing_end_time = datetime.now().isoformat() self.aggregate_metrics.processing_duration_seconds = time.time() - start_time - - # Print summary self._print_summary() - - # Save metrics + if self.config.metrics_enabled: metrics_path = output_dir / self.config.metrics_output_file with open(metrics_path, 'w', encoding="utf-8") as f: json.dump(self.aggregate_metrics.to_dict(), f, indent=2) console.print(f"\n💾 Metrics saved to {metrics_path}") - + def _print_summary(self): """Print comprehensive compression summary statistics.""" m = self.aggregate_metrics.to_dict() - - # Calculate some additional stats + total = m['summary']['total_trajectories'] compressed = m['summary']['trajectories_compressed'] skipped = m['summary']['trajectories_skipped_under_target'] over_limit = m['summary']['trajectories_still_over_limit'] failed = m['summary']['trajectories_failed'] - - # Token stats + tokens_before = m['tokens']['total_before'] tokens_after = m['tokens']['total_after'] tokens_saved = m['tokens']['total_saved'] - - # Calculate percentages + compressed_pct = (compressed / max(total, 1)) * 100 skipped_pct = (skipped / max(total, 1)) * 100 over_limit_pct = (over_limit / max(total, 1)) * 100 - + print("\n") print(f"╔{'═'*70}╗") print(f"║{'TRAJECTORY COMPRESSION REPORT':^70}║") print(f"╠{'═'*70}╣") - - # Trajectories section + print(f"║{'':2}📁 TRAJECTORIES{' '*54}║") print(f"║{'─'*70}║") print(f"║{'':4}Total Processed: {total:>10,}{' '*32}║") @@ -1344,33 +933,30 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" print(f"║{'':4}├─ Skipped (under limit):{skipped:>9,} ({skipped_pct:>5.1f}%){' '*18}║") print(f"║{'':4}├─ Still over limit: {over_limit:>10,} ({over_limit_pct:>5.1f}%){' '*18}║") print(f"║{'':4}└─ Failed: {failed:>10,}{' '*32}║") - + print(f"╠{'═'*70}╣") - - # Tokens section + print(f"║{'':2}🔢 TOKENS{' '*60}║") print(f"║{'─'*70}║") print(f"║{'':4}Before Compression: {tokens_before:>15,} tokens{' '*21}║") print(f"║{'':4}After Compression: {tokens_after:>15,} tokens{' '*21}║") print(f"║{'':4}Total Saved: {tokens_saved:>15,} tokens{' '*21}║") print(f"║{'':4}Overall Compression: {m['tokens']['overall_compression_ratio']:>14.1%}{' '*28}║") - + if tokens_before > 0: savings_pct = (tokens_saved / tokens_before) * 100 print(f"║{'':4}Space Savings: {savings_pct:>14.1f}%{' '*28}║") - + print(f"╠{'═'*70}╣") - - # Turns section + print(f"║{'':2}💬 CONVERSATION TURNS{' '*48}║") print(f"║{'─'*70}║") print(f"║{'':4}Before Compression: {m['turns']['total_before']:>15,} turns{' '*22}║") print(f"║{'':4}After Compression: {m['turns']['total_after']:>15,} turns{' '*22}║") print(f"║{'':4}Total Removed: {m['turns']['total_removed']:>15,} turns{' '*22}║") - + print(f"╠{'═'*70}╣") - - # Averages section (for compressed trajectories only) + print(f"║{'':2}📈 AVERAGES (Compressed Trajectories Only){' '*27}║") print(f"║{'─'*70}║") if compressed > 0: @@ -1379,46 +965,145 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" print(f"║{'':4}Avg Turns Removed: {m['averages']['avg_turns_removed_per_compressed']:>14.1f}{' '*28}║") else: print(f"║{'':4}No trajectories were compressed{' '*38}║") - + print(f"╠{'═'*70}╣") - - # Summarization API section + print(f"║{'':2}🤖 SUMMARIZATION API{' '*49}║") print(f"║{'─'*70}║") print(f"║{'':4}API Calls Made: {m['summarization']['total_api_calls']:>15,}{' '*27}║") print(f"║{'':4}Errors: {m['summarization']['total_errors']:>15,}{' '*27}║") print(f"║{'':4}Success Rate: {m['summarization']['success_rate']:>14.1%}{' '*28}║") - + print(f"╠{'═'*70}╣") - - # Processing time section + duration = m['processing']['duration_seconds'] - if duration > 60: - time_str = f"{duration/60:.1f} minutes" - else: - time_str = f"{duration:.1f} seconds" - + time_str = f"{duration/60:.1f} minutes" if duration > 60 else f"{duration:.1f} seconds" throughput = total / max(duration, 0.001) - + print(f"║{'':2}⏱️ PROCESSING TIME{' '*51}║") print(f"║{'─'*70}║") print(f"║{'':4}Duration: {time_str:>20}{' '*22}║") print(f"║{'':4}Throughput: {throughput:>15.1f} traj/sec{' '*18}║") print(f"║{'':4}Started: {m['processing']['start_time'][:19]:>20}{' '*22}║") print(f"║{'':4}Finished: {m['processing']['end_time'][:19]:>20}{' '*22}║") - + print(f"╚{'═'*70}╝") - - # Distribution summary if we have data + if self.aggregate_metrics.compression_ratios: ratios = self.aggregate_metrics.compression_ratios tokens_saved_list = self.aggregate_metrics.tokens_saved_list - + print("\n📊 Distribution Summary:") print(f" Compression ratios: min={min(ratios):.2%}, max={max(ratios):.2%}, median={sorted(ratios)[len(ratios)//2]:.2%}") print(f" Tokens saved: min={min(tokens_saved_list):,}, max={max(tokens_saved_list):,}, median={sorted(tokens_saved_list)[len(tokens_saved_list)//2]:,}") +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def _load_cli_config(config: str, target_max_tokens: Optional[int], tokenizer: Optional[str]) -> CompressionConfig: + """Load the YAML config (defaults if missing) and apply CLI overrides.""" + if Path(config).exists(): + print(f"📋 Loading config from {config}") + compression_config = CompressionConfig.from_yaml(config) + else: + print(f"⚠️ Config not found at {config}, using defaults") + compression_config = CompressionConfig() + if target_max_tokens: + compression_config.target_max_tokens = target_max_tokens + if tokenizer: + compression_config.tokenizer_name = tokenizer + return compression_config + + +def _print_dry_run(icon: str, target: Any, output_path: Path) -> None: + print("\n🔍 DRY RUN MODE - analyzing without writing") + print(f"{icon} Would process: {target}") + print(f"{icon} Would output to: {output_path}") + + +def _run_file_mode(input_path: Path, output: Optional[str], compression_config: CompressionConfig, + sample_percent: Optional[float], seed: int, dry_run: bool) -> None: + """Single-file input: (sample,) compress via a temp directory, merge into one output file.""" + print("📄 Input mode: Single JSONL file") + output_path = Path(output) if output else input_path.parent / (input_path.stem + compression_config.output_suffix + ".jsonl") + + entries = [entry for _, entry in _load_jsonl( + input_path, lambda n, e: print(f"⚠️ Skipping invalid JSON at line {n}: {e}"), start=1)] + total_entries = len(entries) + print(f" Loaded {total_entries:,} trajectories from {input_path.name}") + + if sample_percent is not None: + random.seed(seed) + entries = random.sample(entries, max(1, int(total_entries * sample_percent / 100))) + print(f" Sampled {len(entries):,} trajectories ({sample_percent}% of {total_entries:,})") + + if dry_run: + _print_dry_run("📄", f"{len(entries):,} trajectories", output_path) + return + + with tempfile.TemporaryDirectory() as temp_dir: + temp_input_dir = Path(temp_dir) / "input" + temp_output_dir = Path(temp_dir) / "output" + temp_input_dir.mkdir() + _write_jsonl(temp_input_dir / "trajectories.jsonl", entries) + + TrajectoryCompressor(compression_config).process_directory(temp_input_dir, temp_output_dir) + + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, 'w', encoding='utf-8') as out_f: + for jsonl_file in sorted(temp_output_dir.glob("*.jsonl")): + with open(jsonl_file, 'r', encoding='utf-8') as in_f: + for line in in_f: + out_f.write(line) + + metrics_file = temp_output_dir / compression_config.metrics_output_file + if metrics_file.exists(): + metrics_output = output_path.parent / (output_path.stem + "_metrics.json") + shutil.copy(metrics_file, metrics_output) + print(f"💾 Metrics saved to {metrics_output}") + + print("\n✅ Compression complete!") + print(f"📄 Output: {output_path}") + + +def _run_dir_mode(input_path: Path, output: Optional[str], compression_config: CompressionConfig, + sample_percent: Optional[float], seed: int, dry_run: bool) -> None: + """Directory input: compress in place, or per-file sample into a temp dir first.""" + print("📁 Input mode: Directory of JSONL files") + output_path = Path(output) if output else input_path.parent / (input_path.name + compression_config.output_suffix) + + if sample_percent is None: + if dry_run: + _print_dry_run("📁", input_path, output_path) + return + TrajectoryCompressor(compression_config).process_directory(input_path, output_path) + else: + print(f"\n⚠️ Sampling from directory: will sample {sample_percent}% from each file") + with tempfile.TemporaryDirectory() as temp_dir: + temp_input_dir = Path(temp_dir) / "input" + temp_input_dir.mkdir() + random.seed(seed) + total_original = 0 + total_sampled = 0 + for jsonl_file in sorted(input_path.glob("*.jsonl")): + entries = [entry for _, entry in _load_jsonl(jsonl_file)] + total_original += len(entries) + sample_size = max(1, int(len(entries) * sample_percent / 100)) + sampled_entries = random.sample(entries, min(sample_size, len(entries))) + total_sampled += len(sampled_entries) + _write_jsonl(temp_input_dir / jsonl_file.name, sampled_entries) + print(f" Sampled {total_sampled:,} from {total_original:,} total trajectories") + + if dry_run: + _print_dry_run("📁", temp_input_dir, output_path) + return + TrajectoryCompressor(compression_config).process_directory(temp_input_dir, output_path) + + print("\n✅ Compression complete!") + + def main( input: str, output: str = None, @@ -1445,195 +1130,25 @@ def main( sample_percent: Sample this percentage of trajectories (1-100) before compression seed: Random seed for sampling reproducibility (default: 42) dry_run: Analyze without compressing (just show what would happen) - - Examples: - # Compress a directory (original behavior) - python trajectory_compressor.py --input=data/my_run - - # Compress a single file - python trajectory_compressor.py --input=data/trajectories.jsonl - - # Compress 15% sample of a file - python trajectory_compressor.py --input=data/trajectories.jsonl --sample_percent=15 - - # Compress 10% sample with custom output - python trajectory_compressor.py --input=data/trajectories.jsonl --sample_percent=10 --output=data/sampled_compressed.jsonl """ - import random - import tempfile - import shutil - print("🗜️ Trajectory Compressor") print("=" * 60) - - # Load configuration - config_path = Path(config) - if config_path.exists(): - print(f"📋 Loading config from {config}") - compression_config = CompressionConfig.from_yaml(config) - else: - print(f"⚠️ Config not found at {config}, using defaults") - compression_config = CompressionConfig() - - # Apply CLI overrides - if target_max_tokens: - compression_config.target_max_tokens = target_max_tokens - if tokenizer: - compression_config.tokenizer_name = tokenizer - - # Validate sample_percent + + compression_config = _load_cli_config(config, target_max_tokens, tokenizer) + if sample_percent is not None: if sample_percent <= 0 or sample_percent > 100: print(f"❌ sample_percent must be between 1 and 100, got {sample_percent}") return print(f"🎲 Will sample {sample_percent}% of trajectories (seed={seed})") - - # Setup paths and determine input type + input_path = Path(input) if not input_path.exists(): print(f"❌ Input not found: {input}") return - - is_file_input = input_path.is_file() - - if is_file_input: - print("📄 Input mode: Single JSONL file") - - # For file input, default output is file with _compressed suffix - if output: - output_path = Path(output) - else: - output_path = input_path.parent / (input_path.stem + compression_config.output_suffix + ".jsonl") - - # Load entries from the single file - entries = [] - with open(input_path, 'r', encoding='utf-8') as f: - for line_num, line in enumerate(f, 1): - line = line.strip() - if line: - try: - entries.append(json.loads(line)) - except json.JSONDecodeError as e: - print(f"⚠️ Skipping invalid JSON at line {line_num}: {e}") - - total_entries = len(entries) - print(f" Loaded {total_entries:,} trajectories from {input_path.name}") - - # Sample if requested - if sample_percent is not None: - random.seed(seed) - sample_size = max(1, int(total_entries * sample_percent / 100)) - entries = random.sample(entries, sample_size) - print(f" Sampled {len(entries):,} trajectories ({sample_percent}% of {total_entries:,})") - - if dry_run: - print("\n🔍 DRY RUN MODE - analyzing without writing") - print(f"📄 Would process: {len(entries):,} trajectories") - print(f"📄 Would output to: {output_path}") - return - - # Create a temporary directory for processing - with tempfile.TemporaryDirectory() as temp_dir: - temp_input_dir = Path(temp_dir) / "input" - temp_output_dir = Path(temp_dir) / "output" - temp_input_dir.mkdir() - - # Write entries to temp file - temp_input_file = temp_input_dir / "trajectories.jsonl" - with open(temp_input_file, 'w', encoding='utf-8') as f: - for entry in entries: - f.write(json.dumps(entry, ensure_ascii=False) + '\n') - - # Initialize compressor and process - compressor = TrajectoryCompressor(compression_config) - compressor.process_directory(temp_input_dir, temp_output_dir) - - # Copy result to output path (merge all files in temp_output_dir) - output_path.parent.mkdir(parents=True, exist_ok=True) - with open(output_path, 'w', encoding='utf-8') as out_f: - for jsonl_file in sorted(temp_output_dir.glob("*.jsonl")): - with open(jsonl_file, 'r', encoding='utf-8') as in_f: - for line in in_f: - out_f.write(line) - - # Copy metrics file if it exists - metrics_file = temp_output_dir / compression_config.metrics_output_file - if metrics_file.exists(): - metrics_output = output_path.parent / (output_path.stem + "_metrics.json") - shutil.copy(metrics_file, metrics_output) - print(f"💾 Metrics saved to {metrics_output}") - - print("\n✅ Compression complete!") - print(f"📄 Output: {output_path}") - - else: - # Directory input - original behavior - print("📁 Input mode: Directory of JSONL files") - - if output: - output_path = Path(output) - else: - output_path = input_path.parent / (input_path.name + compression_config.output_suffix) - - # If sampling is requested for directory mode, we need to handle it differently - if sample_percent is not None: - print(f"\n⚠️ Sampling from directory: will sample {sample_percent}% from each file") - - # Create a temp directory with sampled files - with tempfile.TemporaryDirectory() as temp_dir: - temp_input_dir = Path(temp_dir) / "input" - temp_input_dir.mkdir() - - random.seed(seed) - total_original = 0 - total_sampled = 0 - - # Sample from each JSONL file - for jsonl_file in sorted(input_path.glob("*.jsonl")): - entries = [] - with open(jsonl_file, 'r', encoding='utf-8') as f: - for line in f: - line = line.strip() - if line: - try: - entries.append(json.loads(line)) - except json.JSONDecodeError: - pass - - total_original += len(entries) - sample_size = max(1, int(len(entries) * sample_percent / 100)) - sampled_entries = random.sample(entries, min(sample_size, len(entries))) - total_sampled += len(sampled_entries) - - # Write sampled entries - temp_file = temp_input_dir / jsonl_file.name - with open(temp_file, 'w', encoding='utf-8') as f: - for entry in sampled_entries: - f.write(json.dumps(entry, ensure_ascii=False) + '\n') - - print(f" Sampled {total_sampled:,} from {total_original:,} total trajectories") - - if dry_run: - print("\n🔍 DRY RUN MODE - analyzing without writing") - print(f"📁 Would process: {temp_input_dir}") - print(f"📁 Would output to: {output_path}") - return - - # Initialize compressor and process the sampled data - compressor = TrajectoryCompressor(compression_config) - compressor.process_directory(temp_input_dir, output_path) - else: - if dry_run: - print("\n🔍 DRY RUN MODE - analyzing without writing") - print(f"📁 Would process: {input_path}") - print(f"📁 Would output to: {output_path}") - return - - # Initialize compressor and process directly - compressor = TrajectoryCompressor(compression_config) - compressor.process_directory(input_path, output_path) - - print("\n✅ Compression complete!") + + run_mode = _run_file_mode if input_path.is_file() else _run_dir_mode + run_mode(input_path, output, compression_config, sample_percent, seed, dry_run) if __name__ == "__main__": From 593aa2d191aadb15956157213517a3d2cd77c521 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:14:12 -0700 Subject: [PATCH 2/4] refactor(mini_swe_runner): split run_task into _call_model/_run_tool_calls, extract tool-response folding and prompt constants, reuse trajectory_compressor temperature helper --- mini_swe_runner.py | 675 +++++++++++++++++---------------------------- 1 file changed, 246 insertions(+), 429 deletions(-) diff --git a/mini_swe_runner.py b/mini_swe_runner.py index 2994bfc6d2..7e9211ae6d 100644 --- a/mini_swe_runner.py +++ b/mini_swe_runner.py @@ -1,28 +1,14 @@ #!/usr/bin/env python3 -""" -SWE Runner with Hermes Trajectory Format +"""SWE Runner with Hermes Trajectory Format -A runner that uses Hermes-Agent's built-in execution environments -(local, docker, modal) and outputs trajectories in the Hermes-Agent format -compatible with batch_runner.py and trajectory_compressor.py. - -Features: -- Uses Hermes-Agent's Docker, Modal, or Local environments for command execution -- Outputs trajectories in Hermes format (from/value pairs with / XML) -- Compatible with the trajectory compression pipeline -- Supports batch processing from JSONL prompt files +Runs tool-calling agent tasks in Hermes-Agent's execution environments (local, +docker, modal) and writes trajectories in Hermes format (from/value pairs with +/ XML), compatible with batch_runner.py and +trajectory_compressor.py. Supports single tasks and JSONL batch mode. Usage: - # Run a single task with local environment python mini_swe_runner.py --task "Create a hello world Python script" --env local - - # Run with Docker python mini_swe_runner.py --task "List files in /tmp" --env docker --image python:3.11-slim - - # Run with Modal (cloud) - python mini_swe_runner.py --task "Install numpy and test it" --env modal --image python:3.11-slim - - # Batch mode from JSONL file python mini_swe_runner.py --prompts_file prompts.jsonl --output_file trajectories.jsonl --env docker """ @@ -44,26 +30,10 @@ def _effective_temperature_for_model( model: str, base_url: Optional[str] = None, ) -> Optional[float]: - """Return a fixed temperature for models with strict sampling contracts. + """Fixed temperature for models with strict sampling contracts, else None (omit the kwarg).""" + from trajectory_compressor import _effective_temperature_for_model as _shared + return _shared(model, None, base_url) - Returns ``None`` when the model manages temperature server-side (Kimi); - callers must omit the ``temperature`` kwarg entirely in that case. - """ - try: - from agent.auxiliary_client import _fixed_temperature_for_model, OMIT_TEMPERATURE - except Exception: - return None - result = _fixed_temperature_for_model(model, base_url) - if result is OMIT_TEMPERATURE: - return None # caller must omit temperature - return result - - - - -# ============================================================================ -# Terminal Tool Definition (matches Hermes-Agent format) -# ============================================================================ TERMINAL_TOOL_DEFINITION = { "type": "function", @@ -109,10 +79,35 @@ TERMINAL_TOOL_DEFINITION = { } } +SYSTEM_PROMPT = """You are an AI agent that can execute bash commands to complete tasks. + +When you need to run commands, use the 'terminal' tool with your bash command. + +**Important:** +- When you have completed the task successfully, run: echo "MINI_SWE_AGENT_FINAL_OUTPUT" followed by a summary +- Be concise and efficient in your approach +- Install any needed tools with apt-get or pip +- Avoid interactive commands (no vim, nano, less, etc.) + +Complete the user's task step by step.""" + +HERMES_SYSTEM_PREFIX = ( + "You are a function calling AI model. You are provided with function signatures within XML tags. " + "You may call one or more functions to assist with the user query. If available tools are not relevant in assisting " + "with user query, just respond in natural conversational language. Don't make assumptions about what values to plug " + "into functions. After calling & executing the functions, you will be provided with function results within " + " XML tags. Here are the available tools:\n" +) +HERMES_SYSTEM_SUFFIX = ( + "For each function call return a JSON object, with the following pydantic model json schema for each:\n" + "{'title': 'FunctionCall', 'type': 'object', 'properties': {'name': {'title': 'Name', 'type': 'string'}, " + "'arguments': {'title': 'Arguments', 'type': 'object'}}, 'required': ['name', 'arguments']}\n" + "Each function call should be enclosed within XML tags.\n" + "Example:\n\n{'name': ,'arguments': }\n" +) + +_OPENROUTER_URL = "https://openrouter.ai/api/v1" -# ============================================================================ -# Environment Factory -# ============================================================================ def create_environment( env_type: str = "local", @@ -121,45 +116,38 @@ def create_environment( timeout: int = 60, **kwargs ): - """ - Create an execution environment using Hermes-Agent's built-in backends. - - Args: - env_type: One of "local", "docker", "modal" - image: Docker/Modal image name (ignored for local) - cwd: Working directory - timeout: Default command timeout - **kwargs: Additional environment-specific options - - Returns: - Environment instance with execute() and cleanup() methods - """ + """Create a Hermes execution environment (``local`` ignores ``image``/``kwargs``).""" if env_type == "local": from tools.environments.local import LocalEnvironment return LocalEnvironment(cwd=cwd, timeout=timeout) - - elif env_type == "docker": + if env_type == "docker": from tools.environments.docker import DockerEnvironment return DockerEnvironment(image=image, cwd=cwd, timeout=timeout, **kwargs) - - elif env_type == "modal": + if env_type == "modal": from tools.environments.modal import ModalEnvironment return ModalEnvironment(image=image, cwd=cwd, timeout=timeout, **kwargs) - - else: - raise ValueError(f"Unknown environment type: {env_type}. Use 'local', 'docker', or 'modal'") + raise ValueError(f"Unknown environment type: {env_type}. Use 'local', 'docker', or 'modal'") -# ============================================================================ -# Mini-SWE Runner with Hermes Trajectory Format -# ============================================================================ +def _parse_json_args(raw: Any) -> Any: + """Decode tool-call arguments; invalid JSON becomes ``{}``.""" + if not isinstance(raw, str): + return raw + try: + return json.loads(raw) + except json.JSONDecodeError: + return {} + + +def _gpt_content(msg: Dict[str, Any], content: str) -> str: + """Prefix ``content`` with a ```` block when the message carries reasoning.""" + think = f"{msg['reasoning']}" if msg.get("reasoning") else "" + return think + content + class MiniSWERunner: - """ - Agent runner that uses Hermes-Agent's built-in execution environments - and outputs trajectories in Hermes-Agent format. - """ - + """Tool-calling agent loop over a Hermes execution environment, emitting Hermes trajectories.""" + def __init__( self, model: str = "anthropic/claude-sonnet-4.6", @@ -172,20 +160,6 @@ class MiniSWERunner: command_timeout: int = 60, verbose: bool = False, ): - """ - Initialize the Mini-SWE Runner. - - Args: - model: Model name for OpenAI-compatible API - base_url: API base URL (optional, uses env vars if not provided) - api_key: API key (optional, uses env vars if not provided) - env_type: Environment type - "local", "docker", or "modal" - image: Docker/Modal image (ignored for local) - cwd: Working directory for commands - max_iterations: Maximum tool-calling iterations - command_timeout: Default timeout for commands - verbose: Enable verbose logging - """ self.model = model self.max_iterations = max_iterations self.command_timeout = command_timeout @@ -193,49 +167,40 @@ class MiniSWERunner: self.env_type = env_type self.image = image self.cwd = cwd - + self.logger = logging.getLogger(__name__) - - # Initialize LLM client via centralized provider router. - # If explicit api_key/base_url are provided (e.g. from CLI args), - # construct directly. Otherwise use the router for OpenRouter. - if api_key or base_url: - from openai import OpenAI - client_kwargs = { - "base_url": base_url or "https://openrouter.ai/api/v1", - "api_key": api_key or os.getenv( - "OPENROUTER_API_KEY", - os.getenv("ANTHROPIC_API_KEY", - os.getenv("OPENAI_API_KEY", ""))), - } - self.client = OpenAI(**client_kwargs) - else: - from agent.auxiliary_client import resolve_provider_client - self.client, _ = resolve_provider_client("openrouter", model=model) - if self.client is None: - # Fallback: try auto-detection - self.client, _ = resolve_provider_client("auto", model=model) - if self.client is None: - from openai import OpenAI - self.client = OpenAI( - base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENROUTER_API_KEY", "")) - - # Environment will be created per-task - self.env = None - - # Tool definition + self.client = self._init_client(base_url, api_key) + self.env = None # created per-task self.tools = [TERMINAL_TOOL_DEFINITION] - + print("🤖 Mini-SWE Runner initialized") print(f" Model: {self.model}") print(f" Environment: {self.env_type}") if self.env_type != "local": print(f" Image: {self.image}") print(f" Max iterations: {self.max_iterations}") - + + def _init_client(self, base_url: Optional[str], api_key: Optional[str]): + """Explicit api_key/base_url -> direct OpenAI client; otherwise the provider router.""" + if api_key or base_url: + from openai import OpenAI + return OpenAI( + base_url=base_url or _OPENROUTER_URL, + api_key=api_key or os.getenv( + "OPENROUTER_API_KEY", + os.getenv("ANTHROPIC_API_KEY", + os.getenv("OPENAI_API_KEY", ""))), + ) + from agent.auxiliary_client import resolve_provider_client + client, _ = resolve_provider_client("openrouter", model=self.model) + if client is None: + client, _ = resolve_provider_client("auto", model=self.model) + if client is None: + from openai import OpenAI + client = OpenAI(base_url=_OPENROUTER_URL, api_key=os.getenv("OPENROUTER_API_KEY", "")) + return client + def _create_env(self): - """Create the execution environment.""" print(f"🔧 Creating {self.env_type} environment...") self.env = create_environment( env_type=self.env_type, @@ -244,30 +209,19 @@ class MiniSWERunner: timeout=self.command_timeout ) print("✅ Environment ready") - + def _cleanup_env(self): - """Cleanup the execution environment.""" if self.env is not None: if hasattr(self.env, 'cleanup'): self.env.cleanup() elif hasattr(self.env, 'stop'): self.env.stop() self.env = None - + def _execute_command(self, command: str, timeout: int = None) -> Dict[str, Any]: - """ - Execute a command in the environment. - - Args: - command: Bash command to execute - timeout: Optional timeout override - - Returns: - Dict with 'output' and 'returncode' - """ + """Run ``command`` in the environment; returns ``{output, exit_code, error}``.""" if self.env is None: self._create_env() - try: result = self.env.execute(command, timeout=timeout or self.command_timeout) return { @@ -276,291 +230,187 @@ class MiniSWERunner: "error": None } except Exception as e: - return { - "output": "", - "exit_code": -1, - "error": str(e) - } - + return {"output": "", "exit_code": -1, "error": str(e)} + def _format_tools_for_system_message(self) -> str: - """Format tool definitions for the system message.""" - formatted_tools = [] - for tool in self.tools: - func = tool["function"] - formatted_tools.append({ - "name": func["name"], - "description": func.get("description", ""), - "parameters": func.get("parameters", {}), - "required": None - }) + formatted_tools = [ + { + "name": tool["function"]["name"], + "description": tool["function"].get("description", ""), + "parameters": tool["function"].get("parameters", {}), + "required": None, + } + for tool in self.tools + ] return json.dumps(formatted_tools, ensure_ascii=False) - + + def _tool_response_turn(self, messages: List[Dict[str, Any]], i: int) -> tuple: + """Fold the tool messages following assistant turn ``i`` into one ``tool`` value. + + Returns ``(value_or_None, index_of_last_consumed_message)``. + """ + tool_calls = messages[i]["tool_calls"] + tool_responses = [] + j = i + 1 + while j < len(messages) and messages[j]["role"] == "tool": + tool_msg = messages[j] + tool_content = tool_msg["content"] + try: + if tool_content.strip().startswith(("{", "[")): + tool_content = json.loads(tool_content) + except (json.JSONDecodeError, AttributeError): + pass + k = len(tool_responses) + tool_responses.append( + "\n" + + json.dumps({ + "tool_call_id": tool_msg.get("tool_call_id", ""), + "name": tool_calls[k]["function"]["name"] if k < len(tool_calls) else "unknown", + "content": tool_content, + }, ensure_ascii=False) + + "\n" + ) + j += 1 + if not tool_responses: + return None, i + return "\n".join(tool_responses), j - 1 + def _convert_to_hermes_format( self, messages: List[Dict[str, Any]], user_query: str, - completed: bool ) -> List[Dict[str, Any]]: - """ - Convert internal message format to Hermes trajectory format. - - This produces the exact format used by batch_runner.py. - """ - trajectory = [] - - # System message with tool definitions + """Convert the OpenAI-style message list to the Hermes trajectory format used by batch_runner.py.""" system_msg = ( - "You are a function calling AI model. You are provided with function signatures within XML tags. " - "You may call one or more functions to assist with the user query. If available tools are not relevant in assisting " - "with user query, just respond in natural conversational language. Don't make assumptions about what values to plug " - "into functions. After calling & executing the functions, you will be provided with function results within " - " XML tags. Here are the available tools:\n" - f"\n{self._format_tools_for_system_message()}\n\n" - "For each function call return a JSON object, with the following pydantic model json schema for each:\n" - "{'title': 'FunctionCall', 'type': 'object', 'properties': {'name': {'title': 'Name', 'type': 'string'}, " - "'arguments': {'title': 'Arguments', 'type': 'object'}}, 'required': ['name', 'arguments']}\n" - "Each function call should be enclosed within XML tags.\n" - "Example:\n\n{'name': ,'arguments': }\n" + HERMES_SYSTEM_PREFIX + + f"\n{self._format_tools_for_system_message()}\n\n" + + HERMES_SYSTEM_SUFFIX ) - - trajectory.append({"from": "system", "value": system_msg}) - trajectory.append({"from": "human", "value": user_query}) - - # Process messages (skip first user message as we already added it) - i = 1 + trajectory = [{"from": "system", "value": system_msg}, {"from": "human", "value": user_query}] + + i = 1 # first user message already added while i < len(messages): msg = messages[i] - if msg["role"] == "assistant": - if "tool_calls" in msg and msg["tool_calls"]: - # Assistant message with tool calls - content = "" - - # Add reasoning if present - if msg.get("reasoning"): - content = f"{msg['reasoning']}" - - if msg.get("content"): - content += msg["content"] + "\n" - - # Add tool calls in XML format + if msg.get("tool_calls"): + content = (msg["content"] + "\n") if msg.get("content") else "" for tool_call in msg["tool_calls"]: - if not tool_call or not isinstance(tool_call, dict): continue - try: - arguments = json.loads(tool_call["function"]["arguments"]) \ - if isinstance(tool_call["function"]["arguments"], str) \ - else tool_call["function"]["arguments"] - except json.JSONDecodeError: - arguments = {} - + if not tool_call or not isinstance(tool_call, dict): + continue tool_call_json = { "name": tool_call["function"]["name"], - "arguments": arguments + "arguments": _parse_json_args(tool_call["function"]["arguments"]), } content += f"\n{json.dumps(tool_call_json, ensure_ascii=False)}\n\n" - - trajectory.append({"from": "gpt", "value": content.rstrip()}) - - # Collect subsequent tool responses - tool_responses = [] - j = i + 1 - while j < len(messages) and messages[j]["role"] == "tool": - tool_msg = messages[j] - tool_content = tool_msg["content"] - - # Try to parse as JSON - try: - if tool_content.strip().startswith(("{", "[")): - tool_content = json.loads(tool_content) - except (json.JSONDecodeError, AttributeError): - pass - - tool_response = "\n" - tool_response += json.dumps({ - "tool_call_id": tool_msg.get("tool_call_id", ""), - "name": msg["tool_calls"][len(tool_responses)]["function"]["name"] \ - if len(tool_responses) < len(msg["tool_calls"]) else "unknown", - "content": tool_content - }, ensure_ascii=False) - tool_response += "\n" - tool_responses.append(tool_response) - j += 1 - - if tool_responses: - trajectory.append({"from": "tool", "value": "\n".join(tool_responses)}) - i = j - 1 - + trajectory.append({"from": "gpt", "value": _gpt_content(msg, content).rstrip()}) + tool_value, i = self._tool_response_turn(messages, i) + if tool_value is not None: + trajectory.append({"from": "tool", "value": tool_value}) else: - # Regular assistant message (no tool calls) - content = "" - if msg.get("reasoning"): - content = f"{msg['reasoning']}" - content += msg.get("content") or "" - trajectory.append({"from": "gpt", "value": content}) - + trajectory.append({"from": "gpt", "value": _gpt_content(msg, msg.get("content") or "")}) elif msg["role"] == "user": trajectory.append({"from": "human", "value": msg["content"]}) - i += 1 - + return trajectory - + + def _call_model(self, messages: List[Dict[str, Any]]): + """One chat completion with the ephemeral system prompt; returns the message or None on API error.""" + api_kwargs = { + "model": self.model, + "messages": [{"role": "system", "content": SYSTEM_PROMPT}] + messages, + "tools": self.tools, + "timeout": 300.0, + } + fixed_temperature = _effective_temperature_for_model( + self.model, + str(getattr(self.client, "base_url", "") or ""), + ) + if fixed_temperature is not None: + api_kwargs["temperature"] = fixed_temperature + try: + return self.client.chat.completions.create(**api_kwargs).choices[0].message + except Exception as e: + self.logger.error("API call failed: %s", e) + return None + + def _run_tool_calls(self, assistant_message, messages: List[Dict[str, Any]]) -> bool: + """Record the assistant turn, execute each terminal call, append results; True if the completion signal fired.""" + print(f"🔧 Tool calls: {len(assistant_message.tool_calls)}") + messages.append({ + "role": "assistant", + "content": assistant_message.content, + "tool_calls": [ + { + "id": tc.id, + "type": tc.type, + "function": {"name": tc.function.name, "arguments": tc.function.arguments}, + } + for tc in assistant_message.tool_calls + ] + }) + + completed = False + for tc in assistant_message.tool_calls: + args = _parse_json_args(tc.function.arguments) + command = args.get("command", "echo 'No command provided'") + timeout = args.get("timeout", self.command_timeout) + print(f" 📞 terminal: {command[:60]}...") + + result = self._execute_command(command, timeout) + result_json = json.dumps({ + "content": { + "output": result["output"], + "exit_code": result["exit_code"], + "error": result["error"] + } + }, ensure_ascii=False) + if "MINI_SWE_AGENT_FINAL_OUTPUT" in result["output"]: + print(" ✅ Task completion signal detected!") + completed = True + messages.append(make_tool_result_message(tc.function.name, result_json, tc.id)) + print(f" ✅ exit_code={result['exit_code']}, output={len(result['output'])} chars") + return completed + def run_task(self, task: str) -> Dict[str, Any]: - """ - Run a single task and return the result with trajectory. - - Args: - task: The task/prompt to execute - - Returns: - Dict with trajectory, completion status, and metadata - """ + """Run one task; returns ``{conversations, completed, api_calls, metadata}``.""" print(f"\n{'='*60}") print(f"📝 Task: {task[:80]}{'...' if len(task) > 80 else ''}") print(f"{'='*60}") - - # Initialize environment + self._create_env() - - # Message history messages = [{"role": "user", "content": task}] - - # System prompt for the LLM (ephemeral - not saved to trajectory) - system_prompt = """You are an AI agent that can execute bash commands to complete tasks. - -When you need to run commands, use the 'terminal' tool with your bash command. - -**Important:** -- When you have completed the task successfully, run: echo "MINI_SWE_AGENT_FINAL_OUTPUT" followed by a summary -- Be concise and efficient in your approach -- Install any needed tools with apt-get or pip -- Avoid interactive commands (no vim, nano, less, etc.) - -Complete the user's task step by step.""" - api_call_count = 0 completed = False - final_response = None - + try: while api_call_count < self.max_iterations: api_call_count += 1 print(f"\n🔄 API call #{api_call_count}/{self.max_iterations}") - - # Prepare API messages - api_messages = [{"role": "system", "content": system_prompt}] + messages - - # Make API call - try: - api_kwargs = { - "model": self.model, - "messages": api_messages, - "tools": self.tools, - "timeout": 300.0, - } - fixed_temperature = _effective_temperature_for_model( - self.model, - str(getattr(self.client, "base_url", "") or ""), - ) - if fixed_temperature is not None: - api_kwargs["temperature"] = fixed_temperature - response = self.client.chat.completions.create(**api_kwargs) - except Exception as e: - self.logger.error("API call failed: %s", e) + assistant_message = self._call_model(messages) + if assistant_message is None: break - - assistant_message = response.choices[0].message - - # Log assistant response if assistant_message.content: print(f"🤖 Assistant: {assistant_message.content[:100]}...") - - # Check for tool calls + if assistant_message.tool_calls: - print(f"🔧 Tool calls: {len(assistant_message.tool_calls)}") - - # Add assistant message with tool calls - messages.append({ - "role": "assistant", - "content": assistant_message.content, - "tool_calls": [ - { - "id": tc.id, - "type": tc.type, - "function": { - "name": tc.function.name, - "arguments": tc.function.arguments - } - } - for tc in assistant_message.tool_calls - ] - }) - - # Execute each tool call - for tc in assistant_message.tool_calls: - try: - args = json.loads(tc.function.arguments) - except json.JSONDecodeError: - args = {} - - command = args.get("command", "echo 'No command provided'") - timeout = args.get("timeout", self.command_timeout) - - print(f" 📞 terminal: {command[:60]}...") - - # Execute command - result = self._execute_command(command, timeout) - - # Format result - result_json = json.dumps({ - "content": { - "output": result["output"], - "exit_code": result["exit_code"], - "error": result["error"] - } - }, ensure_ascii=False) - - # Check for task completion signal - if "MINI_SWE_AGENT_FINAL_OUTPUT" in result["output"]: - print(" ✅ Task completion signal detected!") - completed = True - - # Add tool response - messages.append(make_tool_result_message( - tc.function.name, result_json, tc.id, - )) - - print(f" ✅ exit_code={result['exit_code']}, output={len(result['output'])} chars") - - # If task completed, we can stop - if completed: - final_response = assistant_message.content + if self._run_tool_calls(assistant_message, messages): + completed = True break - else: - # No tool calls - final response - final_response = assistant_message.content or "" - messages.append({ - "role": "assistant", - "content": final_response - }) + messages.append({"role": "assistant", "content": assistant_message.content or ""}) completed = True print("🎉 Agent finished (no more tool calls)") break - + if api_call_count >= self.max_iterations: print(f"⚠️ Reached max iterations ({self.max_iterations})") - finally: - # Cleanup environment self._cleanup_env() - - # Convert to Hermes trajectory format - trajectory = self._convert_to_hermes_format(messages, task, completed) - + return { - "conversations": trajectory, + "conversations": self._convert_to_hermes_format(messages, task), "completed": completed, "api_calls": api_call_count, "metadata": { @@ -569,63 +419,59 @@ Complete the user's task step by step.""" "timestamp": datetime.now().isoformat() } } - + def run_batch( self, prompts: List[str], output_file: str ) -> List[Dict[str, Any]]: - """ - Run multiple tasks and save trajectories to a JSONL file. - - Args: - prompts: List of task prompts - output_file: Output JSONL file path - - Returns: - List of results - """ + """Run every prompt, appending each result to ``output_file`` as it finishes.""" results = [] - + print(f"\n📦 Running batch of {len(prompts)} tasks") print(f"📁 Output: {output_file}") - + with open(output_file, 'w', encoding='utf-8') as f: for i, prompt in enumerate(prompts, 1): print(f"\n{'='*60}") print(f"📋 Task {i}/{len(prompts)}") print(f"{'='*60}") - + try: result = self.run_task(prompt) - results.append(result) - - # Write to file immediately - f.write(json.dumps(result, ensure_ascii=False) + "\n") - f.flush() - print(f"✅ Task {i} completed (api_calls={result['api_calls']})") - except Exception as e: self.logger.error("Error on task %s: %s", i, e) - error_result = { + result = { "conversations": [], "completed": False, "api_calls": 0, "error": str(e), "metadata": {"timestamp": datetime.now().isoformat()} } - results.append(error_result) - f.write(json.dumps(error_result, ensure_ascii=False) + "\n") - f.flush() - + results.append(result) + f.write(json.dumps(result, ensure_ascii=False) + "\n") + f.flush() + print(f"\n✅ Batch complete! {len(results)} trajectories saved to {output_file}") return results -# ============================================================================ -# CLI Interface -# ============================================================================ +def _load_prompts(prompts_file: str) -> List[str]: + """One prompt per non-blank line: JSON ``{"prompt"|"task": ...}`` or raw text.""" + prompts = [] + with open(prompts_file, 'r', encoding='utf-8') as f: + for line in f: + line = line.strip() + if not line: + continue + try: + entry = json.loads(line) + prompts.append(entry.get("prompt", entry.get("task", ""))) + except json.JSONDecodeError: + prompts.append(line) + return prompts + def main( task: str = None, @@ -657,28 +503,17 @@ def main( max_iterations: Maximum tool-calling iterations (default: 15) timeout: Command timeout in seconds (default: 60) verbose: Enable verbose logging - - Examples: - # Single task with local environment - python mini_swe_runner.py --task "Create hello.py that prints Hello World" - - # Single task with Docker - python mini_swe_runner.py --task "List files" --env docker - - # Batch from file - python mini_swe_runner.py --prompts_file tasks.jsonl --output_file results.jsonl """ print("🚀 Mini-SWE Runner with Hermes Trajectory Format") print("=" * 60) - + # Configure root logging at the entry point (not in library __init__). logging.basicConfig( level=logging.DEBUG if verbose else logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', datefmt='%H:%M:%S' ) - - # Initialize runner + runner = MiniSWERunner( model=model, base_url=base_url, @@ -690,39 +525,21 @@ def main( command_timeout=timeout, verbose=verbose, ) - + if task: - # Single task mode result = runner.run_task(task) - - # Save to file with open(output_file, 'w', encoding='utf-8') as f: f.write(json.dumps(result, ensure_ascii=False) + "\n") - print(f"\n📁 Trajectory saved to: {output_file}") print(f"✅ Completed: {result['completed']}") print(f"📞 API calls: {result['api_calls']}") print(f"💬 Turns: {len(result['conversations'])}") - elif prompts_file: - # Batch mode - prompts = [] - with open(prompts_file, 'r', encoding='utf-8') as f: - for line in f: - line = line.strip() - if line: - try: - entry = json.loads(line) - prompts.append(entry.get("prompt", entry.get("task", ""))) - except json.JSONDecodeError: - prompts.append(line) - + prompts = _load_prompts(prompts_file) if not prompts: print(f"❌ No prompts found in {prompts_file}") return - runner.run_batch(prompts, output_file) - else: print("❌ Please provide either --task or --prompts_file") print(" Example: python mini_swe_runner.py --task 'Create a hello world script'") From a4d85a4e19d3b841cef70ac9de403ec7369c2856 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:20:48 -0700 Subject: [PATCH 3/4] refactor(trajectory_compressor): compact metrics serialization, report sections, and helper signatures --- trajectory_compressor.py | 328 +++++++++++---------------------------- 1 file changed, 87 insertions(+), 241 deletions(-) diff --git a/trajectory_compressor.py b/trajectory_compressor.py index 163204b046..9dd7c65a51 100644 --- a/trajectory_compressor.py +++ b/trajectory_compressor.py @@ -22,7 +22,7 @@ import logging import asyncio from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Tuple -from dataclasses import dataclass, field +from dataclasses import asdict, dataclass, field from datetime import datetime from utils import base_url_host_matches, base_url_hostname @@ -47,27 +47,15 @@ def _response_finish_reason(response: Any) -> str: standalone CLI deliberately avoids importing the heavy context compressor. """ try: - if isinstance(response, dict): - choices = response.get("choices") or [{}] - first = choices[0] if choices else {} - reason = ( - first.get("finish_reason") - if isinstance(first, dict) - else getattr(first, "finish_reason", None) - ) - else: - choices = getattr(response, "choices", None) or [] - reason = getattr(choices[0], "finish_reason", None) if choices else None + choices = (response.get("choices") if isinstance(response, dict) else getattr(response, "choices", None)) or [] + first = choices[0] if choices else None + reason = first.get("finish_reason") if isinstance(first, dict) else getattr(first, "finish_reason", None) return str(reason).strip().lower() if reason else "" except Exception: return "" -def _effective_temperature_for_model( - model: str, - requested_temperature: Optional[float], - base_url: Optional[str] = None, -) -> Optional[float]: +def _effective_temperature_for_model(model: str, requested_temperature: Optional[float], base_url: Optional[str] = None) -> Optional[float]: """Apply fixed model temperature contracts to direct client calls. Returns ``None`` when the model manages temperature server-side (Kimi); @@ -82,16 +70,10 @@ def _effective_temperature_for_model( fixed_temperature = _fixed_temperature_for_model(model, base_url) if fixed_temperature is OMIT_TEMPERATURE: return None # caller must omit temperature - if fixed_temperature is not None: - return fixed_temperature - return requested_temperature + return requested_temperature if fixed_temperature is None else fixed_temperature -def _load_jsonl( - path: Path, - on_error: Optional[Callable[[int, json.JSONDecodeError], None]] = None, - start: int = 0, -) -> List[Tuple[int, Any]]: +def _load_jsonl(path: Path, on_error: Optional[Callable[[int, json.JSONDecodeError], None]] = None, start: int = 0) -> List[Tuple[int, Any]]: """Return ``(line_num, entry)`` for each non-blank line; bad lines go to ``on_error``.""" entries = [] with open(path, 'r', encoding='utf-8') as f: @@ -113,44 +95,16 @@ def _write_jsonl(path: Path, entries) -> None: f.write(json.dumps(entry, ensure_ascii=False) + '\n') -# YAML section -> ((yaml_key, CompressionConfig attribute), ...) -_YAML_SECTIONS: Dict[str, Tuple[Tuple[str, str], ...]] = { - "tokenizer": (("name", "tokenizer_name"), ("trust_remote_code", "trust_remote_code")), - "compression": ( - ("target_max_tokens", "target_max_tokens"), - ("summary_target_tokens", "summary_target_tokens"), - ), - "protected_turns": ( - ("first_system", "protect_first_system"), - ("first_human", "protect_first_human"), - ("first_gpt", "protect_first_gpt"), - ("first_tool", "protect_first_tool"), - ("last_n_turns", "protect_last_n_turns"), - ), - "summarization": ( - ("model", "summarization_model"), - ("base_url", "base_url"), - ("api_key_env", "api_key_env"), - ("temperature", "temperature"), - ("max_retries", "max_retries"), - ("retry_delay", "retry_delay"), - ), - "output": ( - ("add_summary_notice", "add_summary_notice"), - ("summary_notice_text", "summary_notice_text"), - ("output_suffix", "output_suffix"), - ), - "processing": ( - ("num_workers", "num_workers"), - ("max_concurrent_requests", "max_concurrent_requests"), - ("skip_under_target", "skip_under_target"), - ("save_over_limit", "save_over_limit"), - ), - "metrics": ( - ("enabled", "metrics_enabled"), - ("per_trajectory", "metrics_per_trajectory"), - ("output_file", "metrics_output_file"), - ), +# YAML section -> keys; "yaml_key:attr" when the config attribute name differs. +_YAML_SECTIONS: Dict[str, Tuple[str, ...]] = { + "tokenizer": ("name:tokenizer_name", "trust_remote_code"), + "compression": ("target_max_tokens", "summary_target_tokens"), + "protected_turns": ("first_system:protect_first_system", "first_human:protect_first_human", + "first_gpt:protect_first_gpt", "first_tool:protect_first_tool", "last_n_turns:protect_last_n_turns"), + "summarization": ("model:summarization_model", "base_url", "api_key_env", "temperature", "max_retries", "retry_delay"), + "output": ("add_summary_notice", "summary_notice_text", "output_suffix"), + "processing": ("num_workers", "max_concurrent_requests", "skip_under_target", "save_over_limit"), + "metrics": ("enabled:metrics_enabled", "per_trajectory:metrics_per_trajectory", "output_file:metrics_output_file"), } @@ -207,7 +161,9 @@ class CompressionConfig: for section, keys in _YAML_SECTIONS.items(): if section not in data: continue - for yaml_key, attr in keys: + for key in keys: + yaml_key, _, attr = key.partition(":") + attr = attr or yaml_key value = data[section].get(yaml_key, getattr(config, attr)) if attr == "base_url": value = value or config.base_url # ``base_url: null`` keeps the default @@ -239,25 +195,14 @@ class TrajectoryMetrics: summarization_errors: int = 0 def to_dict(self) -> Dict[str, Any]: - return { - "original_tokens": self.original_tokens, - "compressed_tokens": self.compressed_tokens, - "tokens_saved": self.tokens_saved, - "compression_ratio": round(self.compression_ratio, 4), - "original_turns": self.original_turns, - "compressed_turns": self.compressed_turns, - "turns_removed": self.turns_removed, - "compression_region": { - "start_idx": self.turns_compressed_start_idx, - "end_idx": self.turns_compressed_end_idx, - "turns_count": self.turns_in_compressed_region, - }, - "was_compressed": self.was_compressed, - "still_over_limit": self.still_over_limit, - "skipped_under_target": self.skipped_under_target, - "summarization_api_calls": self.summarization_api_calls, - "summarization_errors": self.summarization_errors, - } + d = asdict(self) + d["compression_ratio"] = round(self.compression_ratio, 4) + region = {"start_idx": d.pop("turns_compressed_start_idx"), "end_idx": d.pop("turns_compressed_end_idx"), + "turns_count": d.pop("turns_in_compressed_region")} + # Insert the region after turns_removed to keep the historical key order. + items = list(d.items()) + items.insert(7, ("compression_region", region)) + return dict(items) def _mean(values, default): @@ -304,18 +249,13 @@ class AggregateMetrics: self.total_turns_removed += metrics.turns_removed self.total_summarization_calls += metrics.summarization_api_calls self.total_summarization_errors += metrics.summarization_errors - if metrics.was_compressed: self.trajectories_compressed += 1 self.compression_ratios.append(metrics.compression_ratio) self.tokens_saved_list.append(metrics.tokens_saved) self.turns_removed_list.append(metrics.turns_removed) - - if metrics.skipped_under_target: - self.trajectories_skipped_under_target += 1 - - if metrics.still_over_limit: - self.trajectories_still_over_limit += 1 + self.trajectories_skipped_under_target += bool(metrics.skipped_under_target) + self.trajectories_still_over_limit += bool(metrics.still_over_limit) def to_dict(self) -> Dict[str, Any]: return { @@ -413,10 +353,7 @@ class TrajectoryCompressor: """Initialize HuggingFace tokenizer for token counting.""" try: from transformers import AutoTokenizer - self.tokenizer = AutoTokenizer.from_pretrained( - self.config.tokenizer_name, - trust_remote_code=self.config.trust_remote_code - ) + self.tokenizer = AutoTokenizer.from_pretrained(self.config.tokenizer_name, trust_remote_code=self.config.trust_remote_code) print(f"✅ Loaded tokenizer: {self.config.tokenizer_name}") except Exception as e: raise RuntimeError(f"Failed to load tokenizer '{self.config.tokenizer_name}': {e}") @@ -428,26 +365,19 @@ class TrajectoryCompressor: self._llm_provider = provider self._use_call_llm = True from agent.auxiliary_client import resolve_provider_client - client, _ = resolve_provider_client( - provider, model=self.config.summarization_model) + client, _ = resolve_provider_client(provider, model=self.config.summarization_model) if client is None: - raise RuntimeError( - f"Provider '{provider}' is not configured. " - f"Check your API key or run: hermes setup") - self.client = None # Not used directly - self.async_client = None # Not used directly + raise RuntimeError(f"Provider '{provider}' is not configured. Check your API key or run: hermes setup") + self.client = self.async_client = None # Not used directly else: # Custom endpoint — use config's raw base_url + api_key_env self._use_call_llm = False api_key = os.getenv(self.config.api_key_env) if not api_key: - raise RuntimeError( - f"Missing API key. Set {self.config.api_key_env} " - f"environment variable.") + raise RuntimeError(f"Missing API key. Set {self.config.api_key_env} environment variable.") from openai import OpenAI from agent.auxiliary_client import _to_openai_base_url - self.client = OpenAI( - api_key=api_key, base_url=_to_openai_base_url(self.config.base_url)) + self.client = OpenAI(api_key=api_key, base_url=_to_openai_base_url(self.config.base_url)) # AsyncOpenAI is created lazily in _get_async_client() so it binds to # the current event loop — each process_directory() call runs its own # asyncio.run(), and a shared client would hit "Event loop is closed". @@ -461,10 +391,7 @@ class TrajectoryCompressor: """Return a fresh AsyncOpenAI client bound to the running event loop.""" from openai import AsyncOpenAI from agent.auxiliary_client import _to_openai_base_url - self.async_client = AsyncOpenAI( - api_key=self._async_client_api_key, - base_url=_to_openai_base_url(self.config.base_url), - ) + self.async_client = AsyncOpenAI(api_key=self._async_client_api_key, base_url=_to_openai_base_url(self.config.base_url)) return self.async_client def _detect_provider(self) -> str: @@ -501,13 +428,8 @@ class TrajectoryCompressor: for i, turn in enumerate(trajectory): first_seen.setdefault(turn.get("from", ""), i) - for role, enabled in ( - ("system", self.config.protect_first_system), - ("human", self.config.protect_first_human), - ("gpt", self.config.protect_first_gpt), - ("tool", self.config.protect_first_tool), - ): - if enabled and role in first_seen: + for role in ("system", "human", "gpt", "tool"): + if getattr(self.config, f"protect_first_{role}") and role in first_seen: protected.add(first_seen[role]) protected.update(range(max(0, n - self.config.protect_last_n_turns), n)) @@ -532,13 +454,7 @@ class TrajectoryCompressor: return idx >= len(trajectory) or trajectory[idx].get("from") != "tool" @classmethod - def _snap_boundary( - cls, - trajectory: List[Dict[str, str]], - idx: int, - min_idx: int, - max_idx: int, - ) -> int: + def _snap_boundary(cls, trajectory: List[Dict[str, str]], idx: int, min_idx: int, max_idx: int) -> int: """Move a boundary onto the nearest clean turn boundary within ``[min_idx, max_idx]``. Forward is preferred (folds an orphaned ``tool`` turn into the region that @@ -601,16 +517,10 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" def _summary_request(self, prompt: str) -> Tuple[Optional[float], Dict[str, Any]]: """Return ``(temperature, create-kwargs)``; temperature None means omit it.""" - temperature = _effective_temperature_for_model( - self.config.summarization_model, - self.config.temperature, - self.config.base_url, - ) - kwargs = { - "model": self.config.summarization_model, - "messages": [{"role": "user", "content": prompt}], - "max_tokens": self.config.summary_target_tokens * 2, - } + cfg = self.config + temperature = _effective_temperature_for_model(cfg.summarization_model, cfg.temperature, cfg.base_url) + kwargs = {"model": cfg.summarization_model, "messages": [{"role": "user", "content": prompt}], + "max_tokens": cfg.summary_target_tokens * 2} if not getattr(self, '_use_call_llm', False) and temperature is not None: kwargs["temperature"] = temperature return temperature, kwargs @@ -620,10 +530,7 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" if _response_finish_reason(response) == "length": # Storing a truncated summary silently corrupts the trajectory's # memory, so raise and let the retry/backoff loop handle it. - raise RuntimeError( - "trajectory summarization hit the output token cap " - "(finish_reason=length); summary is incomplete" - ) + raise RuntimeError("trajectory summarization hit the output token cap (finish_reason=length); summary is incomplete") summary = self._coerce_summary_content(response.choices[0].message.content) return self._ensure_summary_prefix(summary) @@ -722,14 +629,8 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" metrics.turns_in_compressed_region = until - start return start, until - def _assemble_compressed( - self, - trajectory: List[Dict[str, str]], - start: int, - until: int, - summary: str, - metrics: TrajectoryMetrics, - ) -> List[Dict[str, str]]: + def _assemble_compressed(self, trajectory: List[Dict[str, str]], start: int, until: int, summary: str, + metrics: TrajectoryMetrics) -> List[Dict[str, str]]: """Head (with summary notice on system) + summary human turn + verbatim tail; finalize metrics.""" compressed = [] for turn in trajectory[:start]: @@ -749,32 +650,24 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" metrics.still_over_limit = metrics.compressed_tokens > self.config.target_max_tokens return compressed - def compress_trajectory( - self, - trajectory: List[Dict[str, str]] - ) -> Tuple[List[Dict[str, str]], TrajectoryMetrics]: + def compress_trajectory(self, trajectory: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], TrajectoryMetrics]: """Compress one trajectory into the target budget; returns ``(trajectory, metrics)``.""" metrics = TrajectoryMetrics() region = self._plan_compression(trajectory, metrics) if region is None: return trajectory, metrics start, until = region - summary = self._generate_summary( - self._extract_turn_content_for_summary(trajectory, start, until), metrics) + summary = self._generate_summary(self._extract_turn_content_for_summary(trajectory, start, until), metrics) return self._assemble_compressed(trajectory, start, until, summary, metrics), metrics - async def compress_trajectory_async( - self, - trajectory: List[Dict[str, str]] - ) -> Tuple[List[Dict[str, str]], TrajectoryMetrics]: + async def compress_trajectory_async(self, trajectory: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], TrajectoryMetrics]: """Async twin of ``compress_trajectory``.""" metrics = TrajectoryMetrics() region = self._plan_compression(trajectory, metrics) if region is None: return trajectory, metrics start, until = region - summary = await self._generate_summary_async( - self._extract_turn_content_for_summary(trajectory, start, until), metrics) + summary = await self._generate_summary_async(self._extract_turn_content_for_summary(trajectory, start, until), metrics) return self._assemble_compressed(trajectory, start, until, summary, metrics), metrics async def process_entry_async(self, entry: Dict[str, Any]) -> Tuple[Dict[str, Any], TrajectoryMetrics]: @@ -792,18 +685,13 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" """Compress every ``*.jsonl`` in ``input_dir`` into ``output_dir`` (async, parallel API calls).""" asyncio.run(self._process_directory_async(input_dir, output_dir)) - async def _process_one( - self, run: _RunProgress, file_path: Path, entry_idx: int, entry: Dict - ) -> Optional[Tuple[Dict[str, Any], TrajectoryMetrics]]: + async def _process_one(self, run: _RunProgress, file_path: Path, entry_idx: int, entry: Dict) -> Optional[Tuple[Dict[str, Any], TrajectoryMetrics]]: """Process one entry under the semaphore/timeout; None means dropped (timed out).""" async with run.semaphore: async with run.lock: run.in_flight += 1 try: - processed_entry, metrics = await asyncio.wait_for( - self.process_entry_async(entry), - timeout=self.config.per_trajectory_timeout - ) + processed_entry, metrics = await asyncio.wait_for(self.process_entry_async(entry), timeout=self.config.per_trajectory_timeout) async with run.lock: self.aggregate_metrics.add_trajectory_metrics(metrics) if metrics.was_compressed: @@ -859,28 +747,16 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" console.print(f"{'='*60}\n") with Progress( - SpinnerColumn(), - TextColumn("[progress.description]{task.description}"), - BarColumn(), - TaskProgressColumn(), - TextColumn("•"), - TimeElapsedColumn(), - TextColumn("•"), - TimeRemainingColumn(), - console=console, - refresh_per_second=10 # Higher refresh for async + SpinnerColumn(), TextColumn("[progress.description]{task.description}"), BarColumn(), TaskProgressColumn(), + TextColumn("•"), TimeElapsedColumn(), TextColumn("•"), TimeRemainingColumn(), + console=console, refresh_per_second=10, # Higher refresh for async ) as progress: run = _RunProgress( - progress=progress, - main_task=progress.add_task(f"[cyan]Compressing {total_entries:,} trajectories", total=total_entries), - status_task=progress.add_task("[dim]Starting...[/dim]", total=None), - lock=asyncio.Lock(), - semaphore=asyncio.Semaphore(self.config.max_concurrent_requests), + progress, progress.add_task(f"[cyan]Compressing {total_entries:,} trajectories", total=total_entries), + progress.add_task("[dim]Starting...[/dim]", total=None), + asyncio.Lock(), asyncio.Semaphore(self.config.max_concurrent_requests), ) - outcomes = await asyncio.gather(*( - self._process_one(run, file_path, entry_idx, entry) - for file_path, entry_idx, entry in all_entries - )) + outcomes = await asyncio.gather(*(self._process_one(run, *item) for item in all_entries)) progress.remove_task(run.status_task) # Write results preserving original order; timed-out entries are dropped. @@ -906,59 +782,44 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" def _print_summary(self): """Print comprehensive compression summary statistics.""" m = self.aggregate_metrics.to_dict() - total = m['summary']['total_trajectories'] compressed = m['summary']['trajectories_compressed'] skipped = m['summary']['trajectories_skipped_under_target'] over_limit = m['summary']['trajectories_still_over_limit'] - failed = m['summary']['trajectories_failed'] - tokens_before = m['tokens']['total_before'] - tokens_after = m['tokens']['total_after'] tokens_saved = m['tokens']['total_saved'] + pct = lambda n: (n / max(total, 1)) * 100 # noqa: E731 - compressed_pct = (compressed / max(total, 1)) * 100 - skipped_pct = (skipped / max(total, 1)) * 100 - over_limit_pct = (over_limit / max(total, 1)) * 100 + def section(title: str, pad: int) -> None: + print(f"╠{'═'*70}╣") + print(f"║{'':2}{title}{' '*pad}║") + print(f"║{'─'*70}║") print("\n") print(f"╔{'═'*70}╗") print(f"║{'TRAJECTORY COMPRESSION REPORT':^70}║") - print(f"╠{'═'*70}╣") - print(f"║{'':2}📁 TRAJECTORIES{' '*54}║") - print(f"║{'─'*70}║") + section("📁 TRAJECTORIES", 54) print(f"║{'':4}Total Processed: {total:>10,}{' '*32}║") - print(f"║{'':4}├─ Compressed: {compressed:>10,} ({compressed_pct:>5.1f}%){' '*18}║") - print(f"║{'':4}├─ Skipped (under limit):{skipped:>9,} ({skipped_pct:>5.1f}%){' '*18}║") - print(f"║{'':4}├─ Still over limit: {over_limit:>10,} ({over_limit_pct:>5.1f}%){' '*18}║") - print(f"║{'':4}└─ Failed: {failed:>10,}{' '*32}║") + print(f"║{'':4}├─ Compressed: {compressed:>10,} ({pct(compressed):>5.1f}%){' '*18}║") + print(f"║{'':4}├─ Skipped (under limit):{skipped:>9,} ({pct(skipped):>5.1f}%){' '*18}║") + print(f"║{'':4}├─ Still over limit: {over_limit:>10,} ({pct(over_limit):>5.1f}%){' '*18}║") + print(f"║{'':4}└─ Failed: {m['summary']['trajectories_failed']:>10,}{' '*32}║") - print(f"╠{'═'*70}╣") - - print(f"║{'':2}🔢 TOKENS{' '*60}║") - print(f"║{'─'*70}║") + section("🔢 TOKENS", 60) print(f"║{'':4}Before Compression: {tokens_before:>15,} tokens{' '*21}║") - print(f"║{'':4}After Compression: {tokens_after:>15,} tokens{' '*21}║") + print(f"║{'':4}After Compression: {m['tokens']['total_after']:>15,} tokens{' '*21}║") print(f"║{'':4}Total Saved: {tokens_saved:>15,} tokens{' '*21}║") print(f"║{'':4}Overall Compression: {m['tokens']['overall_compression_ratio']:>14.1%}{' '*28}║") - if tokens_before > 0: - savings_pct = (tokens_saved / tokens_before) * 100 - print(f"║{'':4}Space Savings: {savings_pct:>14.1f}%{' '*28}║") + print(f"║{'':4}Space Savings: {(tokens_saved / tokens_before) * 100:>14.1f}%{' '*28}║") - print(f"╠{'═'*70}╣") - - print(f"║{'':2}💬 CONVERSATION TURNS{' '*48}║") - print(f"║{'─'*70}║") + section("💬 CONVERSATION TURNS", 48) print(f"║{'':4}Before Compression: {m['turns']['total_before']:>15,} turns{' '*22}║") print(f"║{'':4}After Compression: {m['turns']['total_after']:>15,} turns{' '*22}║") print(f"║{'':4}Total Removed: {m['turns']['total_removed']:>15,} turns{' '*22}║") - print(f"╠{'═'*70}╣") - - print(f"║{'':2}📈 AVERAGES (Compressed Trajectories Only){' '*27}║") - print(f"║{'─'*70}║") + section("📈 AVERAGES (Compressed Trajectories Only)", 27) if compressed > 0: print(f"║{'':4}Avg Compression Ratio: {m['averages']['avg_compression_ratio']:>14.1%}{' '*28}║") print(f"║{'':4}Avg Tokens Saved: {m['averages']['avg_tokens_saved_per_compressed']:>14,.0f}{' '*28}║") @@ -966,36 +827,26 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" else: print(f"║{'':4}No trajectories were compressed{' '*38}║") - print(f"╠{'═'*70}╣") - - print(f"║{'':2}🤖 SUMMARIZATION API{' '*49}║") - print(f"║{'─'*70}║") + section("🤖 SUMMARIZATION API", 49) print(f"║{'':4}API Calls Made: {m['summarization']['total_api_calls']:>15,}{' '*27}║") print(f"║{'':4}Errors: {m['summarization']['total_errors']:>15,}{' '*27}║") print(f"║{'':4}Success Rate: {m['summarization']['success_rate']:>14.1%}{' '*28}║") - print(f"╠{'═'*70}╣") - duration = m['processing']['duration_seconds'] time_str = f"{duration/60:.1f} minutes" if duration > 60 else f"{duration:.1f} seconds" - throughput = total / max(duration, 0.001) - - print(f"║{'':2}⏱️ PROCESSING TIME{' '*51}║") - print(f"║{'─'*70}║") + section("⏱️ PROCESSING TIME", 51) print(f"║{'':4}Duration: {time_str:>20}{' '*22}║") - print(f"║{'':4}Throughput: {throughput:>15.1f} traj/sec{' '*18}║") + print(f"║{'':4}Throughput: {total / max(duration, 0.001):>15.1f} traj/sec{' '*18}║") print(f"║{'':4}Started: {m['processing']['start_time'][:19]:>20}{' '*22}║") print(f"║{'':4}Finished: {m['processing']['end_time'][:19]:>20}{' '*22}║") - print(f"╚{'═'*70}╝") - if self.aggregate_metrics.compression_ratios: - ratios = self.aggregate_metrics.compression_ratios - tokens_saved_list = self.aggregate_metrics.tokens_saved_list - + ratios = self.aggregate_metrics.compression_ratios + if ratios: + saved = self.aggregate_metrics.tokens_saved_list print("\n📊 Distribution Summary:") print(f" Compression ratios: min={min(ratios):.2%}, max={max(ratios):.2%}, median={sorted(ratios)[len(ratios)//2]:.2%}") - print(f" Tokens saved: min={min(tokens_saved_list):,}, max={max(tokens_saved_list):,}, median={sorted(tokens_saved_list)[len(tokens_saved_list)//2]:,}") + print(f" Tokens saved: min={min(saved):,}, max={max(saved):,}, median={sorted(saved)[len(saved)//2]:,}") # --------------------------------------------------------------------------- @@ -1023,8 +874,7 @@ def _print_dry_run(icon: str, target: Any, output_path: Path) -> None: print(f"{icon} Would output to: {output_path}") -def _run_file_mode(input_path: Path, output: Optional[str], compression_config: CompressionConfig, - sample_percent: Optional[float], seed: int, dry_run: bool) -> None: +def _run_file_mode(input_path: Path, output: Optional[str], compression_config: CompressionConfig, sample_percent: Optional[float], seed: int, dry_run: bool) -> None: """Single-file input: (sample,) compress via a temp directory, merge into one output file.""" print("📄 Input mode: Single JSONL file") output_path = Path(output) if output else input_path.parent / (input_path.stem + compression_config.output_suffix + ".jsonl") @@ -1055,8 +905,7 @@ def _run_file_mode(input_path: Path, output: Optional[str], compression_config: with open(output_path, 'w', encoding='utf-8') as out_f: for jsonl_file in sorted(temp_output_dir.glob("*.jsonl")): with open(jsonl_file, 'r', encoding='utf-8') as in_f: - for line in in_f: - out_f.write(line) + shutil.copyfileobj(in_f, out_f) metrics_file = temp_output_dir / compression_config.metrics_output_file if metrics_file.exists(): @@ -1068,8 +917,7 @@ def _run_file_mode(input_path: Path, output: Optional[str], compression_config: print(f"📄 Output: {output_path}") -def _run_dir_mode(input_path: Path, output: Optional[str], compression_config: CompressionConfig, - sample_percent: Optional[float], seed: int, dry_run: bool) -> None: +def _run_dir_mode(input_path: Path, output: Optional[str], compression_config: CompressionConfig, sample_percent: Optional[float], seed: int, dry_run: bool) -> None: """Directory input: compress in place, or per-file sample into a temp dir first.""" print("📁 Input mode: Directory of JSONL files") output_path = Path(output) if output else input_path.parent / (input_path.name + compression_config.output_suffix) @@ -1085,13 +933,11 @@ def _run_dir_mode(input_path: Path, output: Optional[str], compression_config: C temp_input_dir = Path(temp_dir) / "input" temp_input_dir.mkdir() random.seed(seed) - total_original = 0 - total_sampled = 0 + total_original = total_sampled = 0 for jsonl_file in sorted(input_path.glob("*.jsonl")): entries = [entry for _, entry in _load_jsonl(jsonl_file)] + sampled_entries = random.sample(entries, min(max(1, int(len(entries) * sample_percent / 100)), len(entries))) total_original += len(entries) - sample_size = max(1, int(len(entries) * sample_percent / 100)) - sampled_entries = random.sample(entries, min(sample_size, len(entries))) total_sampled += len(sampled_entries) _write_jsonl(temp_input_dir / jsonl_file.name, sampled_entries) print(f" Sampled {total_sampled:,} from {total_original:,} total trajectories") From 6d81b7e7e1aaa0cc8079475754400c520a48f5a1 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Wed, 2 Sep 2026 18:22:53 -0700 Subject: [PATCH 4/4] refactor(mini_swe_runner): unify temperature helper with trajectory_compressor, compact literals and signatures --- mini_swe_runner.py | 205 ++++++++++----------------------------------- 1 file changed, 45 insertions(+), 160 deletions(-) diff --git a/mini_swe_runner.py b/mini_swe_runner.py index 7e9211ae6d..fc7dc5099d 100644 --- a/mini_swe_runner.py +++ b/mini_swe_runner.py @@ -21,20 +21,12 @@ from typing import List, Dict, Any, Optional import fire from dotenv import load_dotenv from agent.tool_dispatch_helpers import make_tool_result_message +from trajectory_compressor import _effective_temperature_for_model # Load environment variables load_dotenv() -def _effective_temperature_for_model( - model: str, - base_url: Optional[str] = None, -) -> Optional[float]: - """Fixed temperature for models with strict sampling contracts, else None (omit the kwarg).""" - from trajectory_compressor import _effective_temperature_for_model as _shared - return _shared(model, None, base_url) - - TERMINAL_TOOL_DEFINITION = { "type": "function", "function": { @@ -65,18 +57,12 @@ TERMINAL_TOOL_DEFINITION = { "parameters": { "type": "object", "properties": { - "command": { - "type": "string", - "description": "The bash command to execute" - }, - "timeout": { - "type": "integer", - "description": "Command timeout in seconds (default: 60)" - } + "command": {"type": "string", "description": "The bash command to execute"}, + "timeout": {"type": "integer", "description": "Command timeout in seconds (default: 60)"}, }, - "required": ["command"] - } - } + "required": ["command"], + }, + }, } SYSTEM_PROMPT = """You are an AI agent that can execute bash commands to complete tasks. @@ -109,13 +95,7 @@ HERMES_SYSTEM_SUFFIX = ( _OPENROUTER_URL = "https://openrouter.ai/api/v1" -def create_environment( - env_type: str = "local", - image: str = "python:3.11-slim", - cwd: str = "/tmp", - timeout: int = 60, - **kwargs -): +def create_environment(env_type: str = "local", image: str = "python:3.11-slim", cwd: str = "/tmp", timeout: int = 60, **kwargs): """Create a Hermes execution environment (``local`` ignores ``image``/``kwargs``).""" if env_type == "local": from tools.environments.local import LocalEnvironment @@ -148,26 +128,11 @@ def _gpt_content(msg: Dict[str, Any], content: str) -> str: class MiniSWERunner: """Tool-calling agent loop over a Hermes execution environment, emitting Hermes trajectories.""" - def __init__( - self, - model: str = "anthropic/claude-sonnet-4.6", - base_url: str = None, - api_key: str = None, - env_type: str = "local", - image: str = "python:3.11-slim", - cwd: str = "/tmp", - max_iterations: int = 15, - command_timeout: int = 60, - verbose: bool = False, - ): - self.model = model - self.max_iterations = max_iterations - self.command_timeout = command_timeout - self.verbose = verbose - self.env_type = env_type - self.image = image - self.cwd = cwd - + def __init__(self, model: str = "anthropic/claude-sonnet-4.6", base_url: str = None, api_key: str = None, + env_type: str = "local", image: str = "python:3.11-slim", cwd: str = "/tmp", + max_iterations: int = 15, command_timeout: int = 60, verbose: bool = False): + self.model, self.max_iterations, self.command_timeout, self.verbose = model, max_iterations, command_timeout, verbose + self.env_type, self.image, self.cwd = env_type, image, cwd self.logger = logging.getLogger(__name__) self.client = self._init_client(base_url, api_key) self.env = None # created per-task @@ -184,13 +149,8 @@ class MiniSWERunner: """Explicit api_key/base_url -> direct OpenAI client; otherwise the provider router.""" if api_key or base_url: from openai import OpenAI - return OpenAI( - base_url=base_url or _OPENROUTER_URL, - api_key=api_key or os.getenv( - "OPENROUTER_API_KEY", - os.getenv("ANTHROPIC_API_KEY", - os.getenv("OPENAI_API_KEY", ""))), - ) + return OpenAI(base_url=base_url or _OPENROUTER_URL, api_key=api_key or os.getenv( + "OPENROUTER_API_KEY", os.getenv("ANTHROPIC_API_KEY", os.getenv("OPENAI_API_KEY", "")))) from agent.auxiliary_client import resolve_provider_client client, _ = resolve_provider_client("openrouter", model=self.model) if client is None: @@ -202,12 +162,7 @@ class MiniSWERunner: def _create_env(self): print(f"🔧 Creating {self.env_type} environment...") - self.env = create_environment( - env_type=self.env_type, - image=self.image, - cwd=self.cwd, - timeout=self.command_timeout - ) + self.env = create_environment(env_type=self.env_type, image=self.image, cwd=self.cwd, timeout=self.command_timeout) print("✅ Environment ready") def _cleanup_env(self): @@ -224,25 +179,16 @@ class MiniSWERunner: self._create_env() try: result = self.env.execute(command, timeout=timeout or self.command_timeout) - return { - "output": result.get("output", ""), - "exit_code": result.get("returncode", 0), - "error": None - } + return {"output": result.get("output", ""), "exit_code": result.get("returncode", 0), "error": None} except Exception as e: return {"output": "", "exit_code": -1, "error": str(e)} def _format_tools_for_system_message(self) -> str: - formatted_tools = [ - { - "name": tool["function"]["name"], - "description": tool["function"].get("description", ""), - "parameters": tool["function"].get("parameters", {}), - "required": None, - } - for tool in self.tools - ] - return json.dumps(formatted_tools, ensure_ascii=False) + return json.dumps([ + {"name": t["function"]["name"], "description": t["function"].get("description", ""), + "parameters": t["function"].get("parameters", {}), "required": None} + for t in self.tools + ], ensure_ascii=False) def _tool_response_turn(self, messages: List[Dict[str, Any]], i: int) -> tuple: """Fold the tool messages following assistant turn ``i`` into one ``tool`` value. @@ -261,31 +207,18 @@ class MiniSWERunner: except (json.JSONDecodeError, AttributeError): pass k = len(tool_responses) - tool_responses.append( - "\n" - + json.dumps({ - "tool_call_id": tool_msg.get("tool_call_id", ""), - "name": tool_calls[k]["function"]["name"] if k < len(tool_calls) else "unknown", - "content": tool_content, - }, ensure_ascii=False) - + "\n" - ) + body = json.dumps({"tool_call_id": tool_msg.get("tool_call_id", ""), + "name": tool_calls[k]["function"]["name"] if k < len(tool_calls) else "unknown", + "content": tool_content}, ensure_ascii=False) + tool_responses.append(f"\n{body}\n") j += 1 if not tool_responses: return None, i return "\n".join(tool_responses), j - 1 - def _convert_to_hermes_format( - self, - messages: List[Dict[str, Any]], - user_query: str, - ) -> List[Dict[str, Any]]: + def _convert_to_hermes_format(self, messages: List[Dict[str, Any]], user_query: str) -> List[Dict[str, Any]]: """Convert the OpenAI-style message list to the Hermes trajectory format used by batch_runner.py.""" - system_msg = ( - HERMES_SYSTEM_PREFIX - + f"\n{self._format_tools_for_system_message()}\n\n" - + HERMES_SYSTEM_SUFFIX - ) + system_msg = HERMES_SYSTEM_PREFIX + f"\n{self._format_tools_for_system_message()}\n\n" + HERMES_SYSTEM_SUFFIX trajectory = [{"from": "system", "value": system_msg}, {"from": "human", "value": user_query}] i = 1 # first user message already added @@ -297,10 +230,7 @@ class MiniSWERunner: for tool_call in msg["tool_calls"]: if not tool_call or not isinstance(tool_call, dict): continue - tool_call_json = { - "name": tool_call["function"]["name"], - "arguments": _parse_json_args(tool_call["function"]["arguments"]), - } + tool_call_json = {"name": tool_call["function"]["name"], "arguments": _parse_json_args(tool_call["function"]["arguments"])} content += f"\n{json.dumps(tool_call_json, ensure_ascii=False)}\n\n" trajectory.append({"from": "gpt", "value": _gpt_content(msg, content).rstrip()}) tool_value, i = self._tool_response_turn(messages, i) @@ -316,16 +246,10 @@ class MiniSWERunner: def _call_model(self, messages: List[Dict[str, Any]]): """One chat completion with the ephemeral system prompt; returns the message or None on API error.""" - api_kwargs = { - "model": self.model, - "messages": [{"role": "system", "content": SYSTEM_PROMPT}] + messages, - "tools": self.tools, - "timeout": 300.0, - } - fixed_temperature = _effective_temperature_for_model( - self.model, - str(getattr(self.client, "base_url", "") or ""), - ) + api_kwargs = {"model": self.model, "messages": [{"role": "system", "content": SYSTEM_PROMPT}] + messages, + "tools": self.tools, "timeout": 300.0} + # requested_temperature=None: only fixed model contracts (Kimi omit / Arcee 0.5) apply here. + fixed_temperature = _effective_temperature_for_model(self.model, None, str(getattr(self.client, "base_url", "") or "")) if fixed_temperature is not None: api_kwargs["temperature"] = fixed_temperature try: @@ -337,18 +261,10 @@ class MiniSWERunner: def _run_tool_calls(self, assistant_message, messages: List[Dict[str, Any]]) -> bool: """Record the assistant turn, execute each terminal call, append results; True if the completion signal fired.""" print(f"🔧 Tool calls: {len(assistant_message.tool_calls)}") - messages.append({ - "role": "assistant", - "content": assistant_message.content, - "tool_calls": [ - { - "id": tc.id, - "type": tc.type, - "function": {"name": tc.function.name, "arguments": tc.function.arguments}, - } - for tc in assistant_message.tool_calls - ] - }) + messages.append({"role": "assistant", "content": assistant_message.content, "tool_calls": [ + {"id": tc.id, "type": tc.type, "function": {"name": tc.function.name, "arguments": tc.function.arguments}} + for tc in assistant_message.tool_calls + ]}) completed = False for tc in assistant_message.tool_calls: @@ -358,13 +274,7 @@ class MiniSWERunner: print(f" 📞 terminal: {command[:60]}...") result = self._execute_command(command, timeout) - result_json = json.dumps({ - "content": { - "output": result["output"], - "exit_code": result["exit_code"], - "error": result["error"] - } - }, ensure_ascii=False) + result_json = json.dumps({"content": {"output": result["output"], "exit_code": result["exit_code"], "error": result["error"]}}, ensure_ascii=False) if "MINI_SWE_AGENT_FINAL_OUTPUT" in result["output"]: print(" ✅ Task completion signal detected!") completed = True @@ -413,18 +323,10 @@ class MiniSWERunner: "conversations": self._convert_to_hermes_format(messages, task), "completed": completed, "api_calls": api_call_count, - "metadata": { - "model": self.model, - "env_type": self.env_type, - "timestamp": datetime.now().isoformat() - } + "metadata": {"model": self.model, "env_type": self.env_type, "timestamp": datetime.now().isoformat()}, } - def run_batch( - self, - prompts: List[str], - output_file: str - ) -> List[Dict[str, Any]]: + def run_batch(self, prompts: List[str], output_file: str) -> List[Dict[str, Any]]: """Run every prompt, appending each result to ``output_file`` as it finishes.""" results = [] @@ -442,13 +344,8 @@ class MiniSWERunner: print(f"✅ Task {i} completed (api_calls={result['api_calls']})") except Exception as e: self.logger.error("Error on task %s: %s", i, e) - result = { - "conversations": [], - "completed": False, - "api_calls": 0, - "error": str(e), - "metadata": {"timestamp": datetime.now().isoformat()} - } + result = {"conversations": [], "completed": False, "api_calls": 0, "error": str(e), + "metadata": {"timestamp": datetime.now().isoformat()}} results.append(result) f.write(json.dumps(result, ensure_ascii=False) + "\n") f.flush() @@ -508,23 +405,11 @@ def main( print("=" * 60) # Configure root logging at the entry point (not in library __init__). - logging.basicConfig( - level=logging.DEBUG if verbose else logging.INFO, - format='%(asctime)s - %(levelname)s - %(message)s', - datefmt='%H:%M:%S' - ) + logging.basicConfig(level=logging.DEBUG if verbose else logging.INFO, + format='%(asctime)s - %(levelname)s - %(message)s', datefmt='%H:%M:%S') - runner = MiniSWERunner( - model=model, - base_url=base_url, - api_key=api_key, - env_type=env, - image=image, - cwd=cwd, - max_iterations=max_iterations, - command_timeout=timeout, - verbose=verbose, - ) + runner = MiniSWERunner(model=model, base_url=base_url, api_key=api_key, env_type=env, image=image, cwd=cwd, + max_iterations=max_iterations, command_timeout=timeout, verbose=verbose) if task: result = runner.run_task(task)