mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-06-27 18:50:41 +00:00
feat: add postgres deps to starter distro (#2360)
Once we have this, we can use the starter distro for the Kubernetes cluster demos.
This commit is contained in:
Ashwin Bharambe
GitHub
parent
e743257d1d
commit
b380cb463f
8 changed files with 68 additions and 16 deletions
|
|
@ -72,15 +72,15 @@ class PostgresKVStoreConfig(CommonConfig):
|
|||
table_name: str = "llamastack_kvstore"
|
||||
|
||||
@classmethod
|
||||
def sample_run_config(cls, table_name: str = "llamastack_kvstore"):
|
||||
def sample_run_config(cls, table_name: str = "llamastack_kvstore", **kwargs):
|
||||
return {
|
||||
"type": "postgres",
|
||||
"namespace": None,
|
||||
"host": "${env.POSTGRES_HOST:localhost}",
|
||||
"port": "${env.POSTGRES_PORT:5432}",
|
||||
"db": "${env.POSTGRES_DB}",
|
||||
"user": "${env.POSTGRES_USER}",
|
||||
"password": "${env.POSTGRES_PASSWORD}",
|
||||
"db": "${env.POSTGRES_DB:llamastack}",
|
||||
"user": "${env.POSTGRES_USER:llamastack}",
|
||||
"password": "${env.POSTGRES_PASSWORD:llamastack}",
|
||||
"table_name": "${env.POSTGRES_TABLE_NAME:" + table_name + "}",
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -74,6 +74,17 @@ class PostgresSqlStoreConfig(SqlAlchemySqlStoreConfig):
|
|||
def pip_packages(self) -> list[str]:
|
||||
return super().pip_packages + ["asyncpg"]
|
||||
|
||||
@classmethod
|
||||
def sample_run_config(cls, **kwargs):
|
||||
return cls(
|
||||
type="postgres",
|
||||
host="${env.POSTGRES_HOST:localhost}",
|
||||
port="${env.POSTGRES_PORT:5432}",
|
||||
db="${env.POSTGRES_DB:llamastack}",
|
||||
user="${env.POSTGRES_USER:llamastack}",
|
||||
password="${env.POSTGRES_PASSWORD:llamastack}",
|
||||
)
|
||||
|
||||
|
||||
SqlStoreConfig = Annotated[
|
||||
SqliteSqlStoreConfig | PostgresSqlStoreConfig,
|
||||
|
|
|
|||
|
|
@ -101,15 +101,7 @@ def get_distribution_template() -> DistributionTemplate:
|
|||
provider_id="vllm-inference",
|
||||
)
|
||||
)
|
||||
postgres_config = {
|
||||
"type": "postgres",
|
||||
"host": "${env.POSTGRES_HOST:localhost}",
|
||||
"port": "${env.POSTGRES_PORT:5432}",
|
||||
"db": "${env.POSTGRES_DB:llamastack}",
|
||||
"user": "${env.POSTGRES_USER:llamastack}",
|
||||
"password": "${env.POSTGRES_PASSWORD:llamastack}",
|
||||
}
|
||||
|
||||
postgres_config = PostgresSqlStoreConfig.sample_run_config()
|
||||
return DistributionTemplate(
|
||||
name=name,
|
||||
distro_type="self_hosted",
|
||||
|
|
@ -147,8 +139,8 @@ def get_distribution_template() -> DistributionTemplate:
|
|||
default_models=default_models,
|
||||
default_tool_groups=default_tool_groups,
|
||||
default_shields=[ShieldInput(shield_id="meta-llama/Llama-Guard-3-8B")],
|
||||
metadata_store=PostgresKVStoreConfig.model_validate(postgres_config),
|
||||
inference_store=PostgresSqlStoreConfig.model_validate(postgres_config),
|
||||
metadata_store=PostgresKVStoreConfig.sample_run_config(),
|
||||
inference_store=postgres_config,
|
||||
),
|
||||
},
|
||||
run_config_env_vars={
|
||||
|
|
|
|||
|
|
@ -79,7 +79,7 @@ metadata_store:
|
|||
db: ${env.POSTGRES_DB:llamastack}
|
||||
user: ${env.POSTGRES_USER:llamastack}
|
||||
password: ${env.POSTGRES_PASSWORD:llamastack}
|
||||
table_name: llamastack_kvstore
|
||||
table_name: ${env.POSTGRES_TABLE_NAME:llamastack_kvstore}
|
||||
inference_store:
|
||||
type: postgres
|
||||
host: ${env.POSTGRES_HOST:localhost}
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ distribution_spec:
|
|||
- remote::gemini
|
||||
- remote::groq
|
||||
- remote::sambanova
|
||||
- remote::remote-vllm
|
||||
- inline::sentence-transformers
|
||||
vector_io:
|
||||
- inline::sqlite-vec
|
||||
|
|
@ -37,4 +38,5 @@ distribution_spec:
|
|||
image_type: conda
|
||||
additional_pip_packages:
|
||||
- aiosqlite
|
||||
- asyncpg
|
||||
- sqlalchemy[asyncio]
|
||||
|
|
|
|||
|
|
@ -39,6 +39,13 @@ providers:
|
|||
config:
|
||||
url: https://api.sambanova.ai/v1
|
||||
api_key: ${env.SAMBANOVA_API_KEY:}
|
||||
- provider_id: remote-vllm
|
||||
provider_type: remote::remote-vllm
|
||||
config:
|
||||
url: ${env.VLLM_URL:http://localhost:8000/v1}
|
||||
max_tokens: ${env.VLLM_MAX_TOKENS:4096}
|
||||
api_token: ${env.VLLM_API_TOKEN:fake}
|
||||
tls_verify: ${env.VLLM_TLS_VERIFY:true}
|
||||
- provider_id: sentence-transformers
|
||||
provider_type: inline::sentence-transformers
|
||||
config: {}
|
||||
|
|
|
|||
|
|
@ -42,11 +42,13 @@ from llama_stack.providers.remote.inference.sambanova.config import SambaNovaImp
|
|||
from llama_stack.providers.remote.inference.sambanova.models import (
|
||||
MODEL_ENTRIES as SAMBANOVA_MODEL_ENTRIES,
|
||||
)
|
||||
from llama_stack.providers.remote.inference.vllm import VLLMInferenceAdapterConfig
|
||||
from llama_stack.providers.remote.vector_io.chroma.config import ChromaVectorIOConfig
|
||||
from llama_stack.providers.remote.vector_io.pgvector.config import (
|
||||
PGVectorVectorIOConfig,
|
||||
)
|
||||
from llama_stack.providers.utils.inference.model_registry import ProviderModelEntry
|
||||
from llama_stack.providers.utils.sqlstore.sqlstore import PostgresSqlStoreConfig
|
||||
from llama_stack.templates.template import (
|
||||
DistributionTemplate,
|
||||
RunConfigSettings,
|
||||
|
|
@ -87,6 +89,13 @@ def get_inference_providers() -> tuple[list[Provider], dict[str, list[ProviderMo
|
|||
SAMBANOVA_MODEL_ENTRIES,
|
||||
SambaNovaImplConfig.sample_run_config(api_key="${env.SAMBANOVA_API_KEY:}"),
|
||||
),
|
||||
(
|
||||
"remote-vllm",
|
||||
[],
|
||||
VLLMInferenceAdapterConfig.sample_run_config(
|
||||
url="${env.VLLM_URL:http://localhost:8000/v1}",
|
||||
),
|
||||
),
|
||||
]
|
||||
inference_providers = []
|
||||
available_models = {}
|
||||
|
|
@ -169,6 +178,8 @@ def get_distribution_template() -> DistributionTemplate:
|
|||
)
|
||||
|
||||
default_models = get_model_registry(available_models)
|
||||
|
||||
postgres_store = PostgresSqlStoreConfig.sample_run_config()
|
||||
return DistributionTemplate(
|
||||
name=name,
|
||||
distro_type="self_hosted",
|
||||
|
|
@ -177,6 +188,7 @@ def get_distribution_template() -> DistributionTemplate:
|
|||
template_path=None,
|
||||
providers=providers,
|
||||
available_models_by_provider=available_models,
|
||||
additional_pip_packages=postgres_store.pip_packages,
|
||||
run_configs={
|
||||
"run.yaml": RunConfigSettings(
|
||||
provider_overrides={
|
||||
|
|
@ -201,5 +213,25 @@ def get_distribution_template() -> DistributionTemplate:
|
|||
"",
|
||||
"OpenAI API Key",
|
||||
),
|
||||
"GROQ_API_KEY": (
|
||||
"",
|
||||
"Groq API Key",
|
||||
),
|
||||
"ANTHROPIC_API_KEY": (
|
||||
"",
|
||||
"Anthropic API Key",
|
||||
),
|
||||
"GEMINI_API_KEY": (
|
||||
"",
|
||||
"Gemini API Key",
|
||||
),
|
||||
"SAMBANOVA_API_KEY": (
|
||||
"",
|
||||
"SambaNova API Key",
|
||||
),
|
||||
"VLLM_URL": (
|
||||
"http://localhost:8000/v1",
|
||||
"VLLM URL",
|
||||
),
|
||||
},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -154,6 +154,11 @@ class DistributionTemplate(BaseModel):
|
|||
|
||||
available_models_by_provider: dict[str, list[ProviderModelEntry]] | None = None
|
||||
|
||||
# we may want to specify additional pip packages without necessarily indicating a
|
||||
# specific "default" inference store (which is what typically used to dictate additional
|
||||
# pip packages)
|
||||
additional_pip_packages: list[str] | None = None
|
||||
|
||||
def build_config(self) -> BuildConfig:
|
||||
additional_pip_packages: list[str] = []
|
||||
for run_config in self.run_configs.values():
|
||||
|
|
@ -161,6 +166,9 @@ class DistributionTemplate(BaseModel):
|
|||
if run_config_.inference_store:
|
||||
additional_pip_packages.extend(run_config_.inference_store.pip_packages)
|
||||
|
||||
if self.additional_pip_packages:
|
||||
additional_pip_packages.extend(self.additional_pip_packages)
|
||||
|
||||
return BuildConfig(
|
||||
distribution_spec=DistributionSpec(
|
||||
description=self.description,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue