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
This commit is contained in:
@@ -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.profile import Profile
|
||||||
from models.schedule_items import ScheduleItem
|
from models.schedule_items import ScheduleItem
|
||||||
from models.schedule_subscriptions import ScheduleSubscription
|
from models.schedule_subscriptions import ScheduleSubscription
|
||||||
|
from models.system_agents import SystemAgents
|
||||||
from models.todos import Todo
|
from models.todos import Todo
|
||||||
from models.todo_sources import TodoSource
|
from models.todo_sources import TodoSource
|
||||||
from models.user_agents import UserAgent
|
|
||||||
from models.user_agent_catalog import UserAgentCatalog
|
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"AgentChatMessage",
|
"AgentChatMessage",
|
||||||
@@ -33,8 +32,7 @@ __all__ = [
|
|||||||
"Profile",
|
"Profile",
|
||||||
"ScheduleItem",
|
"ScheduleItem",
|
||||||
"ScheduleSubscription",
|
"ScheduleSubscription",
|
||||||
|
"SystemAgents",
|
||||||
"Todo",
|
"Todo",
|
||||||
"TodoSource",
|
"TodoSource",
|
||||||
"UserAgent",
|
|
||||||
"UserAgentCatalog",
|
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -9,8 +9,8 @@ from sqlalchemy.orm import Mapped, mapped_column
|
|||||||
from core.db.base import Base, TimestampMixin
|
from core.db.base import Base, TimestampMixin
|
||||||
|
|
||||||
|
|
||||||
class UserAgentCatalog(TimestampMixin, Base):
|
class SystemAgents(TimestampMixin, Base):
|
||||||
__tablename__: str = "user_agent_catalog"
|
__tablename__: str = "system_agents"
|
||||||
|
|
||||||
agent_type: Mapped[str] = mapped_column(
|
agent_type: Mapped[str] = mapped_column(
|
||||||
String(20),
|
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,
|
|
||||||
)
|
|
||||||
Reference in New Issue
Block a user