From e700ebd3fbd2f3e8056cb26c24d99a1c674f60f5 Mon Sep 17 00:00:00 2001 From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> Date: Mon, 9 Jun 2025 15:36:51 -0400 Subject: [PATCH] Initial Model Changes --- .../d429de0e5338_adopt_peer_paradigm.py | 209 ++++++++++++++++++ src/models.py | 160 +++++--------- 2 files changed, 259 insertions(+), 110 deletions(-) create mode 100644 migrations/versions/d429de0e5338_adopt_peer_paradigm.py diff --git a/migrations/versions/d429de0e5338_adopt_peer_paradigm.py b/migrations/versions/d429de0e5338_adopt_peer_paradigm.py new file mode 100644 index 00000000..6125eb4b --- /dev/null +++ b/migrations/versions/d429de0e5338_adopt_peer_paradigm.py @@ -0,0 +1,209 @@ +"""adopt peer paradigm + +Revision ID: d429de0e5338 +Revises: 66e63cf2cf77 +Create Date: 2025-06-09 15:16:38.164067 + +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op +from nanoid import generate as generate_nanoid +from sqlalchemy.dialects import postgresql + +from migrations.utils import column_exists, fk_exists, index_exists + +# revision identifiers, used by Alembic. +revision: str = 'd429de0e5338' +down_revision: Union[str, None] = '66e63cf2cf77' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + inspector = sa.inspect(op.get_bind()) + + # Rename tables if they exist + if inspector.has_table("apps"): + op.rename_table('apps', 'workspaces') + if inspector.has_table("users"): + op.rename_table('users', 'peers') + + # Alter columns in renamed tables + if column_exists('peers', 'app_id', inspector): + op.alter_column('peers', 'app_id', new_column_name='workspace_id') + + # Add workspace_id to sessions table + if not column_exists('sessions', 'workspace_id', inspector): + op.add_column('sessions', sa.Column('workspace_id', sa.TEXT(), nullable=True)) + + # Data Migration: Populate sessions.workspace_id + if column_exists('sessions', 'app_id', inspector): + op.execute("UPDATE sessions SET workspace_id = app_id") + + op.alter_column('sessions', 'workspace_id', existing_type=sa.TEXT(), nullable=False) + + # Create session_peers table for many-to-many relationship + if not inspector.has_table('session_peers'): + op.create_table('session_peers', + sa.Column('session_public_id', sa.TEXT(), nullable=False), + sa.Column('peer_public_id', sa.TEXT(), nullable=False), + sa.ForeignKeyConstraint(['peer_public_id'], ['peers.public_id'], ), + sa.ForeignKeyConstraint(['session_public_id'], ['sessions.public_id'], ), + sa.PrimaryKeyConstraint('session_public_id', 'peer_public_id') + ) + + # --- Data Migration --- + conn = op.get_bind() + + # Create an 'agent' peer for each workspace + workspaces = conn.execute(sa.text("SELECT public_id FROM workspaces")).fetchall() + agent_peers_map = {} + for workspace_data in workspaces: + workspace_id = workspace_data[0] + agent_peer_public_id = generate_nanoid() + agent_peers_map[workspace_id] = agent_peer_public_id + op.execute( + sa.text( + "INSERT INTO peers (public_id, name, workspace_id, created_at, metadata) VALUES (:public_id, 'agent', :workspace_id, NOW(), '{}'::jsonb)" + ).bindparams(public_id=agent_peer_public_id, workspace_id=workspace_id) + ) + + # Fetch existing sessions to migrate relationships + if column_exists('sessions', 'user_id', inspector) and column_exists('sessions', 'workspace_id', inspector): + sessions = conn.execute(sa.text("SELECT public_id, user_id, workspace_id FROM sessions")).fetchall() + + # Populate session_peers with both the original user and the new agent + for session_data in sessions: + session_id, user_id, workspace_id = session_data + if user_id: + op.execute(sa.text("INSERT INTO session_peers (session_public_id, peer_public_id) VALUES (:sid, :pid)").bindparams(sid=session_id, pid=user_id)) + agent_peer_id = agent_peers_map.get(workspace_id) + if agent_peer_id: + op.execute(sa.text("INSERT INTO session_peers (session_public_id, peer_public_id) VALUES (:sid, :pid)").bindparams(sid=session_id, pid=agent_peer_id)) + + # Add and populate the new sender_id column in messages + if not column_exists('messages', 'sender_id', inspector): + op.add_column('messages', sa.Column('sender_id', sa.TEXT(), nullable=True)) + + if column_exists('messages', 'is_user', inspector): + # Set sender_id for user messages + op.execute(""" + UPDATE messages + SET sender_id = s.user_id + FROM sessions s + WHERE messages.session_id = s.public_id AND messages.is_user = TRUE + """) + # Set sender_id for agent messages + op.execute(""" + UPDATE messages m + SET sender_id = p.public_id + FROM sessions s + JOIN peers p ON s.workspace_id = p.workspace_id + WHERE m.session_id = s.public_id + AND p.name = 'agent' + AND m.is_user = FALSE + """) + + # Finalize schema: add constraints, drop old columns + op.alter_column('messages', 'sender_id', existing_type=sa.TEXT(), nullable=False) + op.alter_column('messages', 'session_id', existing_type=sa.TEXT(), nullable=True) + + if not fk_exists('messages', 'messages_sender_id_fkey', inspector): + op.create_foreign_key('messages_sender_id_fkey', 'messages', 'peers', ['sender_id'], ['public_id']) + + if column_exists('messages', 'is_user', inspector): + op.drop_column('messages', 'is_user') + + if index_exists('messages', 'idx_messages_session_lookup', inspector): + op.drop_index('idx_messages_session_lookup', table_name='messages') + op.create_index('idx_messages_session_lookup', 'messages', ['session_id', 'id'], unique=False, postgresql_include=['public_id', 'created_at']) + + # Clean up sessions table + if fk_exists('sessions', 'sessions_user_id_fkey', inspector): + op.drop_constraint('sessions_user_id_fkey', 'sessions', type_='foreignkey') + if fk_exists('sessions', 'sessions_app_id_fkey', inspector): + op.drop_constraint('sessions_app_id_fkey', 'sessions', type_='foreignkey') + if column_exists('sessions', 'user_id', inspector): + op.drop_column('sessions', 'user_id') + if column_exists('sessions', 'app_id', inspector): + op.drop_column('sessions', 'app_id') + + if not fk_exists('sessions', 'sessions_workspace_id_fkey', inspector): + op.create_foreign_key('sessions_workspace_id_fkey', 'sessions', 'workspaces', ['workspace_id'], ['public_id']) + + # Drop metamessages table + if inspector.has_table('metamessages'): + op.drop_table('metamessages') + + # Update other tables + if column_exists('collections', 'user_id', inspector): + op.alter_column('collections', 'user_id', new_column_name='peer_id') + if column_exists('collections', 'app_id', inspector): + op.alter_column('collections', 'app_id', new_column_name='workspace_id') + + if column_exists('documents', 'user_id', inspector): + op.alter_column('documents', 'user_id', new_column_name='peer_id') + if column_exists('documents', 'app_id', inspector): + op.alter_column('documents', 'app_id', new_column_name='workspace_id') + + +def downgrade() -> None: + # Note: This downgrade path only reverts schema changes and does not restore data. + inspector = sa.inspect(op.get_bind()) + + if column_exists('documents', 'peer_id', inspector): + op.alter_column('documents', 'peer_id', new_column_name='user_id') + if column_exists('documents', 'workspace_id', inspector): + op.alter_column('documents', 'workspace_id', new_column_name='app_id') + + if column_exists('collections', 'peer_id', inspector): + op.alter_column('collections', 'peer_id', new_column_name='user_id') + if column_exists('collections', 'workspace_id', inspector): + op.alter_column('collections', 'workspace_id', new_column_name='app_id') + + if not inspector.has_table('metamessages'): + op.create_table('metamessages', + sa.Column('id', sa.BIGINT(), autoincrement=True, nullable=False), + sa.Column('public_id', sa.TEXT(), autoincrement=False, nullable=False), + sa.Column('session_id', sa.TEXT(), autoincrement=False, nullable=False), + sa.Column('content', sa.TEXT(), autoincrement=False, nullable=False), + sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()), autoincrement=False, nullable=False), + sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), autoincrement=False, nullable=False), + sa.ForeignKeyConstraint(['session_id'], ['sessions.public_id'], name='metamessages_session_id_fkey'), + sa.PrimaryKeyConstraint('id', name='metamessages_pkey') + ) + + if fk_exists('sessions', 'sessions_workspace_id_fkey', inspector): + op.drop_constraint('sessions_workspace_id_fkey', 'sessions', type_='foreignkey') + + if not column_exists('sessions', 'app_id', inspector): + op.add_column('sessions', sa.Column('app_id', sa.TEXT(), autoincrement=False, nullable=True)) + if column_exists('sessions', 'workspace_id', inspector): + op.execute("UPDATE sessions SET app_id = workspace_id") + op.drop_column('sessions', 'workspace_id') + + if not column_exists('sessions', 'user_id', inspector): + op.add_column('sessions', sa.Column('user_id', sa.TEXT(), autoincrement=False, nullable=True)) + + if not column_exists('messages', 'is_user', inspector): + op.add_column('messages', sa.Column('is_user', sa.BOOLEAN(), autoincrement=False, nullable=True)) + + if fk_exists('messages', 'messages_sender_id_fkey', inspector): + op.drop_constraint('messages_sender_id_fkey', 'messages', type_='foreignkey') + + op.alter_column('messages', 'session_id', existing_type=sa.TEXT(), nullable=False) + if column_exists('messages', 'sender_id', inspector): + op.drop_column('messages', 'sender_id') + + if inspector.has_table('session_peers'): + op.drop_table('session_peers') + + if inspector.has_table("peers"): + op.rename_table('peers', 'users') + if inspector.has_table("workspaces"): + op.rename_table('workspaces', 'apps') + + if column_exists('users', 'workspace_id', inspector): + op.alter_column('users', 'workspace_id', new_column_name='app_id') diff --git a/src/models.py b/src/models.py index 71587653..69f0cc7f 100644 --- a/src/models.py +++ b/src/models.py @@ -12,7 +12,6 @@ from sqlalchemy import ( Identity, Index, UniqueConstraint, - text, ) from sqlalchemy.dialects.postgresql import JSONB, TEXT from sqlalchemy.orm import Mapped, mapped_column, relationship @@ -23,8 +22,8 @@ from .db import Base load_dotenv() -class App(Base): - __tablename__ = "apps" +class Workspace(Base): + __tablename__ = "workspaces" id: Mapped[int] = mapped_column( BigInteger, Identity(), primary_key=True, index=True, autoincrement=True ) @@ -32,7 +31,7 @@ class App(Base): TEXT, index=True, unique=True, default=generate_nanoid ) name: Mapped[str] = mapped_column(TEXT, index=True, unique=True) - users = relationship("User", back_populates="app") + peers = relationship("Peer", back_populates="workspace") created_at: Mapped[datetime.datetime] = mapped_column( DateTime(timezone=True), index=True, default=func.now() ) @@ -45,8 +44,18 @@ class App(Base): ) -class User(Base): - __tablename__ = "users" +class SessionPeer(Base): + __tablename__ = "session_peers" + session_public_id: Mapped[str] = mapped_column( + TEXT, ForeignKey("sessions.public_id"), primary_key=True + ) + peer_public_id: Mapped[str] = mapped_column( + TEXT, ForeignKey("peers.public_id"), primary_key=True + ) + + +class Peer(Base): + __tablename__ = "peers" id: Mapped[int] = mapped_column( BigInteger, Identity(), primary_key=True, index=True, autoincrement=True ) @@ -58,22 +67,23 @@ class User(Base): created_at: Mapped[datetime.datetime] = mapped_column( DateTime(timezone=True), index=True, default=func.now() ) - app_id: Mapped[str] = mapped_column(ForeignKey("apps.public_id"), index=True) - app = relationship("App", back_populates="users") - sessions = relationship("Session", back_populates="user") - collections = relationship("Collection", back_populates="user") - metamessages = relationship("Metamessage", back_populates="user") + workspace_id: Mapped[str] = mapped_column( + ForeignKey("workspaces.public_id"), index=True + ) + workspace = relationship("Workspace", back_populates="peers") + sessions = relationship("Session", secondary="session_peers", back_populates="peers") + collections = relationship("Collection", back_populates="peer") __table_args__ = ( - UniqueConstraint("name", "app_id", name="unique_name_app_user"), + UniqueConstraint("name", "workspace_id", name="unique_name_workspace_peer"), CheckConstraint("length(public_id) = 21", name="public_id_length"), CheckConstraint("length(name) <= 512", name="name_length"), CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"), - Index("idx_users_app_lookup", "app_id", "public_id"), + Index("idx_peers_workspace_lookup", "workspace_id", "public_id"), ) def __repr__(self) -> str: - return f"User(id={self.id}, app_id={self.app_id}, public_id={self.public_id} created_at={self.created_at}, h_metadata={self.h_metadata})" + return f"Peer(id={self.id}, workspace_id={self.workspace_id}, public_id={self.public_id} created_at={self.created_at}, h_metadata={self.h_metadata})" class Session(Base): @@ -90,19 +100,18 @@ class Session(Base): DateTime(timezone=True), index=True, default=func.now() ) messages = relationship("Message", back_populates="session") - metamessages = relationship("Metamessage", back_populates="session") - user_id: Mapped[str] = mapped_column(ForeignKey("users.public_id"), index=True) - app_id: Mapped[str] = mapped_column(ForeignKey("apps.public_id"), index=True) - user = relationship("User", back_populates="sessions") + workspace_id: Mapped[str] = mapped_column( + ForeignKey("workspaces.public_id"), index=True + ) + peers = relationship("Peer", secondary="session_peers", back_populates="sessions") __table_args__ = ( CheckConstraint("length(public_id) = 21", name="public_id_length"), CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"), - Index("idx_sessions_user_lookup", "user_id", "public_id"), ) def __repr__(self) -> str: - return f"Session(id={self.id}, user_id={self.user_id}, is_active={self.is_active}, created_at={self.created_at}, h_metadata={self.h_metadata})" + return f"Session(id={self.id}, is_active={self.is_active}, created_at={self.created_at}, h_metadata={self.h_metadata})" class Message(Base): @@ -113,10 +122,9 @@ class Message(Base): public_id: Mapped[str] = mapped_column( TEXT, index=True, unique=True, default=generate_nanoid ) - session_id: Mapped[str] = mapped_column( + session_id: Mapped[str | None] = mapped_column( ForeignKey("sessions.public_id"), index=True ) - is_user: Mapped[bool] content: Mapped[str] = mapped_column(TEXT) h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default={}) @@ -124,9 +132,10 @@ class Message(Base): DateTime(timezone=True), index=True, default=func.now() ) session = relationship("Session", back_populates="messages") - metamessages = relationship("Metamessage", back_populates="message") - user_id: Mapped[str] = mapped_column(ForeignKey("users.public_id"), index=True) - app_id: Mapped[str] = mapped_column(ForeignKey("apps.public_id"), index=True) + sender_id: Mapped[str] = mapped_column(ForeignKey("peers.public_id"), index=True) + workspace_id: Mapped[str] = mapped_column( + ForeignKey("workspaces.public_id"), index=True + ) __table_args__ = ( CheckConstraint("length(public_id) = 21", name="public_id_length"), @@ -136,85 +145,12 @@ class Message(Base): "idx_messages_session_lookup", "session_id", "id", - postgresql_include=["public_id", "is_user", "created_at"], + postgresql_include=["public_id", "created_at"], ), ) def __repr__(self) -> str: - return f"Message(id={self.id}, session_id={self.session_id}, is_user={self.is_user}, content={self.content[10:]})" - - -class Metamessage(Base): - __tablename__ = "metamessages" - id: Mapped[int] = mapped_column( - BigInteger, Identity(), primary_key=True, index=True, autoincrement=True - ) - public_id: Mapped[str] = mapped_column( - TEXT, index=True, unique=True, default=generate_nanoid - ) - label: Mapped[str] = mapped_column(TEXT, index=True) - content: Mapped[str] = mapped_column(TEXT) - - # Foreign keys - message_id is now optional - user_id: Mapped[str] = mapped_column(ForeignKey("users.public_id"), index=True) - app_id: Mapped[str] = mapped_column(ForeignKey("apps.public_id"), index=True) - session_id: Mapped[str | None] = mapped_column( - ForeignKey("sessions.public_id"), index=True, nullable=True - ) - message_id: Mapped[str | None] = mapped_column( - ForeignKey("messages.public_id"), index=True, nullable=True - ) - - # Relationships - user = relationship("User", back_populates="metamessages") - session = relationship("Session", back_populates="metamessages") - message = relationship("Message", back_populates="metamessages") - - created_at: Mapped[datetime.datetime] = mapped_column( - DateTime(timezone=True), index=True, default=func.now() - ) - h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default={}) - - __table_args__ = ( - CheckConstraint("length(public_id) = 21", name="public_id_length"), - CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"), - CheckConstraint("length(content) <= 65535", name="content_length"), - CheckConstraint("length(label) <= 512", name="label_length"), - # Added constraints to ensure consistency - CheckConstraint( - "(message_id IS NULL) OR (session_id IS NOT NULL)", - name="message_requires_session", - ), - # Keep existing index - Index( - "idx_metamessages_lookup", - "label", - text("id DESC"), - postgresql_include=["public_id", "message_id", "created_at"], - ), - # Indices for user, session, and message lookups - Index( - "idx_metamessages_user_lookup", - "user_id", - "label", - text("id DESC"), - ), - Index( - "idx_metamessages_session_lookup", - "session_id", - "label", - text("id DESC"), - ), - Index( - "idx_metamessages_message_lookup", - "message_id", - "label", - text("id DESC"), - ), - ) - - def __repr__(self) -> str: - return f"Metamessages(id={self.id}, user_id={self.user_id}, session_id={self.session_id}, message_id={self.message_id}, label={self.label})" + return f"Message(id={self.id}, session_id={self.session_id}, content={self.content[10:]})" class Collection(Base): @@ -234,14 +170,16 @@ class Collection(Base): documents = relationship( "Document", back_populates="collection", cascade="all, delete, delete-orphan" ) - user = relationship("User", back_populates="collections") - user_id: Mapped[str] = mapped_column( - TEXT, ForeignKey("users.public_id"), index=True + peer = relationship("Peer", back_populates="collections") + peer_id: Mapped[str] = mapped_column( + TEXT, ForeignKey("peers.public_id"), index=True + ) + workspace_id: Mapped[str] = mapped_column( + ForeignKey("workspaces.public_id"), index=True ) - app_id: Mapped[str] = mapped_column(ForeignKey("apps.public_id"), index=True) __table_args__ = ( - UniqueConstraint("name", "user_id", name="unique_name_collection_user"), + UniqueConstraint("name", "peer_id", name="unique_name_collection_peer"), CheckConstraint("length(public_id) = 21", name="public_id_length"), CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"), CheckConstraint("length(name) <= 512", name="name_length"), @@ -266,8 +204,10 @@ class Document(Base): collection_id: Mapped[str] = mapped_column( TEXT, ForeignKey("collections.public_id"), index=True ) - user_id: Mapped[str] = mapped_column(ForeignKey("users.public_id"), index=True) - app_id: Mapped[str] = mapped_column(ForeignKey("apps.public_id"), index=True) + peer_id: Mapped[str] = mapped_column(ForeignKey("peers.public_id"), index=True) + workspace_id: Mapped[str] = mapped_column( + ForeignKey("workspaces.public_id"), index=True + ) collection = relationship("Collection", back_populates="documents") __table_args__ = ( @@ -278,9 +218,9 @@ class Document(Base): Index( "idx_documents_embedding_hnsw", "embedding", - postgresql_using="hnsw", # HNSW index type - postgresql_with={"m": 16, "ef_construction": 64}, # HNSW parameters - postgresql_ops={"embedding": "vector_cosine_ops"}, # Cosine distance operator + postgresql_using="hnsw", # HNSW index type + postgresql_with={"m": 16, "ef_construction": 64}, # HNSW parameters + postgresql_ops={"embedding": "vector_cosine_ops"}, # Cosine distance operator ), )