diff --git a/packages/prime/src/prime_cli/api/rl.py b/packages/prime/src/prime_cli/api/rl.py index 511ebc7f..b24988b7 100644 --- a/packages/prime/src/prime_cli/api/rl.py +++ b/packages/prime/src/prime_cli/api/rl.py @@ -59,6 +59,25 @@ def resolve_prices(self, category: str) -> tuple[Optional[float], Optional[float return getattr(self, legacy_attr), getattr(self, effective_attr) +class RLTeacherModelPricing(BaseModel): + """Inference pricing for a Hosted Training teacher model.""" + + prompt: Optional[float] = None + completion: Optional[float] = None + + +class RLTeacherModel(BaseModel): + """Generate-capable inference model available as a Hosted Training teacher.""" + + id: str = Field(..., description="Inference model ID") + name: str = Field(..., description="Display name") + description: Optional[str] = None + pricing: RLTeacherModelPricing = Field(default_factory=RLTeacherModelPricing) + generate_supported: bool = Field(False, alias="generateSupported") + + model_config = ConfigDict(populate_by_name=True) + + class RLRun(BaseModel): """Hosted Training run.""" @@ -168,6 +187,20 @@ def list_models(self, team_id: Optional[str] = None) -> List[RLModel]: raise APIError(f"Failed to list Hosted Training models: {e.response.text}") raise APIError(f"Failed to list Hosted Training models: {str(e)}") + def list_teacher_models(self, team_id: Optional[str] = None) -> List[RLTeacherModel]: + """List generate-capable inference models for Hosted Training teachers.""" + try: + params = {} + if team_id: + params["team_id"] = team_id + response = self.client.get("/rft/teacher-models", params=params if params else None) + models_data = response.get("models", []) + return [RLTeacherModel.model_validate(model) for model in models_data] + except Exception as e: + if hasattr(e, "response") and hasattr(e.response, "text"): + raise APIError(f"Failed to list Hosted Training teacher models: {e.response.text}") + raise APIError(f"Failed to list Hosted Training teacher models: {str(e)}") + def list_runs(self, team_id: Optional[str] = None) -> List[RLRun]: """List Hosted Training runs for the authenticated user.""" try: diff --git a/packages/prime/src/prime_cli/commands/rl.py b/packages/prime/src/prime_cli/commands/rl.py index d055bbd8..97175a17 100644 --- a/packages/prime/src/prime_cli/commands/rl.py +++ b/packages/prime/src/prime_cli/commands/rl.py @@ -53,6 +53,8 @@ "effective_training_price_per_mtok?, " "effective_inference_input_price_per_mtok?, " "effective_inference_output_price_per_mtok?, promo_label?}", + "with --teachers: .models[] = {id, name, description?, " + "pricing: {prompt?, completion?}, generate_supported}", ) RL_LIST_JSON_HELP = json_output_help( @@ -1502,6 +1504,11 @@ def _format(list_p: Any, eff_p: Any) -> str: @app.command("models", rich_help_panel="Commands", epilog=RL_MODELS_JSON_HELP) def list_models( output: str = typer.Option("table", "--output", "-o", help="Output format: table or json"), + teachers: bool = typer.Option( + False, + "--teachers", + help="List generate-capable teacher models instead of trainable student models.", + ), ) -> None: """List available models for Hosted Training.""" validate_output_format(output, console) @@ -1511,6 +1518,40 @@ def list_models( rl_client = RLClient(api_client) config = Config() + if teachers: + models = rl_client.list_teacher_models(team_id=config.team_id) + + if output == "json": + output_data_as_json({"models": [m.model_dump() for m in models]}, console) + return + + if not models: + console.print("[yellow]No teacher models available for Hosted Training.[/yellow]") + console.print( + "[dim]This could mean no tenant-visible inference models " + "support /generate.[/dim]" + ) + return + + table = Table(title="Hosted Training - Teacher Models") + table.add_column("Model", style="cyan") + table.add_column("Input", style="green", justify="right") + table.add_column("Output", style="green", justify="right") + + for model in sorted(models, key=lambda model: _model_name_sort_key(model.id)): + table.add_row( + model.id, + format_promo_price(model.pricing.prompt, None) or "-", + format_promo_price(model.pricing.completion, None) or "-", + ) + + table.caption = ( + "[dim]Prices are per 1M tokens. All listed models support token-level " + "/generate.[/dim]" + ) + console.print(table) + return + models = rl_client.list_models(team_id=config.team_id) if output == "json": diff --git a/packages/prime/tests/test_rl_api.py b/packages/prime/tests/test_rl_api.py index b49e7f5a..8c22d888 100644 --- a/packages/prime/tests/test_rl_api.py +++ b/packages/prime/tests/test_rl_api.py @@ -12,6 +12,17 @@ def __init__(self) -> None: def get(self, endpoint: str, params: dict[str, Any] | None = None) -> dict[str, Any]: self.requests.append((endpoint, params)) + if endpoint == "/rft/teacher-models": + return { + "models": [ + { + "id": "prime/generate-model", + "name": "Prime Generate Model", + "pricing": {"prompt": 1.5, "completion": 2.5}, + "generateSupported": True, + } + ] + } return { "chartData": { "histogramData": [ @@ -104,6 +115,18 @@ def test_create_run_sends_max_inflight_rollouts() -> None: assert run.max_inflight_rollouts == 96 +def test_list_teacher_models_forwards_team_context() -> None: + api_client = FakeAPIClient() + client = RLClient(api_client) # type: ignore[arg-type] + + models = client.list_teacher_models(team_id="team-1") + + assert api_client.requests[-1] == ("/rft/teacher-models", {"team_id": "team-1"}) + assert models[0].id == "prime/generate-model" + assert models[0].pricing.prompt == 1.5 + assert models[0].generate_supported is True + + def test_create_run_sends_sft_loss_and_teacher_config() -> None: api_client = FakeAPIClient() client = RLClient(api_client) # type: ignore[arg-type] diff --git a/packages/prime/tests/test_rl_models.py b/packages/prime/tests/test_rl_models.py index 1101e769..77715182 100644 --- a/packages/prime/tests/test_rl_models.py +++ b/packages/prime/tests/test_rl_models.py @@ -37,11 +37,33 @@ def _models_payload() -> Dict[str, Any]: } +def _teacher_models_payload() -> Dict[str, Any]: + return { + "models": [ + { + "id": "prime/generate-model", + "name": "Prime Generate Model", + "description": "Generate-capable teacher", + "pricing": {"prompt": 1.5, "completion": 2.5}, + "generateSupported": True, + }, + { + "id": "prime/cheap-teacher", + "name": "Prime Cheap Teacher", + "pricing": {"prompt": None, "completion": 0.25}, + "generateSupported": True, + }, + ] + } + + def _mock_get_factory(calls: List[str]): def mock_get(self: Any, endpoint: str, params: Dict[str, Any] | None = None) -> Dict[str, Any]: calls.append(endpoint) if endpoint == "/rft/models": return _models_payload() + if endpoint == "/rft/teacher-models": + return _teacher_models_payload() raise AssertionError(f"Unexpected endpoint: {endpoint}") return mock_get @@ -76,6 +98,35 @@ def test_models_json_includes_pricing(monkeypatch: pytest.MonkeyPatch) -> None: assert data["models"][1]["training_price_per_mtok"] is None +def test_teacher_models_table_renders_pricing(monkeypatch: pytest.MonkeyPatch) -> None: + calls: List[str] = [] + monkeypatch.setattr("prime_cli.core.APIClient.get", _mock_get_factory(calls)) + + result = CliRunner().invoke(app, ["train", "models", "--teachers"], env={"COLUMNS": "200"}) + + assert result.exit_code == 0, result.output + assert "prime/generate-model" in result.output + assert "$1.5" in result.output + assert "$2.5" in result.output + assert "prime/cheap-teacher" in result.output + assert "$0.25" in result.output + assert calls == ["/rft/teacher-models"] + + +def test_teacher_models_json_includes_generate_support( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("prime_cli.core.APIClient.get", _mock_get_factory([])) + + result = CliRunner().invoke(app, ["train", "models", "--teachers", "--output", "json"]) + + assert result.exit_code == 0, result.output + data = json.loads(result.output) + assert data["models"][0]["id"] == "prime/generate-model" + assert data["models"][0]["pricing"]["prompt"] == 1.5 + assert data["models"][0]["generate_supported"] is True + + def test_models_handles_backend_without_pricing_fields( monkeypatch: pytest.MonkeyPatch, ) -> None: