diff --git a/trajectory_compressor.py b/trajectory_compressor.py index 9dd7c65a51..19892bb80f 100644 --- a/trajectory_compressor.py +++ b/trajectory_compressor.py @@ -31,13 +31,10 @@ from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskPr from rich.console import Console from hermes_constants import OPENROUTER_BASE_URL, get_hermes_home from agent.retry_utils import jittered_backoff - -# Load .env from HERMES_HOME first, then project root as a dev fallback. from hermes_cli.env_loader import load_hermes_dotenv -_hermes_home = get_hermes_home() -_project_env = Path(__file__).parent / ".env" -load_hermes_dotenv(hermes_home=_hermes_home, project_env=_project_env) +# Load .env from HERMES_HOME first, then project root as a dev fallback. +load_hermes_dotenv(hermes_home=get_hermes_home(), project_env=Path(__file__).parent / ".env") def _response_finish_reason(response: Any) -> str: @@ -66,7 +63,6 @@ def _effective_temperature_for_model(model: str, requested_temperature: Optional from agent.auxiliary_client import _fixed_temperature_for_model, OMIT_TEMPERATURE except Exception: return requested_temperature - fixed_temperature = _fixed_temperature_for_model(model, base_url) if fixed_temperature is OMIT_TEMPERATURE: return None # caller must omit temperature @@ -78,8 +74,7 @@ def _load_jsonl(path: Path, on_error: Optional[Callable[[int, json.JSONDecodeErr entries = [] with open(path, 'r', encoding='utf-8') as f: for line_num, line in enumerate(f, start): - line = line.strip() - if not line: + if not line.strip(): continue try: entries.append((line_num, json.loads(line))) @@ -110,43 +105,30 @@ _YAML_SECTIONS: Dict[str, Tuple[str, ...]] = { @dataclass class CompressionConfig: - """Configuration for trajectory compression.""" - # Tokenizer + """Configuration for trajectory compression (tokenizer / targets / protected turns / summarizer / output / processing / metrics).""" 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 api_key_env: str = "OPENROUTER_API_KEY" 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 + per_trajectory_timeout: int = 300 # seconds (default: 5 min) metrics_enabled: bool = True metrics_per_trajectory: bool = True metrics_output_file: str = "compression_metrics.json" @@ -156,12 +138,9 @@ class CompressionConfig: """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() for section, keys in _YAML_SECTIONS.items(): - if section not in data: - continue - for key in keys: + for key in keys if section in data else (): yaml_key, _, attr = key.partition(":") attr = attr or yaml_key value = data[section].get(yaml_key, getattr(config, attr)) @@ -178,19 +157,15 @@ 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 @@ -199,9 +174,8 @@ class TrajectoryMetrics: 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)) + items.insert(7, ("compression_region", region)) # after turns_removed: historical key order return dict(items) @@ -217,23 +191,17 @@ 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 @@ -259,53 +227,27 @@ class AggregateMetrics: def to_dict(self) -> Dict[str, Any]: return { - "summary": { - "total_trajectories": self.total_trajectories, - "trajectories_compressed": self.trajectories_compressed, - "trajectories_skipped_under_target": self.trajectories_skipped_under_target, - "trajectories_still_over_limit": self.trajectories_still_over_limit, - "trajectories_failed": self.trajectories_failed, - "compression_rate": round(self.trajectories_compressed / max(self.total_trajectories, 1), 4), - }, - "tokens": { - "total_before": self.total_tokens_before, - "total_after": self.total_tokens_after, - "total_saved": self.total_tokens_saved, - "overall_compression_ratio": round(self.total_tokens_after / max(self.total_tokens_before, 1), 4), - }, - "turns": { - "total_before": self.total_turns_before, - "total_after": self.total_turns_after, - "total_removed": self.total_turns_removed, - }, - "averages": { - "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, - "total_errors": self.total_summarization_errors, - "success_rate": round(1 - (self.total_summarization_errors / max(self.total_summarization_calls, 1)), 4), - }, - "processing": { - "start_time": self.processing_start_time, - "end_time": self.processing_end_time, - "duration_seconds": round(self.processing_duration_seconds, 2), - }, + "summary": {"total_trajectories": self.total_trajectories, "trajectories_compressed": self.trajectories_compressed, + "trajectories_skipped_under_target": self.trajectories_skipped_under_target, + "trajectories_still_over_limit": self.trajectories_still_over_limit, "trajectories_failed": self.trajectories_failed, + "compression_rate": round(self.trajectories_compressed / max(self.total_trajectories, 1), 4)}, + "tokens": {"total_before": self.total_tokens_before, "total_after": self.total_tokens_after, "total_saved": self.total_tokens_saved, + "overall_compression_ratio": round(self.total_tokens_after / max(self.total_tokens_before, 1), 4)}, + "turns": {"total_before": self.total_turns_before, "total_after": self.total_turns_after, "total_removed": self.total_turns_removed}, + "averages": {"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, "total_errors": self.total_summarization_errors, + "success_rate": round(1 - (self.total_summarization_errors / max(self.total_summarization_calls, 1)), 4)}, + "processing": {"start_time": self.processing_start_time, "end_time": self.processing_end_time, + "duration_seconds": round(self.processing_duration_seconds, 2)}, } # 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"), + ("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"), ) @@ -327,7 +269,9 @@ class _RunProgress: in_flight: int = 0 timeouts: int = 0 - def advance(self, update_status: bool = True) -> None: + def finish(self, update_status: bool = True) -> None: + """Retire one in-flight entry and advance the bar (caller holds ``lock``).""" + self.in_flight -= 1 self.progress.advance(self.main_task) if update_status: self.progress.update(self.status_task, description=_STATUS_FMT.format( @@ -361,9 +305,9 @@ class TrajectoryCompressor: def _init_summarizer(self): """Route summarization through call_llm for known providers, else a raw client.""" provider = self._detect_provider() + self._use_call_llm = bool(provider) if provider: 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) if client is None: @@ -371,19 +315,17 @@ class TrajectoryCompressor: 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} 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)) - # 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". + # AsyncOpenAI is created lazily in _get_async_client() so it binds to the current event + # loop — each process_directory() runs its own asyncio.run(); a shared client would hit + # "Event loop is closed". self.async_client = None self._async_client_api_key = api_key - print(f"✅ Initialized summarizer client: {self.config.summarization_model}") print(f" Max concurrent requests: {self.config.max_concurrent_requests}") @@ -399,10 +341,7 @@ class TrajectoryCompressor: url = self.config.base_url or "" if base_url_hostname(url) == "chatgpt.com" and "/backend-api/codex" in url.lower(): return "codex" - for host, provider in _PROVIDER_HOSTS: - if base_url_host_matches(url, host): - return provider - return "" + return next((provider for host, provider in _PROVIDER_HOSTS if base_url_host_matches(url, host)), "") def count_tokens(self, text: str) -> int: """Token count via the configured tokenizer; falls back to len//4.""" @@ -422,51 +361,36 @@ class TrajectoryCompressor: def _find_protected_indices(self, trajectory: List[Dict[str, str]]) -> Tuple[set, int, int]: """Return ``(protected_set, compressible_start, compressible_end)``.""" n = len(trajectory) - protected = set() - first_seen: Dict[str, int] = {} for i, turn in enumerate(trajectory): first_seen.setdefault(turn.get("from", ""), i) - - 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 = {first_seen[role] for role in ("system", "human", "gpt", "tool") + if getattr(self.config, f"protect_first_{role}") and role in first_seen} 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 - - return protected, compressible_start, compressible_end + return protected, max(head_protected) + 1 if head_protected else 0, min(tail_protected) if tail_protected else n @staticmethod - def _is_boundary_clean(trajectory: List[Dict[str, str]], idx: int) -> bool: - """True if a boundary at ``idx`` does not split a gpt /tool pair. + def _snap_boundary(trajectory: List[Dict[str, str]], idx: int, min_idx: int, max_idx: int) -> int: + """Move a boundary onto the nearest turn boundary within ``[min_idx, max_idx]`` that does not + split a gpt /tool pair. - 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" - - @classmethod - 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 - holds its ``gpt`` turn); backward only when nothing clean lies ahead. + 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. + Forward is preferred (folds an orphaned ``tool`` turn into the region that holds its ``gpt`` + turn); backward only when nothing clean lies ahead. """ + def clean(i: int) -> bool: + return i >= len(trajectory) or trajectory[i].get("from") != "tool" forward = idx - while forward < max_idx and not cls._is_boundary_clean(trajectory, forward): + while forward < max_idx and not clean(forward): forward += 1 - if cls._is_boundary_clean(trajectory, forward): + if clean(forward): return forward backward = idx - while backward > min_idx and not cls._is_boundary_clean(trajectory, backward): + while backward > min_idx and not clean(backward): backward -= 1 return backward @@ -475,28 +399,12 @@ class TrajectoryCompressor: parts = [] for i in range(start, end): turn = trajectory[i] - role = turn.get("from", "unknown") value = turn.get("value", "") if len(value) > 3000: value = value[:1500] + "\n...[truncated]...\n" + value[-500:] - parts.append(f"[Turn {i} - {role.upper()}]:\n{value}") + parts.append(f"[Turn {i} - {turn.get('from', 'unknown').upper()}]:\n{value}") return "\n\n".join(parts) - @staticmethod - def _coerce_summary_content(content: Any) -> str: - """Normalize summary-model output to a safe string.""" - if not isinstance(content, str): - content = str(content) if content else "" - return content.strip() - - @staticmethod - def _ensure_summary_prefix(summary: str) -> str: - """Normalize summary text to include the expected prefix exactly once.""" - text = (summary or "").strip() - if text.startswith("[CONTEXT SUMMARY]:"): - return text - return "[CONTEXT SUMMARY]:" if not text else f"[CONTEXT SUMMARY]: {text}" - 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. @@ -526,13 +434,16 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" return temperature, kwargs def _finish_summary(self, response: Any) -> str: - """Extract the summary text; a ``length`` stop is a failure (partial summary).""" + """Extract the summary text with the ``[CONTEXT SUMMARY]:`` prefix exactly once; a ``length`` stop is a failure.""" 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. + # 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) + content = response.choices[0].message.content + text = (content if isinstance(content, str) else str(content) if content else "").strip() + if text.startswith("[CONTEXT SUMMARY]:"): + return text + return "[CONTEXT SUMMARY]:" if not text else f"[CONTEXT SUMMARY]: {text}" 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.""" @@ -580,9 +491,7 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" return _SUMMARY_FALLBACK await asyncio.sleep(delay) - def _plan_compression( - self, trajectory: List[Dict[str, str]], metrics: TrajectoryMetrics - ) -> Optional[Tuple[int, int]]: + 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 @@ -594,18 +503,15 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" 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 @@ -617,15 +523,12 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" 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_compressed_start_idx, metrics.turns_compressed_end_idx = start, until metrics.turns_in_compressed_region = until - start return start, until @@ -640,7 +543,6 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" 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 @@ -656,9 +558,8 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" 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) - return self._assemble_compressed(trajectory, start, until, summary, metrics), metrics + summary = self._generate_summary(self._extract_turn_content_for_summary(trajectory, *region), metrics) + return self._assemble_compressed(trajectory, *region, summary, metrics), metrics async def compress_trajectory_async(self, trajectory: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], TrajectoryMetrics]: """Async twin of ``compress_trajectory``.""" @@ -666,17 +567,15 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" 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) - return self._assemble_compressed(trajectory, start, until, summary, metrics), metrics + summary = await self._generate_summary_async(self._extract_turn_content_for_summary(trajectory, *region), metrics) + return self._assemble_compressed(trajectory, *region, summary, metrics), metrics async def process_entry_async(self, entry: Dict[str, Any]) -> Tuple[Dict[str, Any], TrajectoryMetrics]: """Compress one JSONL entry's ``conversations``; attach metrics when compressed.""" if "conversations" not in entry: return entry, TrajectoryMetrics() compressed_trajectory, metrics = await self.compress_trajectory_async(entry["conversations"]) - result = entry.copy() - result["conversations"] = compressed_trajectory + result = dict(entry, conversations=compressed_trajectory) if self.config.metrics_per_trajectory and metrics.was_compressed: result["compression_metrics"] = metrics.to_dict() return result, metrics @@ -697,32 +596,27 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" 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() + run.skipped += bool(metrics.skipped_under_target) + run.finish() 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() + run.finish() 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) + run.finish(update_status=False) return entry, TrajectoryMetrics() # keep the original on error async def _process_directory_async(self, input_dir: Path, output_dir: Path): console = Console() self.aggregate_metrics.processing_start_time = datetime.now().isoformat() start_time = time.time() - jsonl_files = sorted(input_dir.glob("*.jsonl")) if not jsonl_files: self.logger.warning("No JSONL files found in %s", input_dir) @@ -772,7 +666,6 @@ Write only the summary, starting with "[CONTEXT SUMMARY]:" prefix.""" self.aggregate_metrics.processing_end_time = datetime.now().isoformat() self.aggregate_metrics.processing_duration_seconds = time.time() - start_time self._print_summary() - if self.config.metrics_enabled: metrics_path = output_dir / self.config.metrics_output_file with open(metrics_path, 'w', encoding="utf-8") as f: @@ -782,63 +675,57 @@ 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'] - tokens_before = m['tokens']['total_before'] - tokens_saved = m['tokens']['total_saved'] + s, t, u, a, z, p = m['summary'], m['tokens'], m['turns'], m['averages'], m['summarization'], m['processing'] + total, compressed = s['total_trajectories'], s['trajectories_compressed'] pct = lambda n: (n / max(total, 1)) * 100 # noqa: E731 + duration = p['duration_seconds'] + time_str = f"{duration/60:.1f} minutes" if duration > 60 else f"{duration:.1f} seconds" - def section(title: str, pad: int) -> None: - print(f"╠{'═'*70}╣") - print(f"║{'':2}{title}{' '*pad}║") - print(f"║{'─'*70}║") - + sections = [ + ("📁 TRAJECTORIES", 54, [ + f"║{'':4}Total Processed: {total:>10,}{' '*32}║", + f"║{'':4}├─ Compressed: {compressed:>10,} ({pct(compressed):>5.1f}%){' '*18}║", + f"║{'':4}├─ Skipped (under limit):{s['trajectories_skipped_under_target']:>9,} ({pct(s['trajectories_skipped_under_target']):>5.1f}%){' '*18}║", + f"║{'':4}├─ Still over limit: {s['trajectories_still_over_limit']:>10,} ({pct(s['trajectories_still_over_limit']):>5.1f}%){' '*18}║", + f"║{'':4}└─ Failed: {s['trajectories_failed']:>10,}{' '*32}║", + ]), + ("🔢 TOKENS", 60, [ + f"║{'':4}Before Compression: {t['total_before']:>15,} tokens{' '*21}║", + f"║{'':4}After Compression: {t['total_after']:>15,} tokens{' '*21}║", + f"║{'':4}Total Saved: {t['total_saved']:>15,} tokens{' '*21}║", + f"║{'':4}Overall Compression: {t['overall_compression_ratio']:>14.1%}{' '*28}║", + ] + ([f"║{'':4}Space Savings: {(t['total_saved'] / t['total_before']) * 100:>14.1f}%{' '*28}║"] if t['total_before'] > 0 else [])), + ("💬 CONVERSATION TURNS", 48, [ + f"║{'':4}Before Compression: {u['total_before']:>15,} turns{' '*22}║", + f"║{'':4}After Compression: {u['total_after']:>15,} turns{' '*22}║", + f"║{'':4}Total Removed: {u['total_removed']:>15,} turns{' '*22}║", + ]), + ("📈 AVERAGES (Compressed Trajectories Only)", 27, [ + f"║{'':4}Avg Compression Ratio: {a['avg_compression_ratio']:>14.1%}{' '*28}║", + f"║{'':4}Avg Tokens Saved: {a['avg_tokens_saved_per_compressed']:>14,.0f}{' '*28}║", + f"║{'':4}Avg Turns Removed: {a['avg_turns_removed_per_compressed']:>14.1f}{' '*28}║", + ] if compressed > 0 else [f"║{'':4}No trajectories were compressed{' '*38}║"]), + ("🤖 SUMMARIZATION API", 49, [ + f"║{'':4}API Calls Made: {z['total_api_calls']:>15,}{' '*27}║", + f"║{'':4}Errors: {z['total_errors']:>15,}{' '*27}║", + f"║{'':4}Success Rate: {z['success_rate']:>14.1%}{' '*28}║", + ]), + ("⏱️ PROCESSING TIME", 51, [ + f"║{'':4}Duration: {time_str:>20}{' '*22}║", + f"║{'':4}Throughput: {total / max(duration, 0.001):>15.1f} traj/sec{' '*18}║", + f"║{'':4}Started: {p['start_time'][:19]:>20}{' '*22}║", + f"║{'':4}Finished: {p['end_time'][:19]:>20}{' '*22}║", + ]), + ] print("\n") print(f"╔{'═'*70}╗") print(f"║{'TRAJECTORY COMPRESSION REPORT':^70}║") - - section("📁 TRAJECTORIES", 54) - print(f"║{'':4}Total Processed: {total:>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}║") - - section("🔢 TOKENS", 60) - print(f"║{'':4}Before Compression: {tokens_before:>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: - print(f"║{'':4}Space Savings: {(tokens_saved / tokens_before) * 100:>14.1f}%{' '*28}║") - - 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}║") - - 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}║") - 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}║") - - 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}║") - - duration = m['processing']['duration_seconds'] - time_str = f"{duration/60:.1f} minutes" if duration > 60 else f"{duration:.1f} seconds" - section("⏱️ PROCESSING TIME", 51) - print(f"║{'':4}Duration: {time_str:>20}{' '*22}║") - 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}║") + for title, pad, rows in sections: + print(f"╠{'═'*70}╣") + print(f"║{'':2}{title}{' '*pad}║") + print(f"║{'─'*70}║") + for row in rows: + print(row) print(f"╚{'═'*70}╝") ratios = self.aggregate_metrics.compression_ratios @@ -874,45 +761,40 @@ def _print_dry_run(icon: str, target: Any, output_path: Path) -> None: print(f"{icon} Would output to: {output_path}") +def _sample(entries: list, sample_percent: float) -> list: + return random.sample(entries, min(max(1, int(len(entries) * sample_percent / 100)), len(entries))) + + 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)] + 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, temp_output_dir = Path(temp_dir) / "input", 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: shutil.copyfileobj(in_f, out_f) - 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}") @@ -921,7 +803,6 @@ def _run_dir_mode(input_path: Path, output: Optional[str], compression_config: C """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) @@ -936,30 +817,20 @@ def _run_dir_mode(input_path: Path, output: Optional[str], compression_config: C 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))) + sampled_entries = _sample(entries, sample_percent) total_original += 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, - config: str = "configs/trajectory_compression.yaml", - target_max_tokens: int = None, - tokenizer: str = None, - sample_percent: float = None, - seed: int = 42, - dry_run: bool = False, -): +def main(input: str, output: str = None, config: str = "configs/trajectory_compression.yaml", target_max_tokens: int = None, + tokenizer: str = None, sample_percent: float = None, seed: int = 42, dry_run: bool = False): """ Compress agent trajectories to fit within a target token budget. @@ -979,20 +850,16 @@ def main( """ print("🗜️ Trajectory Compressor") print("=" * 60) - 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})") - input_path = Path(input) if not input_path.exists(): print(f"❌ Input not found: {input}") return - 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)