diff --git a/hermes_time.py b/hermes_time.py index ac833cec54..26bcea13de 100644 --- a/hermes_time.py +++ b/hermes_time.py @@ -61,16 +61,6 @@ def _resolve_timezone_name() -> str: return "" -def _get_zoneinfo(name: str) -> Optional[ZoneInfo]: - if not name: - return None - try: - return ZoneInfo(name) - except Exception as exc: - logger.warning("Invalid timezone '%s': %s. Falling back to server local time.", name, exc) - return None - - def get_timezone() -> Optional[ZoneInfo]: """Return the active profile's configured ZoneInfo, or None (server-local).""" cache_identity = _timezone_cache_identity() @@ -81,7 +71,12 @@ def get_timezone() -> Optional[ZoneInfo]: # Resolve outside the lock (config file I/O); first writer wins so concurrent resolvers of the # same identity converge on one ZoneInfo object. name = _resolve_timezone_name() - tz = _get_zoneinfo(name) + tz = None + if name: + try: + tz = ZoneInfo(name) + except Exception as exc: + logger.warning("Invalid timezone '%s': %s. Falling back to server local time.", name, exc) with _cache_lock: return _tz_cache.setdefault(cache_identity, (name, tz))[1] diff --git a/mini_swe_runner.py b/mini_swe_runner.py index fc7dc5099d..62fb8179a8 100644 --- a/mini_swe_runner.py +++ b/mini_swe_runner.py @@ -12,6 +12,7 @@ Usage: python mini_swe_runner.py --prompts_file prompts.jsonl --output_file trajectories.jsonl --env docker """ +import importlib import json import logging import os @@ -100,13 +101,10 @@ def create_environment(env_type: str = "local", image: str = "python:3.11-slim", if env_type == "local": from tools.environments.local import LocalEnvironment return LocalEnvironment(cwd=cwd, timeout=timeout) - if env_type == "docker": - from tools.environments.docker import DockerEnvironment - return DockerEnvironment(image=image, cwd=cwd, timeout=timeout, **kwargs) - if env_type == "modal": - from tools.environments.modal import ModalEnvironment - return ModalEnvironment(image=image, cwd=cwd, timeout=timeout, **kwargs) - raise ValueError(f"Unknown environment type: {env_type}. Use 'local', 'docker', or 'modal'") + if env_type not in ("docker", "modal"): + raise ValueError(f"Unknown environment type: {env_type}. Use 'local', 'docker', or 'modal'") + module = importlib.import_module(f"tools.environments.{env_type}") + return getattr(module, f"{env_type.capitalize()}Environment")(image=image, cwd=cwd, timeout=timeout, **kwargs) def _parse_json_args(raw: Any) -> Any: @@ -121,8 +119,7 @@ def _parse_json_args(raw: Any) -> Any: 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 + return (f"{msg['reasoning']}" if msg.get("reasoning") else "") + content class MiniSWERunner: @@ -137,7 +134,6 @@ class MiniSWERunner: 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}") @@ -167,10 +163,9 @@ class MiniSWERunner: def _cleanup_env(self): if self.env is not None: - if hasattr(self.env, 'cleanup'): - self.env.cleanup() - elif hasattr(self.env, 'stop'): - self.env.stop() + stop = getattr(self.env, 'cleanup', None) or getattr(self.env, 'stop', None) + if stop: + stop() self.env = None def _execute_command(self, command: str, timeout: int = None) -> Dict[str, Any]: @@ -212,36 +207,30 @@ class MiniSWERunner: "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 + return ("\n".join(tool_responses), j - 1) if tool_responses else (None, i) 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 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 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 + if msg["role"] == "user": + trajectory.append({"from": "human", "value": msg["content"]}) + elif msg["role"] == "assistant" and not msg.get("tool_calls"): + trajectory.append({"from": "gpt", "value": _gpt_content(msg, msg.get("content") or "")}) + elif msg["role"] == "assistant": + content = (msg["content"] + "\n") if msg.get("content") else "" + for tool_call in msg["tool_calls"]: + if isinstance(tool_call, dict) and tool_call: 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) - if tool_value is not None: - trajectory.append({"from": "tool", "value": tool_value}) - else: - trajectory.append({"from": "gpt", "value": _gpt_content(msg, msg.get("content") or "")}) - elif msg["role"] == "user": - trajectory.append({"from": "human", "value": msg["content"]}) + 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}) i += 1 - return trajectory def _call_model(self, messages: List[Dict[str, Any]]): @@ -256,7 +245,6 @@ class MiniSWERunner: 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.""" @@ -270,15 +258,12 @@ class MiniSWERunner: 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) + result = self._execute_command(command, args.get("timeout", self.command_timeout)) 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)) + messages.append(make_tool_result_message(tc.function.name, json.dumps({"content": result}, ensure_ascii=False), tc.id)) print(f" āœ… exit_code={result['exit_code']}, output={len(result['output'])} chars") return completed @@ -287,58 +272,44 @@ class MiniSWERunner: print(f"\n{'='*60}") print(f"šŸ“ Task: {task[:80]}{'...' if len(task) > 80 else ''}") print(f"{'='*60}") - self._create_env() messages = [{"role": "user", "content": task}] api_call_count = 0 completed = False - try: while api_call_count < self.max_iterations: api_call_count += 1 print(f"\nšŸ”„ API call #{api_call_count}/{self.max_iterations}") - assistant_message = self._call_model(messages) if assistant_message is None: break if assistant_message.content: print(f"šŸ¤– Assistant: {assistant_message.content[:100]}...") - - if assistant_message.tool_calls: - if self._run_tool_calls(assistant_message, messages): - completed = True - break - else: + if not assistant_message.tool_calls: messages.append({"role": "assistant", "content": assistant_message.content or ""}) completed = True print("šŸŽ‰ Agent finished (no more tool calls)") break - + if self._run_tool_calls(assistant_message, messages): + completed = True + break if api_call_count >= self.max_iterations: print(f"āš ļø Reached max iterations ({self.max_iterations})") finally: self._cleanup_env() - - return { - "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()}, - } + return {"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()}} 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 = [] - 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) print(f"āœ… Task {i} completed (api_calls={result['api_calls']})") @@ -349,7 +320,6 @@ class MiniSWERunner: 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 @@ -403,14 +373,11 @@ def main( """ 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') - 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) with open(output_file, 'w', encoding='utf-8') as f: diff --git a/registration_lifecycle.py b/registration_lifecycle.py index 21601e144c..6724cc2fe9 100644 --- a/registration_lifecycle.py +++ b/registration_lifecycle.py @@ -50,10 +50,8 @@ class ReplacementCoordinator: with self._lock: yield - def acquire( - self, slot: Hashable, *, current: Any, previous: Any, restore: Callable[[Any], bool], - finalize: Callable[[], None] | None = None, - ) -> ReplacementLease: + def acquire(self, slot: Hashable, *, current: Any, previous: Any, restore: Callable[[Any], bool], + finalize: Callable[[], None] | None = None) -> ReplacementLease: """Attach a new live generation to the matching active predecessor.""" with self._lock: leases = self._active.setdefault(slot, []) @@ -63,7 +61,10 @@ class ReplacementCoordinator: return lease def dispose(self, lease: ReplacementLease) -> None: - """Remove *lease*, restoring the nearest still-live predecessor.""" + """Remove *lease*, restoring the nearest still-live predecessor. + + ``restore`` -> ``finalize`` -> slot pruning each run even when an earlier step raises. + """ with self._lock: if not lease.active: return @@ -75,25 +76,20 @@ class ReplacementCoordinator: try: try: if latest is lease: - replacement = lease.previous - predecessor = lease.predecessor - while predecessor is not None: - if predecessor.active: - replacement = predecessor.current - break - replacement = predecessor.previous - predecessor = predecessor.predecessor - lease.restore(replacement) + # Restore the nearest still-live predecessor, else the last dead generation's previous. + replacement, predecessor = lease.previous, lease.predecessor + while predecessor is not None and not predecessor.active: + replacement, predecessor = predecessor.previous, predecessor.predecessor + lease.restore(predecessor.current if predecessor is not None else replacement) finally: if lease.finalize is not None: lease.finalize() finally: - if leases: - live = [item for item in leases if item.active] - if live: - self._active[lease.slot] = live - else: - self._active.pop(lease.slot, None) + live = [item for item in leases if item.active] + if live: + self._active[lease.slot] = live + elif leases: + self._active.pop(lease.slot, None) replacement_coordinator = ReplacementCoordinator() diff --git a/toolset_distributions.py b/toolset_distributions.py index c20372524d..409d7ce818 100644 --- a/toolset_distributions.py +++ b/toolset_distributions.py @@ -15,29 +15,20 @@ def _dist(description: str, **toolsets: int) -> Dict[str, object]: DISTRIBUTIONS = { - "default": _dist("All available tools, all the time", - web=100, vision=100, image_gen=100, terminal=100, file=100, browser=100), - "image_gen": _dist("Heavy focus on image generation with vision and web support", - image_gen=90, vision=90, web=55, terminal=45), - "research": _dist("Web research with vision analysis and reasoning", - web=90, browser=70, vision=50, terminal=10), + "default": _dist("All available tools, all the time", web=100, vision=100, image_gen=100, terminal=100, file=100, browser=100), + "image_gen": _dist("Heavy focus on image generation with vision and web support", image_gen=90, vision=90, web=55, terminal=45), + "research": _dist("Web research with vision analysis and reasoning", web=90, browser=70, vision=50, terminal=10), "science": _dist("Scientific research with web, terminal, file, and browser capabilities", web=94, terminal=94, file=94, vision=65, browser=50, image_gen=15), - "development": _dist("Terminal, file tools, and reasoning with occasional web lookup", - terminal=80, file=80, web=30, vision=10), - "safe": _dist("All tools except terminal for safety", - web=80, browser=70, vision=60, image_gen=60), - "balanced": _dist("Equal probability of all toolsets", - web=50, vision=50, image_gen=50, terminal=50, file=50, browser=50), + "development": _dist("Terminal, file tools, and reasoning with occasional web lookup", terminal=80, file=80, web=30, vision=10), + "safe": _dist("All tools except terminal for safety", web=80, browser=70, vision=60, image_gen=60), + "balanced": _dist("Equal probability of all toolsets", web=50, vision=50, image_gen=50, terminal=50, file=50, browser=50), "minimal": _dist("Only web tools for basic research", web=100), "terminal_only": _dist("Terminal and file tools for code execution tasks", terminal=100, file=100), - "terminal_web": _dist("Terminal and file tools with web search for documentation lookup", - terminal=100, file=100, web=100), + "terminal_web": _dist("Terminal and file tools with web search for documentation lookup", terminal=100, file=100, web=100), "creative": _dist("Image generation and vision analysis focus", image_gen=90, vision=90, web=30), - "reasoning": _dist("Heavy research/reasoning distribution with minimal other tools", - web=90, file=60, terminal=20), - "browser_use": _dist("Full browser-based web interaction with search, vision, and page control", - browser=100, web=80, vision=70), + "reasoning": _dist("Heavy research/reasoning distribution with minimal other tools", web=90, file=60, terminal=20), + "browser_use": _dist("Full browser-based web interaction with search, vision, and page control", browser=100, web=80, vision=70), "browser_only": _dist("Only browser automation tools for pure web interaction tasks", browser=100), # browser-use-tasks.jsonl: the browser toolset includes web_search since Google blocks direct browser searches "browser_tasks": _dist( @@ -45,15 +36,11 @@ DISTRIBUTIONS = { browser=97, vision=12, terminal=15, ), # nous-terminal-tasks.jsonl - "terminal_tasks": _dist( - "Terminal-focused distribution with high terminal/file availability, occasional other tools", - terminal=97, file=97, web=97, browser=75, vision=50, image_gen=10, - ), + "terminal_tasks": _dist("Terminal-focused distribution with high terminal/file availability, occasional other tools", + terminal=97, file=97, web=97, browser=75, vision=50, image_gen=10), # mixed-browser-terminal-tasks.jsonl - "mixed_tasks": _dist( - "Mixed distribution with high browser, terminal, and file availability for complex tasks", - browser=92, terminal=92, file=92, web=35, vision=15, image_gen=15, - ), + "mixed_tasks": _dist("Mixed distribution with high browser, terminal, and file availability for complex tasks", + browser=92, terminal=92, file=92, web=35, vision=15, image_gen=15), } @@ -66,6 +53,10 @@ def list_distributions() -> Dict[str, Dict]: return DISTRIBUTIONS.copy() +def validate_distribution(distribution_name: str) -> bool: + return distribution_name in DISTRIBUTIONS + + def sample_toolsets_from_distribution(distribution_name: str) -> List[str]: """Sample toolset names, each included independently with its % probability. @@ -75,34 +66,25 @@ def sample_toolsets_from_distribution(distribution_name: str) -> List[str]: dist = get_distribution(distribution_name) if not dist: raise ValueError(f"Unknown distribution: {distribution_name}") - selected_toolsets = [] for toolset_name, probability in dist["toolsets"].items(): if not validate_toolset(toolset_name): print(f"āš ļø Warning: Toolset '{toolset_name}' in distribution '{distribution_name}' is not valid") - continue - if random.random() * 100 < probability: + elif random.random() * 100 < probability: selected_toolsets.append(toolset_name) - if not selected_toolsets and dist["toolsets"]: highest_prob_toolset = max(dist["toolsets"].items(), key=lambda x: x[1])[0] if validate_toolset(highest_prob_toolset): selected_toolsets.append(highest_prob_toolset) - return selected_toolsets -def validate_distribution(distribution_name: str) -> bool: - return distribution_name in DISTRIBUTIONS - - def print_distribution_info(distribution_name: str) -> None: """Print a distribution's description and toolset probabilities (highest first).""" dist = get_distribution(distribution_name) if not dist: print(f"āŒ Unknown distribution: {distribution_name}") return - print(f"\nšŸ“Š Distribution: {distribution_name}") print(f" Description: {dist['description']}") print(" Toolsets:") 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) diff --git a/utils.py b/utils.py index 9de067d059..b9a9c534b2 100644 --- a/utils.py +++ b/utils.py @@ -77,27 +77,21 @@ def _restore_file_mode(path: Path, mode: "int | None") -> None: _IS_WINDOWS = os.name == "nt" - # Windows rename failures possibly caused by another handle on the target. CPython opens files -# without FILE_SHARE_DELETE, so ``os.replace`` onto an open file is denied with: -# 5 ERROR_ACCESS_DENIED ā€” what a held *target* handle actually reports (measured: a plain -# reader, in- or cross-process, yields 5, NOT 32) -# 32 ERROR_SHARING_VIOLATION ā€” the *source* temp file is held -# 33 ERROR_LOCK_VIOLATION ā€” byte-range lock on the target -# These are ambiguous (a real ACL denial is also 5), so recovery is bounded and a still-failing -# write is re-raised unchanged rather than classified up front. +# without FILE_SHARE_DELETE, so ``os.replace`` onto an open file is denied with 5 ERROR_ACCESS_DENIED +# (what a held *target* handle actually reports ā€” measured: a plain reader yields 5, NOT 32), +# 32 ERROR_SHARING_VIOLATION (the *source* temp file is held) or 33 ERROR_LOCK_VIOLATION (byte-range +# lock on the target). Ambiguous (a real ACL denial is also 5), so recovery is bounded and a +# still-failing write is re-raised unchanged rather than classified up front. _WINDOWS_CONTENDED_REPLACE_ERRORS = frozenset({5, 32, 33}) - # Retry budget for the atomic rename. A rename that wins here keeps the write fully atomic, so the -# budget covers a realistic hold (desktop auth-init holds auth.json >100 ms; a status read is -# ~0.05 ms): ~200 ms recovered atomically, ~310 ms worst case. The cap matters as much as the -# count ā€” gateway_state.json is rewritten every turn, so a permanently-held target pays the full -# budget per write (~0.3 s here vs ~1.3 s for 6 x 20..400 ms). Jittered so concurrent writers -# don't retry in lockstep. +# budget covers a realistic hold (desktop auth-init holds auth.json >100 ms): ~200 ms recovered +# atomically, ~310 ms worst case. The cap matters as much as the count ā€” gateway_state.json is +# rewritten every turn, so a permanently-held target pays the full budget per write. Jittered so +# concurrent writers don't retry in lockstep. _REPLACE_RETRY_ATTEMPTS = 4 _REPLACE_RETRY_BASE_DELAY_S = 0.02 _REPLACE_RETRY_MAX_DELAY_S = 0.1 - _CROSS_DEVICE_ERRNOS = (errno.EXDEV, errno.EBUSY) @@ -160,11 +154,8 @@ def atomic_replace(tmp_path: Union[str, Path], target: Union[str, Path]) -> str: if contended: # Lazy: keeps ``utils`` free of a package-level dependency on ``agent``. from agent.retry_utils import jittered_backoff - for attempt in range(1, _REPLACE_RETRY_ATTEMPTS + 1): - time.sleep(jittered_backoff( - attempt, base_delay=_REPLACE_RETRY_BASE_DELAY_S, max_delay=_REPLACE_RETRY_MAX_DELAY_S - )) + time.sleep(jittered_backoff(attempt, base_delay=_REPLACE_RETRY_BASE_DELAY_S, max_delay=_REPLACE_RETRY_MAX_DELAY_S)) try: os.replace(tmp_str, real_path) return real_path @@ -175,22 +166,15 @@ def atomic_replace(tmp_path: Union[str, Path], target: Union[str, Path]) -> str: break if not _is_contended_windows_replace_error(retry_exc): raise - logger.debug( - "atomic_replace: %s -> %s failed with %s; falling back to %s", tmp_str, real_path, - getattr(exc, "winerror", None) or errno.errorcode.get(exc.errno or 0, exc.errno), - "in-place rewrite" if contended else "copy", - ) - if contended: - # Re-raises the rewrite's own error, so an ACL denial is reported as such, not as contention. - _rewrite_in_place(tmp_str, real_path) - else: - _copy_fallback(tmp_str, real_path) + logger.debug("atomic_replace: %s -> %s failed with %s; falling back to %s", tmp_str, real_path, + getattr(exc, "winerror", None) or errno.errorcode.get(exc.errno or 0, exc.errno), + "in-place rewrite" if contended else "copy") + # The rewrite re-raises its own error, so an ACL denial is reported as such, not as contention. + (_rewrite_in_place if contended else _copy_fallback)(tmp_str, real_path) return real_path -def _atomic_write( - path: Path, write, *, prefix: str, encoding: str = "utf-8", mode: "int | None" = None, preserve_owner: bool = True -) -> None: +def _atomic_write(path: Path, write, *, prefix: str, encoding: str = "utf-8", mode: "int | None" = None, preserve_owner: bool = True) -> None: """Temp file + fsync + :func:`atomic_replace`, then re-apply owner/mode. *write(f)* emits the payload into the open text handle. *mode* is fchmod'd onto the temp fd @@ -208,8 +192,7 @@ def _atomic_write( write(f) f.flush() os.fsync(f.fileno()) - # Preserve symlinks ā€” swap in place on the real file. - _restore_file_metadata(Path(atomic_replace(tmp_path, path)), original_owner, mode) + _restore_file_metadata(Path(atomic_replace(tmp_path, path)), original_owner, mode) # symlink-preserving except BaseException: with suppress(OSError): os.unlink(tmp_path) @@ -222,36 +205,26 @@ def _mode_for_write(path: Path, create_mode: "int | None", preserve: bool = True return mode if mode is not None or path.exists() else create_mode -def atomic_write_text( - path: Union[str, Path], content: str, *, encoding: str = "utf-8", tmp_prefix: str = ".tmp_", - preserve_mode: bool = False, create_mode: "int | None" = None, -) -> None: +def atomic_write_text(path: Union[str, Path], content: str, *, encoding: str = "utf-8", tmp_prefix: str = ".tmp_", + preserve_mode: bool = False, create_mode: "int | None" = None) -> None: """Write *content* to *path* via temp file + fsync + atomic rename. The target is never left partially written on crash/interrupt. Shared by every destructive file rewrite (memory store, skill manager, agent importer, ...). """ path = Path(path) - _atomic_write( - path, lambda f: f.write(content), prefix=tmp_prefix, encoding=encoding, - mode=_mode_for_write(path, create_mode, preserve=preserve_mode), preserve_owner=preserve_mode, - ) + _atomic_write(path, lambda f: f.write(content), prefix=tmp_prefix, encoding=encoding, + mode=_mode_for_write(path, create_mode, preserve=preserve_mode), preserve_owner=preserve_mode) -def atomic_json_write( - path: Union[str, Path], data: Any, *, indent: int = 2, mode: int | None = None, **dump_kwargs: Any -) -> None: +def atomic_json_write(path: Union[str, Path], data: Any, *, indent: int = 2, mode: int | None = None, **dump_kwargs: Any) -> None: """Write JSON to *path* atomically (temp file + fsync + replace).""" path = Path(path) - _atomic_write( - path, lambda f: json.dump(data, f, indent=indent, ensure_ascii=False, **dump_kwargs), - prefix=f".{path.stem}_", mode=mode if mode is not None else _preserve_file_mode(path), - ) + _atomic_write(path, lambda f: json.dump(data, f, indent=indent, ensure_ascii=False, **dump_kwargs), + prefix=f".{path.stem}_", mode=mode if mode is not None else _preserve_file_mode(path)) -def warn_if_credential_file_broadly_readable( - path: Union[str, Path], *, label: str = "", log: logging.Logger | None = None -) -> bool: +def warn_if_credential_file_broadly_readable(path: Union[str, Path], *, label: str = "", log: logging.Logger | None = None) -> bool: """Warn when a credential file is group/world-readable; True when a warning was emitted. Hand-made secret files (or ones older Hermes wrote without an explicit mode) commonly end up @@ -266,10 +239,8 @@ def warn_if_credential_file_broadly_readable( return False if os.name != "posix" or not (file_mode & (stat.S_IRGRP | stat.S_IROTH)): return False - (log or logger).warning( - "%s%s is group/world-readable (mode 0%o) and contains secrets. Run: chmod 600 %s", - f"{label} " if label else "", p.name, stat.S_IMODE(file_mode), p, - ) + (log or logger).warning("%s%s is group/world-readable (mode 0%o) and contains secrets. Run: chmod 600 %s", + f"{label} " if label else "", p.name, stat.S_IMODE(file_mode), p) return True @@ -285,10 +256,8 @@ class IndentDumper(yaml.SafeDumper): return super().increase_indent(flow, False) -def atomic_yaml_write( - path: Union[str, Path], data: Any, *, default_flow_style: bool = False, sort_keys: bool = False, - extra_content: str | None = None, create_mode: "int | None" = None, -) -> None: +def atomic_yaml_write(path: Union[str, Path], data: Any, *, default_flow_style: bool = False, sort_keys: bool = False, + extra_content: str | None = None, create_mode: "int | None" = None) -> None: """Write YAML to *path* atomically (temp file + fsync + replace).""" path = Path(path) @@ -296,39 +265,30 @@ def atomic_yaml_write( # allow_unicode=True writes emoji/kaomoji as real UTF-8. Without it PyYAML emits astral # chars as `\UXXXXXXXX` escapes inside `\`-continued double-quoted strings ā€” a structure # stricter parsers and hand-edits routinely break into unclosed quotes, corrupting the config. - yaml.dump( - data, f, Dumper=IndentDumper, default_flow_style=default_flow_style, sort_keys=sort_keys, allow_unicode=True - ) + yaml.dump(data, f, Dumper=IndentDumper, default_flow_style=default_flow_style, sort_keys=sort_keys, allow_unicode=True) if extra_content: f.write(extra_content) _atomic_write(path, _write, prefix=f".{path.stem}_", mode=_mode_for_write(path, create_mode)) -def _roundtrip_yaml(): - """ruamel round-trip ``YAML`` configured to keep quotes/Unicode with 2-space indents.""" +def _roundtrip_load(path: Path): + """``(yaml_rt, CommentedMap)``: a ruamel round-trip loader keeping quotes/Unicode with 2-space + indents, plus *path* loaded through it (empty map when missing/blank).""" from ruamel.yaml import YAML + from ruamel.yaml.comments import CommentedMap yaml_rt = YAML(typ="rt") yaml_rt.preserve_quotes = True yaml_rt.allow_unicode = True yaml_rt.default_flow_style = False yaml_rt.indent(mapping=2, sequence=4, offset=2) - return yaml_rt - - -def _load_commented_map(yaml_rt, path: Path): - """Load *path* with *yaml_rt* as a ``CommentedMap`` (empty when missing/blank).""" - from ruamel.yaml.comments import CommentedMap - data = yaml_rt.load(path.read_text(encoding="utf-8")) if path.exists() else None - return data if isinstance(data, CommentedMap) else CommentedMap(data or {}) + return yaml_rt, data if isinstance(data, CommentedMap) else CommentedMap(data or {}) def _roundtrip_dump(path: Path, yaml_rt, config) -> None: - _atomic_write( - path, lambda f: yaml_rt.dump(config, f), prefix=f".{path.stem}_", mode=_preserve_file_mode(path) - ) + _atomic_write(path, lambda f: yaml_rt.dump(config, f), prefix=f".{path.stem}_", mode=_preserve_file_mode(path)) def atomic_roundtrip_yaml_update(path: Union[str, Path], key_path: str, value: Any) -> None: @@ -338,7 +298,6 @@ def atomic_roundtrip_yaml_update(path: Union[str, Path], key_path: str, value: A single setting mutation must not disturb the rest. Still writes via temp file + atomic replace. """ from ruamel.yaml.comments import CommentedMap - # Honor escaped dots and prefer existing literal dotted keys (model IDs like ``glm-5.3``) over # blind splitting ā€” same navigation as ``hermes config set``'s ``_set_nested``; otherwise # /model + TUI persistence wrote ``glm-5: {'3': ...}`` phantom siblings. @@ -346,9 +305,7 @@ def atomic_roundtrip_yaml_update(path: Union[str, Path], key_path: str, value: A path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) - yaml_rt = _roundtrip_yaml() - config = _load_commented_map(yaml_rt, path) - + yaml_rt, config = _roundtrip_load(path) current = config keys = _split_key_path(key_path) i = 0 @@ -364,7 +321,6 @@ def atomic_roundtrip_yaml_update(path: Union[str, Path], key_path: str, value: A current[seg] = next_value current = next_value i += consumed - _roundtrip_dump(path, yaml_rt, config) @@ -384,14 +340,12 @@ def atomic_roundtrip_yaml_save(path: Union[str, Path], new_state: dict) -> None: """ from ruamel.yaml.comments import CommentedMap from ruamel.yaml.scalarstring import DoubleQuotedScalarString - from hermes_cli.config import require_readable_config_before_write path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) require_readable_config_before_write(path) - yaml_rt = _roundtrip_yaml() - existing = _load_commented_map(yaml_rt, path) + yaml_rt, existing = _roundtrip_load(path) def _merge(dst: CommentedMap, src: dict) -> None: for key, value in src.items(): @@ -435,10 +389,8 @@ def fast_safe_load(stream: Any) -> Any: def _env_number(key: str, default, cast): raw = os.getenv(key, "").strip() - if not raw: - return default try: - return cast(raw) + return cast(raw) if raw else default except (ValueError, TypeError): return default @@ -481,9 +433,7 @@ def normalize_proxy_env_vars() -> None: def _parse_base_url(base_url: str): """``urlparse`` that tolerates a bare ``host[:port][/path]`` (no scheme).""" raw = (base_url or "").strip() - if not raw: - return None - return urlparse(raw if "://" in raw else f"//{raw}") + return urlparse(raw if "://" in raw else f"//{raw}") if raw else None def _hostname_of(parsed) -> str: @@ -523,9 +473,7 @@ def base_url_origin(base_url: str) -> tuple[str, str, int]: port = parsed.port except ValueError: # out-of-range or non-numeric port ā€” not a usable origin return ("", "", 0) - if port is None: - port = {"https": 443, "http": 80}.get(scheme, 0) - return (scheme, hostname, port) + return (scheme, hostname, {"https": 443, "http": 80}.get(scheme, 0) if port is None else port) def base_url_host_matches(base_url: str, domain: str) -> bool: