diff --git a/llama_stack/distribution/build.py b/llama_stack/distribution/build.py index 5906614ed..4f9091a5d 100644 --- a/llama_stack/distribution/build.py +++ b/llama_stack/distribution/build.py @@ -43,23 +43,12 @@ def get_provider_dependencies( config: BuildConfig | DistributionTemplate, ) -> tuple[list[str], list[str]]: """Get normal and special dependencies from provider configuration.""" - # Extract providers based on config type if isinstance(config, DistributionTemplate): - providers = config.providers + config = config.build_config() + + providers = config.distribution_spec.providers + additional_pip_packages = config.additional_pip_packages - # TODO: This is a hack to get the dependencies for internal APIs into build - # We should have a better way to do this by formalizing the concept of "internal" APIs - # and providers, with a way to specify dependencies for them. - run_configs = config.run_configs - additional_pip_packages: list[str] = [] - if run_configs: - for run_config in run_configs.values(): - run_config_ = run_config.run_config(name="", providers={}, container_image=None) - if run_config_.inference_store: - additional_pip_packages.extend(run_config_.inference_store.pip_packages) - elif isinstance(config, BuildConfig): - providers = config.distribution_spec.providers - additional_pip_packages = config.additional_pip_packages deps = [] registry = get_provider_registry(config) for api_str, provider_or_providers in providers.items(): @@ -87,8 +76,7 @@ def get_provider_dependencies( else: normal_deps.append(package) - if additional_pip_packages: - normal_deps.extend(additional_pip_packages) + normal_deps.extend(additional_pip_packages or []) return list(set(normal_deps)), list(set(special_deps)) diff --git a/llama_stack/providers/registry/agents.py b/llama_stack/providers/registry/agents.py index e0801a8d1..e47f84c65 100644 --- a/llama_stack/providers/registry/agents.py +++ b/llama_stack/providers/registry/agents.py @@ -24,7 +24,7 @@ def available_providers() -> list[ProviderSpec]: "pandas", "scikit-learn", ] - + kvstore_dependencies(), + + kvstore_dependencies(), # TODO make this dynamic based on the kvstore config module="llama_stack.providers.inline.agents.meta_reference", config_class="llama_stack.providers.inline.agents.meta_reference.MetaReferenceAgentsImplConfig", api_dependencies=[ diff --git a/llama_stack/providers/utils/kvstore/config.py b/llama_stack/providers/utils/kvstore/config.py index 18a6eb8dd..e966e13ba 100644 --- a/llama_stack/providers/utils/kvstore/config.py +++ b/llama_stack/providers/utils/kvstore/config.py @@ -36,6 +36,10 @@ class RedisKVStoreConfig(CommonConfig): def url(self) -> str: return f"redis://{self.host}:{self.port}" + @property + def pip_packages(self) -> list[str]: + return ["redis"] + @classmethod def sample_run_config(cls): return { @@ -53,6 +57,10 @@ class SqliteKVStoreConfig(CommonConfig): description="File path for the sqlite database", ) + @property + def pip_packages(self) -> list[str]: + return ["aiosqlite"] + @classmethod def sample_run_config(cls, __distro_dir__: str, db_name: str = "kvstore.db"): return { @@ -100,6 +108,10 @@ class PostgresKVStoreConfig(CommonConfig): raise ValueError("Table name must be less than 63 characters") return v + @property + def pip_packages(self) -> list[str]: + return ["psycopg2-binary"] + class MongoDBKVStoreConfig(CommonConfig): type: Literal[KVStoreType.mongodb.value] = KVStoreType.mongodb.value @@ -110,6 +122,10 @@ class MongoDBKVStoreConfig(CommonConfig): password: str | None = None collection_name: str = "llamastack_kvstore" + @property + def pip_packages(self) -> list[str]: + return ["pymongo"] + @classmethod def sample_run_config(cls, collection_name: str = "llamastack_kvstore"): return { diff --git a/llama_stack/providers/utils/kvstore/kvstore.py b/llama_stack/providers/utils/kvstore/kvstore.py index 3a1ee8a26..426523d8e 100644 --- a/llama_stack/providers/utils/kvstore/kvstore.py +++ b/llama_stack/providers/utils/kvstore/kvstore.py @@ -10,6 +10,13 @@ from .config import KVStoreConfig, KVStoreType def kvstore_dependencies(): + """ + Returns all possible kvstore dependencies for registry/provider specifications. + + NOTE: For specific kvstore implementations, use config.pip_packages instead. + This function returns the union of all dependencies for cases where the specific + kvstore type is not known at declaration time (e.g., provider registries). + """ return ["aiosqlite", "psycopg2-binary", "redis", "pymongo"] diff --git a/llama_stack/templates/postgres-demo/build.yaml b/llama_stack/templates/postgres-demo/build.yaml index a7dee0787..6416cd00f 100644 --- a/llama_stack/templates/postgres-demo/build.yaml +++ b/llama_stack/templates/postgres-demo/build.yaml @@ -21,4 +21,5 @@ distribution_spec: image_type: conda additional_pip_packages: - asyncpg +- psycopg2-binary - sqlalchemy[asyncio] diff --git a/llama_stack/templates/template.py b/llama_stack/templates/template.py index b1ffba5be..712d2dcb4 100644 --- a/llama_stack/templates/template.py +++ b/llama_stack/templates/template.py @@ -186,8 +186,14 @@ class DistributionTemplate(BaseModel): additional_pip_packages: list[str] = [] for run_config in self.run_configs.values(): run_config_ = run_config.run_config(self.name, self.providers, self.container_image) + + # TODO: This is a hack to get the dependencies for internal APIs into build + # We should have a better way to do this by formalizing the concept of "internal" APIs + # and providers, with a way to specify dependencies for them. if run_config_.inference_store: additional_pip_packages.extend(run_config_.inference_store.pip_packages) + if run_config_.metadata_store: + additional_pip_packages.extend(run_config_.metadata_store.pip_packages) if self.additional_pip_packages: additional_pip_packages.extend(self.additional_pip_packages)