feat(db): add migration to simplify agent architecture

- Add migration to delete user_agents table and memories.agent_id - Rename user_agent_catalog to system_agents - Remove UserAgents model - Rename UserAgentCatalog to SystemAgents model
2026-03-04 11:04:13 +08:00
parent 7cc333a862
commit 2d410e8e84
4 changed files with 323 additions and 68 deletions
@@ -0,0 +1,319 @@
+"""simplify_agent_architecture
+
+Revision ID: 435419f8121c
+Revises: 50ae013ce530
+Create Date: 2026-03-04 11:01:11.127364
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+
+# revision identifiers, used by Alembic.
+revision: str = "435419f8121c"
+down_revision: Union[str, Sequence[str], None] = "50ae013ce530"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # 1. Delete memories.agent_id dependencies
+    op.drop_constraint("fk_memories_agent_id", "memories", type_="foreignkey")
+    op.drop_constraint("chk_memory_type_agent_id", "memories", type_="check")
+    op.execute("DROP INDEX IF EXISTS ix_memories_agent_type_status")
+    op.drop_column("memories", "agent_id")
+
+    # 2. Delete user_agents table
+    _drop_rls("user_agents")
+
+    op.drop_constraint("fk_user_agents_updated_by", "user_agents", type_="foreignkey")
+    op.drop_constraint("fk_user_agents_created_by", "user_agents", type_="foreignkey")
+    op.drop_constraint("fk_user_agents_llm_id", "user_agents", type_="foreignkey")
+    op.drop_constraint("fk_user_agents_user_id", "user_agents", type_="foreignkey")
+    op.drop_constraint("chk_agent_type", "user_agents", type_="check")
+    op.drop_constraint(
+        "uq_user_agents_user_id_agent_type", "user_agents", type_="unique"
+    )
+
+    op.execute("DROP INDEX IF EXISTS ix_user_agents_status")
+    op.execute("DROP INDEX IF EXISTS ix_user_agents_agent_type")
+
+    op.drop_table("user_agents")
+
+    # 3. Rename user_agent_catalog to system_agents
+    _drop_rls("user_agent_catalog")
+
+    op.rename_table("user_agent_catalog", "system_agents")
+
+    op.execute(
+        "ALTER TABLE system_agents RENAME CONSTRAINT fk_user_agent_catalog_llm_id "
+        "TO fk_system_agents_llm_id"
+    )
+    op.execute(
+        "ALTER TABLE system_agents RENAME CONSTRAINT chk_user_agent_catalog_status "
+        "TO chk_system_agents_status"
+    )
+
+    _enable_rls("system_agents")
+
+    # 4. Update trigger
+    op.execute("DROP TRIGGER IF EXISTS on_auth_user_created ON auth.users")
+    op.execute("DROP FUNCTION IF EXISTS public.create_profile_for_new_user()")
+
+    op.execute("""
+        CREATE OR REPLACE FUNCTION public.create_profile_for_new_user()
+        RETURNS trigger
+        LANGUAGE plpgsql
+        SECURITY DEFINER
+        SET search_path = public
+        AS $$
+        BEGIN
+            INSERT INTO public.profiles (id, username, avatar_url, bio, settings, created_at, updated_at)
+            VALUES (
+                NEW.id,
+                COALESCE(
+                    NEW.raw_user_meta_data ->> 'username',
+                    split_part(NEW.email, '@', 1),
+                    'user_' || substring(NEW.id::text, 1, 8)
+                ),
+                NULL,
+                NULL,
+                '{"agent_prompts": {}}'::jsonb,
+                now(),
+                now()
+            )
+            ON CONFLICT (id) DO NOTHING;
+            
+            RETURN NEW;
+        END;
+        $$
+    """)
+
+    op.execute("""
+        CREATE TRIGGER on_auth_user_created
+            AFTER INSERT ON auth.users
+            FOR EACH ROW EXECUTE FUNCTION public.create_profile_for_new_user()
+    """)
+
+    # 5. Update existing profiles.settings
+    op.execute("""
+        UPDATE profiles 
+        SET settings = jsonb_set(
+            COALESCE(settings, '{}'::jsonb),
+            '{agent_prompts}',
+            '{}'::jsonb
+        )
+        WHERE NOT settings ? 'agent_prompts'
+    """)
+
+
+def downgrade() -> None:
+    # 1. Revert trigger
+    op.execute("DROP TRIGGER IF EXISTS on_auth_user_created ON auth.users")
+    op.execute("DROP FUNCTION IF EXISTS public.create_profile_for_new_user()")
+
+    op.execute("""
+        CREATE OR REPLACE FUNCTION public.create_profile_for_new_user()
+        RETURNS trigger
+        LANGUAGE plpgsql
+        SECURITY DEFINER
+        SET search_path = public
+        AS $$
+        BEGIN
+            INSERT INTO public.profiles (id, username, avatar_url, bio, settings, created_at, updated_at)
+            VALUES (
+                NEW.id,
+                COALESCE(
+                    NEW.raw_user_meta_data ->> 'username',
+                    split_part(NEW.email, '@', 1),
+                    'user_' || substring(NEW.id::text, 1, 8)
+                ),
+                NULL,
+                NULL,
+                '{}'::jsonb,
+                now(),
+                now()
+            )
+            ON CONFLICT (id) DO NOTHING;
+
+            INSERT INTO public.user_agents (id, user_id, llm_id, agent_type, config, status, created_by, updated_by)
+            SELECT 
+                gen_random_uuid(),
+                NEW.id,
+                uac.llm_id,
+                uac.agent_type,
+                uac.config,
+                uac.status,
+                NEW.id,
+                NEW.id
+            FROM public.user_agent_catalog uac;
+            
+            RETURN NEW;
+        END;
+        $$
+    """)
+
+    op.execute("""
+        CREATE TRIGGER on_auth_user_created
+            AFTER INSERT ON auth.users
+            FOR EACH ROW EXECUTE FUNCTION public.create_profile_for_new_user()
+    """)
+
+    # 2. Revert rename: system_agents -> user_agent_catalog
+    _drop_rls("system_agents")
+
+    op.rename_table("system_agents", "user_agent_catalog")
+
+    op.execute(
+        "ALTER TABLE user_agent_catalog RENAME CONSTRAINT fk_system_agents_llm_id "
+        "TO fk_user_agent_catalog_llm_id"
+    )
+    op.execute(
+        "ALTER TABLE user_agent_catalog RENAME CONSTRAINT chk_system_agents_status "
+        "TO chk_user_agent_catalog_status"
+    )
+
+    _enable_rls("user_agent_catalog")
+
+    # 3. Recreate user_agents table
+    op.create_table(
+        "user_agents",
+        sa.Column("id", sa.UUID(), nullable=False),
+        sa.Column("user_id", sa.UUID(), nullable=False),
+        sa.Column("llm_id", sa.UUID(), nullable=False),
+        sa.Column("agent_type", sa.String(length=20), nullable=False),
+        sa.Column(
+            "config",
+            postgresql.JSONB(astext_type=sa.Text()),
+            server_default="{}",
+            nullable=False,
+        ),
+        sa.Column("status", sa.String(length=20), nullable=False),
+        sa.Column("created_by", sa.UUID(), nullable=True),
+        sa.Column("updated_by", sa.UUID(), nullable=True),
+        sa.Column(
+            "created_at",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("now()"),
+            nullable=False,
+        ),
+        sa.Column(
+            "updated_at",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("now()"),
+            nullable=False,
+        ),
+        sa.Column("deleted_at", sa.DateTime(timezone=True), nullable=True),
+        sa.PrimaryKeyConstraint("id"),
+    )
+
+    op.create_unique_constraint(
+        "uq_user_agents_user_id_agent_type", "user_agents", ["user_id", "agent_type"]
+    )
+
+    op.execute("CREATE INDEX ix_user_agents_agent_type ON user_agents (agent_type)")
+    op.execute("CREATE INDEX ix_user_agents_status ON user_agents (status)")
+
+    op.execute(
+        "ALTER TABLE user_agents ADD CONSTRAINT chk_agent_type "
+        "CHECK (agent_type IN ('INTENT_RECOGNITION', 'TASK_EXECUTION', 'RESULT_REPORTING'))"
+    )
+
+    op.create_foreign_key(
+        "fk_user_agents_user_id",
+        "user_agents",
+        "users",
+        ["user_id"],
+        ["id"],
+        referent_schema="auth",
+        ondelete="CASCADE",
+    )
+    op.create_foreign_key(
+        "fk_user_agents_llm_id",
+        "user_agents",
+        "llms",
+        ["llm_id"],
+        ["id"],
+        ondelete="RESTRICT",
+    )
+    op.create_foreign_key(
+        "fk_user_agents_created_by",
+        "user_agents",
+        "users",
+        ["created_by"],
+        ["id"],
+        referent_schema="auth",
+        ondelete="SET NULL",
+    )
+    op.create_foreign_key(
+        "fk_user_agents_updated_by",
+        "user_agents",
+        "users",
+        ["updated_by"],
+        ["id"],
+        referent_schema="auth",
+        ondelete="SET NULL",
+    )
+
+    _enable_rls("user_agents")
+
+    # 4. Recreate memories.agent_id
+    op.add_column("memories", sa.Column("agent_id", sa.UUID(), nullable=True))
+
+    op.create_foreign_key(
+        "fk_memories_agent_id",
+        "memories",
+        "user_agents",
+        ["agent_id"],
+        ["id"],
+        ondelete="CASCADE",
+    )
+
+    op.execute(
+        "CREATE INDEX ix_memories_agent_type_status ON memories (agent_id, memory_type, status)"
+    )
+
+    op.execute(
+        "ALTER TABLE memories ADD CONSTRAINT chk_memory_type_agent_id "
+        "CHECK ((memory_type = 'work' AND agent_id IS NOT NULL) OR "
+        "(memory_type = 'user' AND agent_id IS NULL))"
+    )
+
+
+def _enable_rls(table_name: str) -> None:
+    for role in ["anon", "authenticated"]:
+        for action in ["select", "insert", "update", "delete"]:
+            op.execute(
+                f"DROP POLICY IF EXISTS {role}_{action}_{table_name} ON {table_name}"
+            )
+    op.execute(f"ALTER TABLE {table_name} ENABLE ROW LEVEL SECURITY")
+    for role in ["anon", "authenticated"]:
+        op.execute(
+            f"CREATE POLICY {role}_select_{table_name} ON {table_name} "
+            f"FOR SELECT TO {role} USING (false)"
+        )
+        op.execute(
+            f"CREATE POLICY {role}_insert_{table_name} ON {table_name} "
+            f"FOR INSERT TO {role} WITH CHECK (false)"
+        )
+        op.execute(
+            f"CREATE POLICY {role}_update_{table_name} ON {table_name} "
+            f"FOR UPDATE TO {role} USING (false) WITH CHECK (false)"
+        )
+        op.execute(
+            f"CREATE POLICY {role}_delete_{table_name} ON {table_name} "
+            f"FOR DELETE TO {role} USING (false)"
+        )
+
+
+def _drop_rls(table_name: str) -> None:
+    for role in ["anon", "authenticated"]:
+        op.execute(f"DROP POLICY IF EXISTS {role}_delete_{table_name} ON {table_name}")
+        op.execute(f"DROP POLICY IF EXISTS {role}_update_{table_name} ON {table_name}")
+        op.execute(f"DROP POLICY IF EXISTS {role}_insert_{table_name} ON {table_name}")
+        op.execute(f"DROP POLICY IF EXISTS {role}_select_{table_name} ON {table_name}")
+    op.execute(f"ALTER TABLE {table_name} DISABLE ROW LEVEL SECURITY")
@@ -13,10 +13,9 @@ from models.memories import Memory
 from models.profile import Profile
 from models.schedule_items import ScheduleItem
 from models.schedule_subscriptions import ScheduleSubscription
+from models.system_agents import SystemAgents
 from models.todos import Todo
 from models.todo_sources import TodoSource
-from models.user_agents import UserAgent
-from models.user_agent_catalog import UserAgentCatalog

 __all__ = [
    "AgentChatMessage",
@@ -33,8 +32,7 @@ __all__ = [
    "Profile",
    "ScheduleItem",
    "ScheduleSubscription",
+    "SystemAgents",
    "Todo",
    "TodoSource",
-    "UserAgent",
-    "UserAgentCatalog",
 ]
@@ -9,8 +9,8 @@ from sqlalchemy.orm import Mapped, mapped_column
 from core.db.base import Base, TimestampMixin


-class UserAgentCatalog(TimestampMixin, Base):
-    __tablename__: str = "user_agent_catalog"
+class SystemAgents(TimestampMixin, Base):
+    __tablename__: str = "system_agents"

    agent_type: Mapped[str] = mapped_column(
        String(20),
@@ -1,62 +0,0 @@
-from __future__ import annotations
-
-import uuid
-from enum import Enum
-
-from sqlalchemy import String
-from sqlalchemy.dialects.postgresql import JSONB, UUID
-from sqlalchemy.orm import Mapped, mapped_column
-
-from core.db.base import Base, SoftDeleteMixin, TimestampMixin
-
-
-class UserAgentStatus(str, Enum):
-    ACTIVE = "active"
-    PAUSED = "paused"
-    MIGRATING = "migrating"
-
-
-class AgentType(str, Enum):
-    INTENT_RECOGNITION = "INTENT_RECOGNITION"
-    TASK_EXECUTION = "TASK_EXECUTION"
-    RESULT_REPORTING = "RESULT_REPORTING"
-
-
-class UserAgent(TimestampMixin, SoftDeleteMixin, Base):
-    __tablename__: str = "user_agents"
-    __table_args__ = {"extend_existing": True}
-
-    id: Mapped[uuid.UUID] = mapped_column(
-        UUID(as_uuid=True), primary_key=True, default=uuid.uuid4
-    )
-    user_id: Mapped[uuid.UUID] = mapped_column(
-        UUID(as_uuid=True),
-        nullable=False,
-        unique=True,
-    )
-    llm_id: Mapped[uuid.UUID] = mapped_column(
-        UUID(as_uuid=True),
-        nullable=False,
-    )
-    agent_type: Mapped[AgentType] = mapped_column(
-        String(20),
-        nullable=False,
-    )
-    config: Mapped[dict] = mapped_column(
-        JSONB,
-        nullable=False,
-        server_default="{}",
-    )
-    status: Mapped[UserAgentStatus] = mapped_column(
-        String(20),
-        nullable=False,
-        default=UserAgentStatus.ACTIVE,
-    )
-    created_by: Mapped[uuid.UUID | None] = mapped_column(
-        UUID(as_uuid=True),
-        nullable=True,
-    )
-    updated_by: Mapped[uuid.UUID | None] = mapped_column(
-        UUID(as_uuid=True),
-        nullable=True,
-    )