models.py 1.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445
  1. from __future__ import annotations
  2. from typing import TYPE_CHECKING
  3. from argparse import ArgumentParser
  4. from .._utils import get_client, print_model
  5. from .._models import BaseModel
  6. if TYPE_CHECKING:
  7. from argparse import _SubParsersAction
  8. def register(subparser: _SubParsersAction[ArgumentParser]) -> None:
  9. sub = subparser.add_parser("models.list")
  10. sub.set_defaults(func=CLIModels.list)
  11. sub = subparser.add_parser("models.retrieve")
  12. sub.add_argument("-i", "--id", required=True, help="The model ID")
  13. sub.set_defaults(func=CLIModels.get, args_model=CLIModelIDArgs)
  14. sub = subparser.add_parser("models.delete")
  15. sub.add_argument("-i", "--id", required=True, help="The model ID")
  16. sub.set_defaults(func=CLIModels.delete, args_model=CLIModelIDArgs)
  17. class CLIModelIDArgs(BaseModel):
  18. id: str
  19. class CLIModels:
  20. @staticmethod
  21. def get(args: CLIModelIDArgs) -> None:
  22. model = get_client().models.retrieve(model=args.id)
  23. print_model(model)
  24. @staticmethod
  25. def delete(args: CLIModelIDArgs) -> None:
  26. model = get_client().models.delete(model=args.id)
  27. print_model(model)
  28. @staticmethod
  29. def list() -> None:
  30. models = get_client().models.list()
  31. for model in models:
  32. print_model(model)