From 7879ba6a773f783c9483563ab74254dd762ff392 Mon Sep 17 00:00:00 2001 From: Evan Lohn Date: Wed, 29 Jan 2025 17:06:41 -0800 Subject: [PATCH] collapsed db migrations post-rebase --- .../1adf5ea20d2b_agent_doc_result_col.py | 29 ------- ...211b14ab5a91_refined_answer_improvement.py | 31 ------- ...925b58bd75b6_agent_metric_col_rename__s.py | 35 -------- ...e58_agent_metric_table_renames__agent__.py | 25 ------ .../versions/98a5008d8711_agent_tracking.py | 42 ---------- ...eb76d618ec_agent_table_renames__agent__.py | 84 ------------------- ...18a25b_agent_table_changes_rename_level.py | 40 --------- ...ed_create_pro_search_persistence_tables.py | 68 --------------- .../onyx/agents/agent_search/db_operations.py | 66 --------------- backend/onyx/db/chat.py | 4 +- backend/onyx/db/models.py | 2 +- 11 files changed, 3 insertions(+), 423 deletions(-) delete mode 100644 backend/alembic/versions/1adf5ea20d2b_agent_doc_result_col.py delete mode 100644 backend/alembic/versions/211b14ab5a91_refined_answer_improvement.py delete mode 100644 backend/alembic/versions/925b58bd75b6_agent_metric_col_rename__s.py delete mode 100644 backend/alembic/versions/9787be927e58_agent_metric_table_renames__agent__.py delete mode 100644 backend/alembic/versions/98a5008d8711_agent_tracking.py delete mode 100644 backend/alembic/versions/bceb76d618ec_agent_table_renames__agent__.py delete mode 100644 backend/alembic/versions/c0132518a25b_agent_table_changes_rename_level.py delete mode 100644 backend/alembic/versions/e9cf2bd7baed_create_pro_search_persistence_tables.py delete mode 100644 backend/onyx/agents/agent_search/db_operations.py diff --git a/backend/alembic/versions/1adf5ea20d2b_agent_doc_result_col.py b/backend/alembic/versions/1adf5ea20d2b_agent_doc_result_col.py deleted file mode 100644 index 62db727f9..000000000 --- a/backend/alembic/versions/1adf5ea20d2b_agent_doc_result_col.py +++ /dev/null @@ -1,29 +0,0 @@ -"""agent_doc_result_col - -Revision ID: 1adf5ea20d2b -Revises: e9cf2bd7baed -Create Date: 2025-01-05 13:14:58.344316 - -""" -from alembic import op -import sqlalchemy as sa -from sqlalchemy.dialects import postgresql - -# revision identifiers, used by Alembic. -revision = "1adf5ea20d2b" -down_revision = "e9cf2bd7baed" -branch_labels = None -depends_on = None - - -def upgrade() -> None: - # Add the new column with JSONB type - op.add_column( - "sub_question", - sa.Column("sub_question_doc_results", postgresql.JSONB(), nullable=True), - ) - - -def downgrade() -> None: - # Drop the column - op.drop_column("sub_question", "sub_question_doc_results") diff --git a/backend/alembic/versions/211b14ab5a91_refined_answer_improvement.py b/backend/alembic/versions/211b14ab5a91_refined_answer_improvement.py deleted file mode 100644 index 5009d0f3b..000000000 --- a/backend/alembic/versions/211b14ab5a91_refined_answer_improvement.py +++ /dev/null @@ -1,31 +0,0 @@ -"""refined answer improvement - -Revision ID: 211b14ab5a91 -Revises: 925b58bd75b6 -Create Date: 2025-01-24 14:05:03.334309 - -""" -from alembic import op -import sqlalchemy as sa - - -# revision identifiers, used by Alembic. -revision = "211b14ab5a91" -down_revision = "925b58bd75b6" -branch_labels = None -depends_on = None - - -def upgrade() -> None: - op.add_column( - "chat_message", - sa.Column( - "refined_answer_improvement", - sa.Boolean(), - nullable=True, - ), - ) - - -def downgrade() -> None: - op.drop_column("chat_message", "refined_answer_improvement") diff --git a/backend/alembic/versions/925b58bd75b6_agent_metric_col_rename__s.py b/backend/alembic/versions/925b58bd75b6_agent_metric_col_rename__s.py deleted file mode 100644 index 6bf501608..000000000 --- a/backend/alembic/versions/925b58bd75b6_agent_metric_col_rename__s.py +++ /dev/null @@ -1,35 +0,0 @@ -"""agent_metric_col_rename__s - -Revision ID: 925b58bd75b6 -Revises: 9787be927e58 -Create Date: 2025-01-06 11:20:26.752441 - -""" -from alembic import op - - -# revision identifiers, used by Alembic. -revision = "925b58bd75b6" -down_revision = "9787be927e58" -branch_labels = None -depends_on = None - - -def upgrade() -> None: - # Rename columns using PostgreSQL syntax - op.alter_column( - "agent__search_metrics", "base_duration_s", new_column_name="base_duration__s" - ) - op.alter_column( - "agent__search_metrics", "full_duration_s", new_column_name="full_duration__s" - ) - - -def downgrade() -> None: - # Revert the column renames - op.alter_column( - "agent__search_metrics", "base_duration__s", new_column_name="base_duration_s" - ) - op.alter_column( - "agent__search_metrics", "full_duration__s", new_column_name="full_duration_s" - ) diff --git a/backend/alembic/versions/9787be927e58_agent_metric_table_renames__agent__.py b/backend/alembic/versions/9787be927e58_agent_metric_table_renames__agent__.py deleted file mode 100644 index 2b605f5b3..000000000 --- a/backend/alembic/versions/9787be927e58_agent_metric_table_renames__agent__.py +++ /dev/null @@ -1,25 +0,0 @@ -"""agent_metric_table_renames__agent__ - -Revision ID: 9787be927e58 -Revises: bceb76d618ec -Create Date: 2025-01-06 11:01:44.210160 - -""" -from alembic import op - - -# revision identifiers, used by Alembic. -revision = "9787be927e58" -down_revision = "bceb76d618ec" -branch_labels = None -depends_on = None - - -def upgrade() -> None: - # Rename table from agent_search_metrics to agent__search_metrics - op.rename_table("agent_search_metrics", "agent__search_metrics") - - -def downgrade() -> None: - # Rename table back from agent__search_metrics to agent_search_metrics - op.rename_table("agent__search_metrics", "agent_search_metrics") diff --git a/backend/alembic/versions/98a5008d8711_agent_tracking.py b/backend/alembic/versions/98a5008d8711_agent_tracking.py deleted file mode 100644 index 0fb4091cd..000000000 --- a/backend/alembic/versions/98a5008d8711_agent_tracking.py +++ /dev/null @@ -1,42 +0,0 @@ -"""agent_tracking - -Revision ID: 98a5008d8711 -Revises: 4d58345da04a -Create Date: 2025-01-04 14:41:52.732238 - -""" -from alembic import op -import sqlalchemy as sa -from sqlalchemy.dialects import postgresql - -# revision identifiers, used by Alembic. -revision = "98a5008d8711" -down_revision = "4d58345da04a" -branch_labels = None -depends_on = None - - -def upgrade() -> None: - op.create_table( - "agent_search_metrics", - sa.Column("id", sa.Integer(), nullable=False), - sa.Column("user_id", postgresql.UUID(as_uuid=True), nullable=True), - sa.Column("persona_id", sa.Integer(), nullable=True), - sa.Column("agent_type", sa.String(), nullable=False), - sa.Column("start_time", sa.DateTime(timezone=True), nullable=False), - sa.Column("base_duration_s", sa.Float(), nullable=False), - sa.Column("full_duration_s", sa.Float(), nullable=False), - sa.Column("base_metrics", postgresql.JSONB(), nullable=True), - sa.Column("refined_metrics", postgresql.JSONB(), nullable=True), - sa.Column("all_metrics", postgresql.JSONB(), nullable=True), - sa.ForeignKeyConstraint( - ["persona_id"], - ["persona.id"], - ), - sa.ForeignKeyConstraint(["user_id"], ["user.id"], ondelete="CASCADE"), - sa.PrimaryKeyConstraint("id"), - ) - - -def downgrade() -> None: - op.drop_table("agent_search_metrics") diff --git a/backend/alembic/versions/bceb76d618ec_agent_table_renames__agent__.py b/backend/alembic/versions/bceb76d618ec_agent_table_renames__agent__.py deleted file mode 100644 index 1c1cb2e0d..000000000 --- a/backend/alembic/versions/bceb76d618ec_agent_table_renames__agent__.py +++ /dev/null @@ -1,84 +0,0 @@ -"""agent_table_renames__agent__ - -Revision ID: bceb76d618ec -Revises: c0132518a25b -Create Date: 2025-01-06 10:50:48.109285 - -""" -from alembic import op - - -# revision identifiers, used by Alembic. -revision = "bceb76d618ec" -down_revision = "c0132518a25b" -branch_labels = None -depends_on = None - - -def upgrade() -> None: - op.drop_constraint( - "sub_query__search_doc_sub_query_id_fkey", - "sub_query__search_doc", - type_="foreignkey", - ) - op.drop_constraint( - "sub_query__search_doc_search_doc_id_fkey", - "sub_query__search_doc", - type_="foreignkey", - ) - # Rename tables - op.rename_table("sub_query", "agent__sub_query") - op.rename_table("sub_question", "agent__sub_question") - op.rename_table("sub_query__search_doc", "agent__sub_query__search_doc") - - # Update both foreign key constraints for agent__sub_query__search_doc - - # Create new foreign keys with updated names - op.create_foreign_key( - "agent__sub_query__search_doc_sub_query_id_fkey", - "agent__sub_query__search_doc", - "agent__sub_query", - ["sub_query_id"], - ["id"], - ) - op.create_foreign_key( - "agent__sub_query__search_doc_search_doc_id_fkey", - "agent__sub_query__search_doc", - "search_doc", # This table name doesn't change - ["search_doc_id"], - ["id"], - ) - - -def downgrade() -> None: - # Update foreign key constraints for sub_query__search_doc - op.drop_constraint( - "agent__sub_query__search_doc_sub_query_id_fkey", - "agent__sub_query__search_doc", - type_="foreignkey", - ) - op.drop_constraint( - "agent__sub_query__search_doc_search_doc_id_fkey", - "agent__sub_query__search_doc", - type_="foreignkey", - ) - - # Rename tables back - op.rename_table("agent__sub_query__search_doc", "sub_query__search_doc") - op.rename_table("agent__sub_question", "sub_question") - op.rename_table("agent__sub_query", "sub_query") - - op.create_foreign_key( - "sub_query__search_doc_sub_query_id_fkey", - "sub_query__search_doc", - "sub_query", - ["sub_query_id"], - ["id"], - ) - op.create_foreign_key( - "sub_query__search_doc_search_doc_id_fkey", - "sub_query__search_doc", - "search_doc", # This table name doesn't change - ["search_doc_id"], - ["id"], - ) diff --git a/backend/alembic/versions/c0132518a25b_agent_table_changes_rename_level.py b/backend/alembic/versions/c0132518a25b_agent_table_changes_rename_level.py deleted file mode 100644 index e84538099..000000000 --- a/backend/alembic/versions/c0132518a25b_agent_table_changes_rename_level.py +++ /dev/null @@ -1,40 +0,0 @@ -"""agent_table_changes_rename_level - -Revision ID: c0132518a25b -Revises: 1adf5ea20d2b -Create Date: 2025-01-05 16:38:37.660152 - -""" -from alembic import op -import sqlalchemy as sa - - -# revision identifiers, used by Alembic. -revision = "c0132518a25b" -down_revision = "1adf5ea20d2b" -branch_labels = None -depends_on = None - - -def upgrade() -> None: - # Add level and level_question_nr columns with NOT NULL constraint - op.add_column( - "sub_question", - sa.Column("level", sa.Integer(), nullable=False, server_default="0"), - ) - op.add_column( - "sub_question", - sa.Column( - "level_question_nr", sa.Integer(), nullable=False, server_default="0" - ), - ) - - # Remove the server_default after the columns are created - op.alter_column("sub_question", "level", server_default=None) - op.alter_column("sub_question", "level_question_nr", server_default=None) - - -def downgrade() -> None: - # Remove the columns - op.drop_column("sub_question", "level_question_nr") - op.drop_column("sub_question", "level") diff --git a/backend/alembic/versions/e9cf2bd7baed_create_pro_search_persistence_tables.py b/backend/alembic/versions/e9cf2bd7baed_create_pro_search_persistence_tables.py deleted file mode 100644 index 6ed9d783f..000000000 --- a/backend/alembic/versions/e9cf2bd7baed_create_pro_search_persistence_tables.py +++ /dev/null @@ -1,68 +0,0 @@ -"""create pro search persistence tables - -Revision ID: e9cf2bd7baed -Revises: 98a5008d8711 -Create Date: 2025-01-02 17:55:56.544246 - -""" -from alembic import op -import sqlalchemy as sa -from sqlalchemy.dialects.postgresql import UUID - - -# revision identifiers, used by Alembic. -revision = "e9cf2bd7baed" -down_revision = "98a5008d8711" -branch_labels = None -depends_on = None - - -def upgrade() -> None: - # Create sub_question table - op.create_table( - "sub_question", - sa.Column("id", sa.Integer, primary_key=True), - sa.Column("primary_question_id", sa.Integer, sa.ForeignKey("chat_message.id")), - sa.Column( - "chat_session_id", UUID(as_uuid=True), sa.ForeignKey("chat_session.id") - ), - sa.Column("sub_question", sa.Text), - sa.Column( - "time_created", sa.DateTime(timezone=True), server_default=sa.func.now() - ), - sa.Column("sub_answer", sa.Text), - ) - - # Create sub_query table - op.create_table( - "sub_query", - sa.Column("id", sa.Integer, primary_key=True), - sa.Column("parent_question_id", sa.Integer, sa.ForeignKey("sub_question.id")), - sa.Column( - "chat_session_id", UUID(as_uuid=True), sa.ForeignKey("chat_session.id") - ), - sa.Column("sub_query", sa.Text), - sa.Column( - "time_created", sa.DateTime(timezone=True), server_default=sa.func.now() - ), - ) - - # Create sub_query__search_doc association table - op.create_table( - "sub_query__search_doc", - sa.Column( - "sub_query_id", sa.Integer, sa.ForeignKey("sub_query.id"), primary_key=True - ), - sa.Column( - "search_doc_id", - sa.Integer, - sa.ForeignKey("search_doc.id"), - primary_key=True, - ), - ) - - -def downgrade() -> None: - op.drop_table("sub_query__search_doc") - op.drop_table("sub_query") - op.drop_table("sub_question") diff --git a/backend/onyx/agents/agent_search/db_operations.py b/backend/onyx/agents/agent_search/db_operations.py deleted file mode 100644 index 3df137b11..000000000 --- a/backend/onyx/agents/agent_search/db_operations.py +++ /dev/null @@ -1,66 +0,0 @@ -from uuid import UUID - -from sqlalchemy.orm import Session - -from onyx.db.models import AgentSubQuery -from onyx.db.models import AgentSubQuestion - - -def create_sub_question( - db_session: Session, - chat_session_id: UUID, - primary_message_id: int, - sub_question: str, - sub_answer: str, -) -> AgentSubQuestion: - """Create a new sub-question record in the database.""" - sub_q = AgentSubQuestion( - chat_session_id=chat_session_id, - primary_question_id=primary_message_id, - sub_question=sub_question, - sub_answer=sub_answer, - ) - db_session.add(sub_q) - db_session.flush() - return sub_q - - -def create_sub_query( - db_session: Session, - chat_session_id: UUID, - parent_question_id: int, - sub_query: str, -) -> AgentSubQuery: - """Create a new sub-query record in the database.""" - sub_q = AgentSubQuery( - chat_session_id=chat_session_id, - parent_question_id=parent_question_id, - sub_query=sub_query, - ) - db_session.add(sub_q) - db_session.flush() - return sub_q - - -def get_sub_questions_for_message( - db_session: Session, - primary_message_id: int, -) -> list[AgentSubQuestion]: - """Get all sub-questions for a given primary message.""" - return ( - db_session.query(AgentSubQuestion) - .filter(AgentSubQuestion.primary_question_id == primary_message_id) - .all() - ) - - -def get_sub_queries_for_question( - db_session: Session, - sub_question_id: int, -) -> list[AgentSubQuery]: - """Get all sub-queries for a given sub-question.""" - return ( - db_session.query(AgentSubQuery) - .filter(AgentSubQuery.parent_question_id == sub_question_id) - .all() - ) diff --git a/backend/onyx/db/chat.py b/backend/onyx/db/chat.py index 4497c6013..f586b294c 100644 --- a/backend/onyx/db/chat.py +++ b/backend/onyx/db/chat.py @@ -889,7 +889,7 @@ def translate_db_sub_questions_to_server_objects( sub_questions.append( SubQuestionDetail( level=sub_question.level, - level_question_nr=sub_question.level_question_nr, + level_question_nr=sub_question.level_question_num, question=sub_question.sub_question, answer=sub_question.sub_answer, sub_queries=sub_queries, @@ -1025,7 +1025,7 @@ def log_agent_sub_question_results( chat_session_id=chat_session_id, primary_question_id=primary_message_id, level=level, - level_question_nr=level_question_nr, + level_question_num=level_question_nr, sub_question=sub_question, sub_answer=sub_answer, sub_question_doc_results=sub_document_results, diff --git a/backend/onyx/db/models.py b/backend/onyx/db/models.py index b90b20f3a..9089b3cd2 100644 --- a/backend/onyx/db/models.py +++ b/backend/onyx/db/models.py @@ -1286,7 +1286,7 @@ class AgentSubQuestion(Base): ) sub_question: Mapped[str] = mapped_column(Text) level: Mapped[int] = mapped_column(Integer) - level_question_nr: Mapped[int] = mapped_column(Integer) + level_question_num: Mapped[int] = mapped_column(Integer) time_created: Mapped[datetime.datetime] = mapped_column( DateTime(timezone=True), server_default=func.now() )