ece8dab05d
Co-authored-by: Shubham Naik <shub@memgpt.ai> Co-authored-by: Matt Zhou <mattzh1314@gmail.com> Co-authored-by: Shubham Naik <shubham.naik10@gmail.com> Co-authored-by: Caren Thomas <caren@letta.com> Co-authored-by: cpacker <packercharles@gmail.com>
55 lines
2.8 KiB
Python
55 lines
2.8 KiB
Python
from typing import TYPE_CHECKING, Dict, List, Optional
|
|
|
|
from sqlalchemy import JSON
|
|
from sqlalchemy import Enum as SqlEnum
|
|
from sqlalchemy import String, UniqueConstraint
|
|
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
|
|
|
from letta.orm.mixins import OrganizationMixin, SandboxConfigMixin
|
|
from letta.orm.sqlalchemy_base import SqlalchemyBase
|
|
from letta.schemas.sandbox_config import SandboxConfig as PydanticSandboxConfig
|
|
from letta.schemas.sandbox_config import SandboxEnvironmentVariable as PydanticSandboxEnvironmentVariable
|
|
from letta.schemas.sandbox_config import SandboxType
|
|
|
|
if TYPE_CHECKING:
|
|
from letta.orm.organization import Organization
|
|
|
|
|
|
class SandboxConfig(SqlalchemyBase, OrganizationMixin):
|
|
"""ORM model for sandbox configurations with JSON storage for arbitrary config data."""
|
|
|
|
__tablename__ = "sandbox_configs"
|
|
__pydantic_model__ = PydanticSandboxConfig
|
|
|
|
# For now, we only allow one type of sandbox config per organization
|
|
__table_args__ = (UniqueConstraint("type", "organization_id", name="uix_type_organization"),)
|
|
|
|
id: Mapped[str] = mapped_column(String, primary_key=True, nullable=False)
|
|
type: Mapped[SandboxType] = mapped_column(SqlEnum(SandboxType), nullable=False, doc="The type of sandbox.")
|
|
config: Mapped[Dict] = mapped_column(JSON, nullable=False, doc="The JSON configuration data.")
|
|
|
|
# relationships
|
|
organization: Mapped["Organization"] = relationship("Organization", back_populates="sandbox_configs")
|
|
sandbox_environment_variables: Mapped[List["SandboxEnvironmentVariable"]] = relationship(
|
|
"SandboxEnvironmentVariable", back_populates="sandbox_config", cascade="all, delete-orphan"
|
|
)
|
|
|
|
|
|
class SandboxEnvironmentVariable(SqlalchemyBase, OrganizationMixin, SandboxConfigMixin):
|
|
"""ORM model for environment variables associated with sandboxes."""
|
|
|
|
__tablename__ = "sandbox_environment_variables"
|
|
__pydantic_model__ = PydanticSandboxEnvironmentVariable
|
|
|
|
# We cannot have duplicate key names in the same sandbox, the env var would get overwritten
|
|
__table_args__ = (UniqueConstraint("key", "sandbox_config_id", name="uix_key_sandbox_config"),)
|
|
|
|
id: Mapped[str] = mapped_column(String, primary_key=True, nullable=False)
|
|
key: Mapped[str] = mapped_column(String, nullable=False, doc="The name of the environment variable.")
|
|
value: Mapped[str] = mapped_column(String, nullable=False, doc="The value of the environment variable.")
|
|
description: Mapped[Optional[str]] = mapped_column(String, nullable=True, doc="An optional description of the environment variable.")
|
|
|
|
# relationships
|
|
organization: Mapped["Organization"] = relationship("Organization", back_populates="sandbox_environment_variables")
|
|
sandbox_config: Mapped["SandboxConfig"] = relationship("SandboxConfig", back_populates="sandbox_environment_variables")
|