diff --git a/lib/cli/src/crewai_cli/benchmark.py b/lib/cli/src/crewai_cli/benchmark.py index fc3d9ebc8..0128d9a49 100644 --- a/lib/cli/src/crewai_cli/benchmark.py +++ b/lib/cli/src/crewai_cli/benchmark.py @@ -3,13 +3,14 @@ from __future__ import annotations import asyncio +from collections.abc import Callable import json +from pathlib import Path import re import time -from pathlib import Path -from typing import Any, Callable +from typing import Any -from pydantic import BaseModel, Field +from pydantic import BaseModel class BenchmarkCase(BaseModel): @@ -340,14 +341,16 @@ async def run_benchmark( return dict(zip(models, all_results)) -class suppress_benchmark_output: +class SuppressBenchmarkOutput: """Context manager that silences console formatter and noisy logging during benchmarks.""" def __enter__(self): import logging self._saved_formatter = None try: - from crewai.events.listeners.tracing.trace_listener import TraceCollectionListener + from crewai.events.listeners.tracing.trace_listener import ( + TraceCollectionListener, + ) listener = TraceCollectionListener._instance if listener: self._saved_formatter = listener.formatter @@ -366,7 +369,9 @@ class suppress_benchmark_output: lg.setLevel(level) if self._saved_formatter is not None: try: - from crewai.events.listeners.tracing.trace_listener import TraceCollectionListener + from crewai.events.listeners.tracing.trace_listener import ( + TraceCollectionListener, + ) listener = TraceCollectionListener._instance if listener: listener.formatter = self._saved_formatter @@ -374,7 +379,7 @@ class suppress_benchmark_output: pass -class artifacts_sandbox: +class ArtifactsSandbox: """Context manager that chdirs into tests/artifacts/ for the benchmark run. Any files created by agent tools land in the artifacts directory instead of diff --git a/lib/cli/src/crewai_cli/cli.py b/lib/cli/src/crewai_cli/cli.py index d317b2909..4b2eb557a 100644 --- a/lib/cli/src/crewai_cli/cli.py +++ b/lib/cli/src/crewai_cli/cli.py @@ -784,11 +784,11 @@ def _test_new_agents( fg="cyan", bold=True, ) - from crewai_cli.benchmark import artifacts_sandbox, suppress_benchmark_output + from crewai_cli.benchmark import ArtifactsSandbox, SuppressBenchmarkOutput progress.start() try: - with artifacts_sandbox(), suppress_benchmark_output(): + with ArtifactsSandbox(), SuppressBenchmarkOutput(): all_results = asyncio.run(_run_all()) finally: progress.stop() @@ -803,7 +803,7 @@ def _test_new_agents( continue agents_tested += 1 - for model_name, results in result.items(): + for results in result.values(): failed = [r for r in results if r.score < job["threshold"]] if failed: all_passed = False @@ -1612,12 +1612,12 @@ def benchmark( click.echo(f"Judge model: {judge_model}") click.echo() - from crewai_cli.benchmark import artifacts_sandbox, suppress_benchmark_output + from crewai_cli.benchmark import ArtifactsSandbox, SuppressBenchmarkOutput progress = _BenchmarkLiveProgress(console=_con) progress.start() try: - with artifacts_sandbox(), suppress_benchmark_output(): + with ArtifactsSandbox(), SuppressBenchmarkOutput(): results_by_model = asyncio.run( run_benchmark( agent_def=agent_path,