feat(benchmark/cli): Add challenge list, challenge info subcommands

- Add `challenge list` command with options `--all`, `--names`, `--json`
   - Add `tabular` dependency
   - Add `.utils.utils.sorted_by_enum_index` function to easily sort lists by an enum value/property based on the order of the enum's definition
- Add `challenge info [name]` command with option `--json`
   - Add `.utils.utils.pretty_print_model` routine to pretty-print Pydantic models
- Refactor `config` subcommand to use `pretty_print_model`
This commit is contained in:
Reinier van der Leer
2024-02-16 15:17:11 +01:00
parent 70e345b2ce
commit 23d58a3cc0
5 changed files with 219 additions and 6 deletions

View File

@@ -3,10 +3,13 @@ import json
import logging
import os
import re
from enum import Enum
from pathlib import Path
from typing import Any, Optional
from typing import Any, Callable, Iterable, Optional, TypeVar, overload
import click
from dotenv import load_dotenv
from pydantic import BaseModel
from agbenchmark.reports.processing.report_types import Test
from agbenchmark.utils.data_types import DIFFICULTY_MAP, DifficultyLevel
@@ -17,6 +20,9 @@ AGENT_NAME = os.getenv("AGENT_NAME")
logger = logging.getLogger(__name__)
T = TypeVar("T")
E = TypeVar("E", bound=Enum)
def replace_backslash(value: Any) -> Any:
if isinstance(value, str):
@@ -124,6 +130,42 @@ def write_pretty_json(data, json_file):
f.write("\n")
def pretty_print_model(model: BaseModel, include_header: bool = True) -> None:
indent = ""
if include_header:
# Try to find the ID and/or name attribute of the model
id, name = None, None
for attr, value in model.dict().items():
if attr == "id" or attr.endswith("_id"):
id = value
if attr.endswith("name"):
name = value
if id and name:
break
identifiers = [v for v in [name, id] if v]
click.echo(
f"{model.__repr_name__()}{repr(identifiers) if identifiers else ''}:"
)
indent = " " * 2
k_col_width = max(len(k) for k in model.dict().keys())
for k, v in model.dict().items():
v_fmt = repr(v)
if v is None or v == "":
v_fmt = click.style(v_fmt, fg="black")
elif type(v) is bool:
v_fmt = click.style(v_fmt, fg="green" if v else "red")
elif type(v) is str and "\n" in v:
v_fmt = f"\n{v}".replace(
"\n", f"\n{indent} {click.style('|', fg='black')} "
)
if isinstance(v, Enum):
v_fmt = click.style(v.value, fg="blue")
elif type(v) is list and len(v) > 0 and isinstance(v[0], Enum):
v_fmt = ", ".join(click.style(lv.value, fg="blue") for lv in v)
click.echo(f"{indent}{k: <{k_col_width}} = {v_fmt}")
def deep_sort(obj):
"""
Recursively sort the keys in JSON object
@@ -133,3 +175,38 @@ def deep_sort(obj):
if isinstance(obj, list):
return [deep_sort(elem) for elem in obj]
return obj
@overload
def sorted_by_enum_index(
sortable: Iterable[E],
enum: type[E],
*,
reverse: bool = False,
) -> list[E]:
...
@overload
def sorted_by_enum_index(
sortable: Iterable[T],
enum: type[Enum],
*,
key: Callable[[T], Enum | None],
reverse: bool = False,
) -> list[T]:
...
def sorted_by_enum_index(
sortable: Iterable[T],
enum: type[Enum],
*,
key: Callable[[T], Enum | None] = lambda x: x, # type: ignore
reverse: bool = False,
) -> list[T]:
return sorted(
sortable,
key=lambda x: enum._member_names_.index(e.name) if (e := key(x)) else 420e3,
reverse=reverse,
)