mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-06-27 18:50:41 +00:00
2a31163178
# What does this PR do? Automatically generates - build.yaml - run.yaml - run-with-safety.yaml - parts of markdown docs for the distributions. ## Test Plan At this point, this only updates the YAMLs and the docs. Some testing (especially with ollama and vllm) has been performed but needs to be much more tested.
112 lines
3.3 KiB
Python
112 lines
3.3 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
import re
|
|
from enum import Enum
|
|
from typing import Literal, Optional, Union
|
|
|
|
from pydantic import BaseModel, Field, field_validator
|
|
from typing_extensions import Annotated
|
|
|
|
from llama_stack.distribution.utils.config_dirs import RUNTIME_BASE_DIR
|
|
|
|
|
|
class KVStoreType(Enum):
|
|
redis = "redis"
|
|
sqlite = "sqlite"
|
|
postgres = "postgres"
|
|
|
|
|
|
class CommonConfig(BaseModel):
|
|
namespace: Optional[str] = Field(
|
|
default=None,
|
|
description="All keys will be prefixed with this namespace",
|
|
)
|
|
|
|
|
|
class RedisKVStoreConfig(CommonConfig):
|
|
type: Literal[KVStoreType.redis.value] = KVStoreType.redis.value
|
|
host: str = "localhost"
|
|
port: int = 6379
|
|
|
|
@property
|
|
def url(self) -> str:
|
|
return f"redis://{self.host}:{self.port}"
|
|
|
|
@classmethod
|
|
def sample_run_config(cls):
|
|
return {
|
|
"type": "redis",
|
|
"namespace": None,
|
|
"host": "${env.REDIS_HOST:localhost}",
|
|
"port": "${env.REDIS_PORT:6379}",
|
|
}
|
|
|
|
|
|
class SqliteKVStoreConfig(CommonConfig):
|
|
type: Literal[KVStoreType.sqlite.value] = KVStoreType.sqlite.value
|
|
db_path: str = Field(
|
|
default=(RUNTIME_BASE_DIR / "kvstore.db").as_posix(),
|
|
description="File path for the sqlite database",
|
|
)
|
|
|
|
@classmethod
|
|
def sample_run_config(
|
|
cls, __distro_dir__: str = "runtime", db_name: str = "kvstore.db"
|
|
):
|
|
return {
|
|
"type": "sqlite",
|
|
"namespace": None,
|
|
"db_path": "${env.SQLITE_STORE_DIR:~/.llama/"
|
|
+ __distro_dir__
|
|
+ "}/"
|
|
+ db_name,
|
|
}
|
|
|
|
|
|
class PostgresKVStoreConfig(CommonConfig):
|
|
type: Literal[KVStoreType.postgres.value] = KVStoreType.postgres.value
|
|
host: str = "localhost"
|
|
port: int = 5432
|
|
db: str = "llamastack"
|
|
user: str
|
|
password: Optional[str] = None
|
|
table_name: str = "llamastack_kvstore"
|
|
|
|
@classmethod
|
|
def sample_run_config(cls, table_name: str = "llamastack_kvstore"):
|
|
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}",
|
|
"table_name": "${env.POSTGRES_TABLE_NAME:" + table_name + "}",
|
|
}
|
|
|
|
@classmethod
|
|
@field_validator("table_name")
|
|
def validate_table_name(cls, v: str) -> str:
|
|
# PostgreSQL identifiers rules:
|
|
# - Must start with a letter or underscore
|
|
# - Can contain letters, numbers, and underscores
|
|
# - Maximum length is 63 bytes
|
|
pattern = r"^[a-zA-Z_][a-zA-Z0-9_]*$"
|
|
if not re.match(pattern, v):
|
|
raise ValueError(
|
|
"Invalid table name. Must start with letter or underscore and contain only letters, numbers, and underscores"
|
|
)
|
|
if len(v) > 63:
|
|
raise ValueError("Table name must be less than 63 characters")
|
|
return v
|
|
|
|
|
|
KVStoreConfig = Annotated[
|
|
Union[RedisKVStoreConfig, SqliteKVStoreConfig, PostgresKVStoreConfig],
|
|
Field(discriminator="type", default=KVStoreType.sqlite.value),
|
|
]
|