llama-stack-mirror/llama_stack/providers/utils/kvstore/config.py
2024-09-22 16:34:44 -07:00

52 lines
1.4 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.
from enum import Enum
from typing import Literal, Optional, Union
from pydantic import BaseModel
from typing_extensions import Annotated
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 RedisKVStoreImplConfig(CommonConfig):
type: Literal[KVStoreType.redis.value] = KVStoreType.redis.value
host: str = "localhost"
port: int = 6379
class SqliteKVStoreImplConfig(CommonConfig):
type: Literal[KVStoreType.sqlite.value] = KVStoreType.sqlite.value
db_path: str = Field(
description="File path for the sqlite database",
)
class PostgresKVStoreImplConfig(CommonConfig):
type: Literal[KVStoreType.postgres.value] = KVStoreType.postgres.value
host: str = "localhost"
port: int = 5432
db: str = "llamastack"
user: str
password: Optional[str] = None
KVStoreConfig = Annotated[
Union[RedisKVStoreImplConfig, SqliteKVStoreImplConfig, PostgresKVStoreImplConfig],
Field(discriminator="type", default=KVStoreType.sqlite.value),
]