From 9fc0696052f40a24e5a7f87e5642452d982f97a9 Mon Sep 17 00:00:00 2001 From: Francisco Javier Arceo Date: Tue, 16 Sep 2025 09:11:17 -0400 Subject: [PATCH] fix: Fixing prompts import warning Signed-off-by: Francisco Javier Arceo --- llama_stack/core/distribution.py | 3 ++- tests/unit/distribution/test_distribution.py | 14 +++++++++++++- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/llama_stack/core/distribution.py b/llama_stack/core/distribution.py index 977eb5393..3e09f0ccc 100644 --- a/llama_stack/core/distribution.py +++ b/llama_stack/core/distribution.py @@ -70,7 +70,8 @@ def builtin_automatically_routed_apis() -> list[AutoRoutedApiInfo]: def providable_apis() -> list[Api]: routing_table_apis = {x.routing_table_api for x in builtin_automatically_routed_apis()} - return [api for api in Api if api not in routing_table_apis and api != Api.inspect and api != Api.providers] + internal_apis = {Api.inspect, Api.providers, Api.prompts} + return [api for api in Api if api not in routing_table_apis and api not in internal_apis] def _load_remote_provider_spec(spec_data: dict[str, Any], api: Api) -> ProviderSpec: diff --git a/tests/unit/distribution/test_distribution.py b/tests/unit/distribution/test_distribution.py index c72106e46..00f7170d5 100644 --- a/tests/unit/distribution/test_distribution.py +++ b/tests/unit/distribution/test_distribution.py @@ -12,7 +12,7 @@ import yaml from pydantic import BaseModel, Field, ValidationError from llama_stack.core.datatypes import Api, Provider, StackRunConfig -from llama_stack.core.distribution import get_provider_registry +from llama_stack.core.distribution import get_provider_registry, providable_apis from llama_stack.providers.datatypes import ProviderSpec @@ -152,6 +152,18 @@ class TestProviderRegistry: assert registry[Api.inference]["test_provider"].provider_type == "test_provider" assert registry[Api.inference]["test_provider"].api == Api.inference + def test_internal_apis_excluded(self): + """Test that internal APIs are excluded and other APIs are included in providable APIs.""" + apis = providable_apis() + internal_apis = {Api.inspect, Api.providers, Api.prompts} + + for internal_api in internal_apis: + assert internal_api not in apis, f"Internal API {internal_api} should not be in providable_apis" + + included_apis = {Api.inference, Api.safety, Api.agents} + for api in included_apis: + assert api in apis, f"API {api} should be in providable_apis" + def test_external_remote_providers(self, api_directories, mock_providers, base_config, provider_spec_yaml): """Test loading external remote providers from YAML files.""" remote_dir, _ = api_directories