Make config multiuser

api/alembic/versions/c78d76a6d65c_add_team_id_to_config_table.py

@@ -0,0 +1,161 @@
+"""add_team_id_to_config_table
+
+Revision ID: c78d76a6d65c
+Revises: 1f7cb465d15a
+Create Date: 2025-12-04 11:23:22.165544
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = "c78d76a6d65c"
+down_revision: Union[str, Sequence[str], None] = "1f7cb465d15a"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    connection = op.get_bind()
+
+    # Check existing columns
+    column_result = connection.execute(sa.text("PRAGMA table_info(config)"))
+    existing_columns = [row[1] for row in column_result.fetchall()]
+
+    # Get existing indexes by querying SQLite directly
+    # SQLite stores unique constraints as unique indexes
+    index_result = connection.execute(
+        sa.text("SELECT name FROM sqlite_master WHERE type='index' AND tbl_name='config'")
+    )
+    existing_index_names = [row[0] for row in index_result.fetchall()]
+
+    # Add columns first (outside batch mode to avoid circular dependency)
+    # Only add if they don't already exist
+    if "user_id" not in existing_columns:
+        op.add_column("config", sa.Column("user_id", sa.String(), nullable=True))
+    if "team_id" not in existing_columns:
+        op.add_column("config", sa.Column("team_id", sa.String(), nullable=True))
+
+    # Handle indexes outside of batch mode to avoid type inference issues
+    # Drop existing unique index on key if it exists (to recreate as non-unique)
+    if "ix_config_key" in existing_index_names:
+        # Check if it's unique by querying the index definition
+        index_info = connection.execute(
+            sa.text("SELECT sql FROM sqlite_master WHERE type='index' AND name='ix_config_key'")
+        ).fetchone()
+        if index_info and index_info[0] and "UNIQUE" in index_info[0].upper():
+            # Drop the unique index using raw SQL to avoid batch mode issues
+            connection.execute(sa.text("DROP INDEX IF EXISTS ix_config_key"))
+            existing_index_names.remove("ix_config_key")  # Update our list
+
+    # Create new indexes (non-unique) - these can be done outside batch mode
+    if "ix_config_key" not in existing_index_names:
+        op.create_index("ix_config_key", "config", ["key"], unique=False)
+    if "ix_config_user_id" not in existing_index_names:
+        op.create_index("ix_config_user_id", "config", ["user_id"], unique=False)
+    if "ix_config_team_id" not in existing_index_names:
+        op.create_index("ix_config_team_id", "config", ["team_id"], unique=False)
+
+    # For SQLite, unique constraints are stored as unique indexes
+    # Create the unique constraint as a unique index using raw SQL to avoid batch mode issues
+    if "uq_config_user_team_key" not in existing_index_names:
+        connection.execute(
+            sa.text("CREATE UNIQUE INDEX IF NOT EXISTS uq_config_user_team_key ON config(user_id, team_id, key)")
+        )
+
+    # Migrate existing configs to admin user's first team
+    # Note: Don't call connection.commit() - Alembic manages transactions
+    connection = op.get_bind()
+    # Find admin user's first team
+    admin_team_result = connection.execute(
+        sa.text("""
+            SELECT ut.team_id
+            FROM users_teams ut
+            JOIN user u ON ut.user_id = u.id
+            WHERE u.email = 'admin@example.com'
+            LIMIT 1
+        """)
+    )
+    admin_team_row = admin_team_result.fetchone()
+
+    if admin_team_row:
+        admin_team_id = admin_team_row[0]
+        # Update all existing configs (where team_id is NULL) to use admin team
+        connection.execute(
+            sa.text("UPDATE config SET team_id = :team_id WHERE team_id IS NULL"), {"team_id": admin_team_id}
+        )
+        print(f"✅ Migrated existing configs to team {admin_team_id}")
+    else:
+        # If no admin team found, try to get any user's first team
+        any_team_result = connection.execute(sa.text("SELECT team_id FROM users_teams LIMIT 1"))
+        any_team_row = any_team_result.fetchone()
+        if any_team_row:
+            any_team_id = any_team_row[0]
+            connection.execute(
+                sa.text("UPDATE config SET team_id = :team_id WHERE team_id IS NULL"), {"team_id": any_team_id}
+            )
+            print(f"✅ Migrated existing configs to team {any_team_id}")
+        else:
+            # No teams found, delete existing configs
+            deleted_count = connection.execute(sa.text("DELETE FROM config WHERE team_id IS NULL")).rowcount
+            print(f"⚠️  No teams found, deleted {deleted_count} config entries")
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    connection = op.get_bind()
+
+    # Check existing indexes
+    index_result = connection.execute(
+        sa.text("SELECT name FROM sqlite_master WHERE type='index' AND tbl_name='config'")
+    )
+    existing_index_names = [row[0] for row in index_result.fetchall()]
+
+    # Check existing columns
+    column_result = connection.execute(sa.text("PRAGMA table_info(config)"))
+    existing_columns = [row[1] for row in column_result.fetchall()]
+
+    # Drop indexes and constraints outside of batch mode to avoid type inference issues
+    # Drop unique constraint (stored as unique index in SQLite)
+    if "uq_config_user_team_key" in existing_index_names:
+        connection.execute(sa.text("DROP INDEX IF EXISTS uq_config_user_team_key"))
+
+    # Drop indexes
+    if "ix_config_team_id" in existing_index_names:
+        op.drop_index("ix_config_team_id", table_name="config")
+    if "ix_config_user_id" in existing_index_names:
+        op.drop_index("ix_config_user_id", table_name="config")
+    if "ix_config_key" in existing_index_names:
+        op.drop_index("ix_config_key", table_name="config")
+
+    # Drop columns using raw SQL to avoid batch mode type inference issues
+    # SQLite doesn't support DROP COLUMN directly, so we recreate the table
+    if "team_id" in existing_columns or "user_id" in existing_columns:
+        # Create new table without user_id and team_id columns
+        connection.execute(
+            sa.text("""
+                CREATE TABLE config_new (
+                    id INTEGER NOT NULL PRIMARY KEY,
+                    key VARCHAR NOT NULL,
+                    value VARCHAR
+                )
+            """)
+        )
+        # Copy data from old table to new table (only id, key, value columns)
+        connection.execute(sa.text("INSERT INTO config_new (id, key, value) SELECT id, key, value FROM config"))
+        # Drop old table (this also drops all indexes)
+        connection.execute(sa.text("DROP TABLE config"))
+        # Rename new table to original name
+        connection.execute(sa.text("ALTER TABLE config_new RENAME TO config"))
+        # Recreate the original unique index on key (it was dropped with the old table)
+        op.create_index("ix_config_key", "config", ["key"], unique=True)
+    else:
+        # If we're not dropping columns, just recreate the unique index on key
+        op.create_index("ix_config_key", "config", ["key"], unique=True)
+    # ### end Alembic commands ###

api/test/api/test_config.py

@@ -2,7 +2,7 @@
 def test_set_config(client):
     response = client.get("/config/set", params={"k": "api_test_key", "v": "test_value"})
     assert response.status_code == 200
-    assert response.json() == {"key": "api_test_key", "value": "test_value"}
+    assert response.json() == {"key": "api_test_key", "value": "test_value", "team_wide": True}
 
 
 def test_get_config(client):

api/test/server/test_config.py

@@ -22,7 +22,7 @@ def test_set(live_server):
 
     response = requests.get(f"{live_server}/config/set", params={"k": "message", "v": "Hello, World!"}, headers=headers)
     assert response.status_code == 200
-    assert response.json() == {"key": "message", "value": "Hello, World!"}
+    assert response.json() == {"key": "message", "value": "Hello, World!", "team_wide": True}
 
 
 @pytest.mark.live_server

api/transformerlab/db/db.py

@@ -1,5 +1,4 @@
 from sqlalchemy import select
-from sqlalchemy.dialects.sqlite import insert  # Correct import for SQLite upsert
 
 # from sqlalchemy import create_engine
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -25,17 +24,93 @@ async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
 ###############
 
 
-async def config_get(key: str):
+async def config_get(key: str, user_id: str | None = None, team_id: str | None = None):
+    """
+    Get config value with priority: user-specific -> team-specific -> global.
+
+    Priority order:
+    1. User-specific (user_id set, team_id matches current team)
+    2. Team-specific (user_id IS NULL, team_id set)
+    """
     async with async_session() as session:
-        result = await session.execute(select(Config.value).where(Config.key == key))
+        # First try user-specific config (if user_id provided)
+        if user_id and team_id:
+            result = await session.execute(
+                select(Config.value)
+                .where(Config.key == key, Config.user_id == user_id, Config.team_id == team_id)
+                .limit(1)
+            )
+            row = result.scalar_one_or_none()
+            if row is not None:
+                return row
+
+        # Then try team-specific config (user_id IS NULL, team_id set)
+        if team_id:
+            result = await session.execute(
+                select(Config.value)
+                .where(Config.key == key, Config.user_id.is_(None), Config.team_id == team_id)
+                .limit(1)
+            )
+            row = result.scalar_one_or_none()
+            if row is not None:
+                return row
+
+        # Finally fallback to global config (user_id IS NULL, team_id IS NULL)
+        result = await session.execute(
+            select(Config.value).where(Config.key == key, Config.user_id.is_(None), Config.team_id.is_(None)).limit(1)
+        )
         row = result.scalar_one_or_none()
         return row
 
 
-async def config_set(key: str, value: str):
-    stmt = insert(Config).values(key=key, value=value)
-    stmt = stmt.on_conflict_do_update(index_elements=["key"], set_={"value": value})
+async def config_set(key: str, value: str, user_id: str | None = None, team_id: str | None = None):
+    """
+    Set config value.
+
+    Args:
+        key: Config key
+        value: Config value
+        user_id: User ID for user-specific config. If None, sets team-wide config.
+        team_id: Team ID for team-specific config. If None, sets global config.
+
+    Config types:
+    - User-specific: user_id is set, team_id is set
+    - Team-wide: user_id is None, team_id is set
+    """
     async with async_session() as session:
-        await session.execute(stmt)
+        # Check if config already exists
+        if user_id is None and team_id is None:
+            # Global config: both user_id and team_id are NULL
+            result = await session.execute(
+                select(Config).where(Config.key == key, Config.user_id.is_(None), Config.team_id.is_(None))
+            )
+        elif user_id is None:
+            # Team-wide config: user_id is NULL, team_id is set
+            result = await session.execute(
+                select(Config).where(Config.key == key, Config.user_id.is_(None), Config.team_id == team_id)
+            )
+        else:
+            # User-specific config: both user_id and team_id are set
+            # Note: team_id should always be set when user_id is set (validated by router)
+            if team_id is None:
+                raise ValueError("team_id is required when user_id is set for user-specific configs")
+            result = await session.execute(
+                select(Config).where(
+                    Config.key == key,
+                    Config.user_id == user_id,
+                    Config.team_id == team_id,
+                )
+            )
+
+        existing = result.scalar_one_or_none()
+
+        if existing:
+            # Update existing config
+            existing.value = value
+        else:
+            # Insert new config
+            new_config = Config(key=key, value=value, user_id=user_id, team_id=team_id)
+            session.add(new_config)
+
         await session.commit()
     return

api/transformerlab/plugin_sdk/plugin_harness.py

@@ -18,14 +18,43 @@
 parser.add_argument("--plugin_dir", type=str, required=True)
 args, unknown = parser.parse_known_args()
 
+
+def set_config_env_vars(env_var: str, target_env_var: str = None, user_id: str = None, team_id: str = None):
+    try:
+        from transformerlab.plugin import get_db_config_value
+
+        value = get_db_config_value(env_var, user_id=user_id, team_id=team_id)
+        if value:
+            os.environ[target_env_var] = value
+            print(f"Set {target_env_var} from {'user' if user_id else 'team'} config: {value}")
+    except Exception as e:
+        print(f"Warning: Could not set {target_env_var} from {'user' if user_id else 'team'} config: {e}")
+
+
 # Set organization context from environment variable if provided
 # This allows plugins to have the correct org context without leaking to the API
 org_id = os.environ.get("_TFL_ORG_ID")
+user_id = os.environ.get("_TFL_USER_ID")  # Optional user_id for user-specific configs
+
 if org_id:
     try:
         from lab.dirs import set_organization_id
 
         set_organization_id(org_id)
+
+        try:
+            # Set HuggingFace token
+            set_config_env_vars("HuggingfaceUserAccessToken", "HF_TOKEN", user_id=user_id, team_id=org_id)
+            # Set WANDB API key
+            set_config_env_vars("WANDB_API_KEY", "WANDB_API_KEY", user_id=user_id, team_id=org_id)
+            # Set AI provider keys
+            set_config_env_vars("OPENAI_API_KEY", "OPENAI_API_KEY", user_id=user_id, team_id=org_id)
+            set_config_env_vars("ANTHROPIC_API_KEY", "ANTHROPIC_API_KEY", user_id=user_id, team_id=org_id)
+            set_config_env_vars("CUSTOM_MODEL_API_KEY", "CUSTOM_MODEL_API_KEY", user_id=user_id, team_id=org_id)
+            # Azure OpenAI details (JSON string)
+            set_config_env_vars("AZURE_OPENAI_DETAILS", "AZURE_OPENAI_DETAILS", user_id=user_id, team_id=org_id)
+        except Exception as e:
+            print(f"Warning: Could not set team/user-specific config env vars: {e}")
     except Exception as e:
         print(f"Warning: Could not set organization context: {e}")
 

api/transformerlab/plugin_sdk/transformerlab/__init__.py

@@ -0,0 +1,2 @@
+# TransformerLab plugin SDK package
+# This allows imports like: from transformerlab.sdk.v1.train import ...