use class methods and construct_stack

src/llama_stack_client/lib/direct/direct.py

@@ -9,8 +9,7 @@
 from llama_stack.distribution.stack import (
     get_stack_run_config_from_template,
 )
-
-from llama_stack.distribution.store.registry import create_dist_registry
+from llama_stack.distribution.stack import construct_stack
 from pydantic import BaseModel
 
 from ..._base_client import ResponseT
@@ -20,22 +19,31 @@
 
 
 class LlamaStackDirectClient(LlamaStackClient):
-    def __init__(self, config: StackRunConfig | str, **kwargs):
+    def __init__(self, config: StackRunConfig, **kwargs):
+        raise TypeError("Use from_yaml() or from_template() instead of direct initialization")
+
+    @classmethod
+    async def from_config(cls, config: StackRunConfig, **kwargs):
+        instance = object.__new__(cls)
+        await instance._initialize(config, **kwargs)
+        return instance
+
+    @classmethod
+    async def from_template(cls, template_name: str, **kwargs):
+        config = get_stack_run_config_from_template(template_name)
+        instance = object.__new__(cls)
+        await instance._initialize(config, **kwargs)
+        return instance
+
+    async def _initialize(self, config: StackRunConfig, **kwargs) -> None:
         super().__init__(**kwargs)
         self.endpoints = get_all_api_endpoints()
-
-        # Allow initialization with template name string
-        if isinstance(config, str):
-            self.config = get_stack_run_config_from_template(config)
-        else:
-            self.config = config
-
-        self.dist_registry = None
+        self.config = config
         self.impls = None
+        await self.initialize()
 
     async def initialize(self) -> None:
-        self.dist_registry, _ = await create_dist_registry(self.config.metadata_store, image_name="")
-        self.impls = await resolve_impls(self.config, get_provider_registry(), self.dist_registry)
+        self.impls = await construct_stack(self.config)
 
     def _convert_param(self, param_type: Any, value: Any) -> Any:
         origin = get_origin(param_type)

src/llama_stack_client/lib/direct/test.py

@@ -12,7 +12,7 @@ async def main(config_path: str):
 
     run_config = parse_and_maybe_upgrade_config(config_dict)
 
-    client = LlamaStackDirectClient(config=run_config)
+    client = await LlamaStackDirectClient.from_config(run_config)
     await client.initialize()
 
     response = await client.models.list()