From c64ce5502c65bbebc0ccd7fe65f908970589bbd2 Mon Sep 17 00:00:00 2001 From: Matt Zhou Date: Fri, 13 Dec 2024 17:20:33 -0800 Subject: [PATCH] drop metadata.py and api keys table --- ...88e702f85e_drop_api_tokens_table_in_oss.py | 42 ++++++++++++++++ ...505cc7eca9_create_a_baseline_migrations.py | 10 ++-- letta/cli/cli.py | 1 - letta/metadata.py | 50 ------------------- letta/services/organization_manager.py | 1 - 5 files changed, 47 insertions(+), 57 deletions(-) create mode 100644 alembic/versions/4e88e702f85e_drop_api_tokens_table_in_oss.py delete mode 100644 letta/metadata.py diff --git a/alembic/versions/4e88e702f85e_drop_api_tokens_table_in_oss.py b/alembic/versions/4e88e702f85e_drop_api_tokens_table_in_oss.py new file mode 100644 index 0000000000..75a90445a0 --- /dev/null +++ b/alembic/versions/4e88e702f85e_drop_api_tokens_table_in_oss.py @@ -0,0 +1,42 @@ +"""Drop api tokens table in OSS + +Revision ID: 4e88e702f85e +Revises: d05669b60ebe +Create Date: 2024-12-13 17:19:55.796210 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "4e88e702f85e" +down_revision: Union[str, None] = "d05669b60ebe" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index("tokens_idx_key", table_name="tokens") + op.drop_index("tokens_idx_user", table_name="tokens") + op.drop_table("tokens") + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "tokens", + sa.Column("id", sa.VARCHAR(), autoincrement=False, nullable=False), + sa.Column("user_id", sa.VARCHAR(), autoincrement=False, nullable=False), + sa.Column("key", sa.VARCHAR(), autoincrement=False, nullable=False), + sa.Column("name", sa.VARCHAR(), autoincrement=False, nullable=True), + sa.PrimaryKeyConstraint("id", name="tokens_pkey"), + ) + op.create_index("tokens_idx_user", "tokens", ["user_id"], unique=False) + op.create_index("tokens_idx_key", "tokens", ["key"], unique=False) + # ### end Alembic commands ### diff --git a/alembic/versions/9a505cc7eca9_create_a_baseline_migrations.py b/alembic/versions/9a505cc7eca9_create_a_baseline_migrations.py index 479ca223e0..21f6a39613 100644 --- a/alembic/versions/9a505cc7eca9_create_a_baseline_migrations.py +++ b/alembic/versions/9a505cc7eca9_create_a_baseline_migrations.py @@ -12,7 +12,7 @@ import sqlalchemy as sa from sqlalchemy.dialects import postgresql -import letta.metadata +import letta.orm from alembic import op # revision identifiers, used by Alembic. @@ -43,8 +43,8 @@ def upgrade() -> None: sa.Column("memory", sa.JSON(), nullable=True), sa.Column("system", sa.String(), nullable=True), sa.Column("agent_type", sa.String(), nullable=True), - sa.Column("llm_config", letta.metadata.LLMConfigColumn(), nullable=True), - sa.Column("embedding_config", letta.metadata.EmbeddingConfigColumn(), nullable=True), + sa.Column("llm_config", letta.orm.custom_columns.LLMConfigColumn(), nullable=True), + sa.Column("embedding_config", letta.orm.custom_columns.EmbeddingConfigColumn(), nullable=True), sa.Column("metadata_", sa.JSON(), nullable=True), sa.Column("tools", sa.JSON(), nullable=True), sa.PrimaryKeyConstraint("id"), @@ -119,7 +119,7 @@ def upgrade() -> None: sa.Column("agent_id", sa.String(), nullable=True), sa.Column("source_id", sa.String(), nullable=True), sa.Column("embedding", pgvector.sqlalchemy.Vector(dim=4096), nullable=True), - sa.Column("embedding_config", letta.metadata.EmbeddingConfigColumn(), nullable=True), + sa.Column("embedding_config", letta.orm.custom_columns.EmbeddingConfigColumn(), nullable=True), sa.Column("metadata_", sa.JSON(), nullable=True), sa.Column("created_at", sa.DateTime(timezone=True), nullable=True), sa.PrimaryKeyConstraint("id"), @@ -131,7 +131,7 @@ def upgrade() -> None: sa.Column("user_id", sa.String(), nullable=False), sa.Column("name", sa.String(), nullable=False), sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=True), - sa.Column("embedding_config", letta.metadata.EmbeddingConfigColumn(), nullable=True), + sa.Column("embedding_config", letta.orm.custom_columns.EmbeddingConfigColumn(), nullable=True), sa.Column("description", sa.String(), nullable=True), sa.Column("metadata_", sa.JSON(), nullable=True), sa.PrimaryKeyConstraint("id"), diff --git a/letta/cli/cli.py b/letta/cli/cli.py index 5fa5d10e17..e5a649f7ce 100644 --- a/letta/cli/cli.py +++ b/letta/cli/cli.py @@ -18,7 +18,6 @@ ) from letta.local_llm.constants import ASSISTANT_MESSAGE_CLI_SYMBOL from letta.log import get_logger -from letta.metadata import MetadataStore from letta.schemas.enums import OptionState from letta.schemas.memory import ChatMemory, Memory from letta.server.server import logger as server_logger diff --git a/letta/metadata.py b/letta/metadata.py deleted file mode 100644 index 363d3ea1fb..0000000000 --- a/letta/metadata.py +++ /dev/null @@ -1,50 +0,0 @@ -""" Metadata store for user/agent/data_source information""" - -from sqlalchemy import JSON, TypeDecorator - -from letta.schemas.embedding_config import EmbeddingConfig -from letta.schemas.llm_config import LLMConfig - - -class LLMConfigColumn(TypeDecorator): - """Custom type for storing LLMConfig as JSON""" - - impl = JSON - cache_ok = True - - def load_dialect_impl(self, dialect): - return dialect.type_descriptor(JSON()) - - def process_bind_param(self, value, dialect): - if value: - # return vars(value) - if isinstance(value, LLMConfig): - return value.model_dump() - return value - - def process_result_value(self, value, dialect): - if value: - return LLMConfig(**value) - return value - - -class EmbeddingConfigColumn(TypeDecorator): - """Custom type for storing EmbeddingConfig as JSON""" - - impl = JSON - cache_ok = True - - def load_dialect_impl(self, dialect): - return dialect.type_descriptor(JSON()) - - def process_bind_param(self, value, dialect): - if value: - # return vars(value) - if isinstance(value, EmbeddingConfig): - return value.model_dump() - return value - - def process_result_value(self, value, dialect): - if value: - return EmbeddingConfig(**value) - return value diff --git a/letta/services/organization_manager.py b/letta/services/organization_manager.py index f98ba65dba..fc86b05fe7 100644 --- a/letta/services/organization_manager.py +++ b/letta/services/organization_manager.py @@ -13,7 +13,6 @@ class OrganizationManager: DEFAULT_ORG_NAME = "default_org" def __init__(self): - # This is probably horrible but we reuse this technique from metadata.py # TODO: Please refactor this out # I am currently working on a ORM refactor and would like to make a more minimal set of changes # - Matt