| 69 | |
| 70 | |
| 71 | def parse_args() -> argparse.Namespace: |
| 72 | parser = argparse.ArgumentParser( |
| 73 | description="List B.AI models and optionally smoke-test model calls." |
| 74 | ) |
| 75 | parser.add_argument( |
| 76 | "--api-key", |
| 77 | help="B.AI API key. Defaults to BAI_API_KEY environment variable.", |
| 78 | ) |
| 79 | parser.add_argument( |
| 80 | "--base-url", |
| 81 | default=os.getenv("BAI_API_BASE", DEFAULT_BASE_URL), |
| 82 | help=f"API base URL. Defaults to {DEFAULT_BASE_URL}.", |
| 83 | ) |
| 84 | parser.add_argument( |
| 85 | "--auth", |
| 86 | choices=("bearer", "x-api-key", "both"), |
| 87 | default="bearer", |
| 88 | help="Authentication header style. Docs support bearer and x-api-key.", |
| 89 | ) |
| 90 | parser.add_argument( |
| 91 | "--endpoint", |
| 92 | choices=("auto", "openai", "anthropic"), |
| 93 | default="auto", |
| 94 | help="Which compatible endpoint to test.", |
| 95 | ) |
| 96 | parser.add_argument( |
| 97 | "--model", |
| 98 | action="append", |
| 99 | help="Model id to test. Can be passed multiple times. Defaults to one model per endpoint.", |
| 100 | ) |
| 101 | parser.add_argument( |
| 102 | "--all", |
| 103 | action="store_true", |
| 104 | help="Smoke-test every listed model that matches --endpoint.", |
| 105 | ) |
| 106 | parser.add_argument( |
| 107 | "--list-only", |
| 108 | action="store_true", |
| 109 | help="Only call GET /v1/models; do not send chat/messages requests.", |
| 110 | ) |
| 111 | parser.add_argument("--prompt", default=DEFAULT_PROMPT) |
| 112 | parser.add_argument("--max-tokens", type=int, default=16) |
| 113 | parser.add_argument("--timeout", type=float, default=60.0) |
| 114 | return parser.parse_args() |
| 115 | |
| 116 | |
| 117 | def build_headers(api_key: str, auth: str) -> dict[str, str]: |