Agent Search

.gitignore

@@ -7,4 +7,6 @@
 .vscode/
 *.sw?
 /backend/tests/regression/answer_quality/search_test_config.yaml
-/web/test-results/
+/web/test-results/
+backend/onyx/agent_search/main/test_data.json
+backend/tests/regression/answer_quality/test_data.json

.vscode/env_template.txt

@@ -51,3 +51,9 @@ BING_API_KEY=<REPLACE THIS>
 # Enable the full set of Danswer Enterprise Edition features
 # NOTE: DO NOT ENABLE THIS UNLESS YOU HAVE A PAID ENTERPRISE LICENSE (or if you are using this for local testing/development)
 ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=False
+
+# Agent Search configs  # TODO: Remove give proper namings
+AGENT_RETRIEVAL_STATS=False   # Note: This setting will incur substantial re-ranking effort
+AGENT_RERANKING_STATS=True
+AGENT_MAX_QUERY_RETRIEVAL_RESULTS=20
+AGENT_RERANKING_MAX_QUERY_RETRIEVAL_RESULTS=20

backend/alembic/versions/027381bce97c_add_shortcut_option_for_users.py

@@ -0,0 +1,29 @@
+"""add shortcut option for users
+
+Revision ID: 027381bce97c
+Revises: 6fc7886d665d
+Create Date: 2025-01-14 12:14:00.814390
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "027381bce97c"
+down_revision = "6fc7886d665d"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "user",
+        sa.Column(
+            "shortcut_enabled", sa.Boolean(), nullable=False, server_default="true"
+        ),
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("user", "shortcut_enabled")

backend/alembic/versions/1adf5ea20d2b_agent_doc_result_col.py

@@ -0,0 +1,29 @@
+"""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")

backend/alembic/versions/3c6531f32351_add_back_input_prompts.py

@@ -0,0 +1,58 @@
+"""add back input prompts
+
+Revision ID: 3c6531f32351
+Revises: aeda5f2df4f6
+Create Date: 2025-01-13 12:49:51.705235
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import fastapi_users_db_sqlalchemy
+
+# revision identifiers, used by Alembic.
+revision = "3c6531f32351"
+down_revision = "aeda5f2df4f6"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "inputprompt",
+        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
+        sa.Column("prompt", sa.String(), nullable=False),
+        sa.Column("content", sa.String(), nullable=False),
+        sa.Column("active", sa.Boolean(), nullable=False),
+        sa.Column("is_public", sa.Boolean(), nullable=False),
+        sa.Column(
+            "user_id",
+            fastapi_users_db_sqlalchemy.generics.GUID(),
+            nullable=True,
+        ),
+        sa.ForeignKeyConstraint(
+            ["user_id"],
+            ["user.id"],
+        ),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_table(
+        "inputprompt__user",
+        sa.Column("input_prompt_id", sa.Integer(), nullable=False),
+        sa.Column(
+            "user_id", fastapi_users_db_sqlalchemy.generics.GUID(), nullable=False
+        ),
+        sa.ForeignKeyConstraint(
+            ["input_prompt_id"],
+            ["inputprompt.id"],
+        ),
+        sa.ForeignKeyConstraint(
+            ["user_id"],
+            ["user.id"],
+        ),
+        sa.PrimaryKeyConstraint("input_prompt_id", "user_id"),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("inputprompt__user")
+    op.drop_table("inputprompt")

backend/alembic/versions/6fc7886d665d_make_categories_labels_and_many_to_many.py

@@ -0,0 +1,79 @@
+"""make categories labels and many to many
+
+Revision ID: 6fc7886d665d
+Revises: 3c6531f32351
+Create Date: 2025-01-13 18:12:18.029112
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "6fc7886d665d"
+down_revision = "3c6531f32351"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # Rename persona_category table to persona_label
+    op.rename_table("persona_category", "persona_label")
+
+    # Create the new association table
+    op.create_table(
+        "persona__persona_label",
+        sa.Column("persona_id", sa.Integer(), nullable=False),
+        sa.Column("persona_label_id", sa.Integer(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["persona_id"],
+            ["persona.id"],
+        ),
+        sa.ForeignKeyConstraint(
+            ["persona_label_id"],
+            ["persona_label.id"],
+        ),
+        sa.PrimaryKeyConstraint("persona_id", "persona_label_id"),
+    )
+
+    # Copy existing relationships to the new table
+    op.execute(
+        """
+        INSERT INTO persona__persona_label (persona_id, persona_label_id)
+        SELECT id, category_id FROM persona WHERE category_id IS NOT NULL
+    """
+    )
+
+    # Remove the old category_id column from persona table
+    op.drop_column("persona", "category_id")
+
+
+def downgrade() -> None:
+    # Rename persona_label table back to persona_category
+    op.rename_table("persona_label", "persona_category")
+
+    # Add back the category_id column to persona table
+    op.add_column("persona", sa.Column("category_id", sa.Integer(), nullable=True))
+    op.create_foreign_key(
+        "persona_category_id_fkey",
+        "persona",
+        "persona_category",
+        ["category_id"],
+        ["id"],
+    )
+
+    # Copy the first label relationship back to the persona table
+    op.execute(
+        """
+        UPDATE persona
+        SET category_id = (
+            SELECT persona_label_id
+            FROM persona__persona_label
+            WHERE persona__persona_label.persona_id = persona.id
+            LIMIT 1
+        )
+    """
+    )
+
+    # Drop the association table
+    op.drop_table("persona__persona_label")

backend/alembic/versions/925b58bd75b6_agent_metric_col_rename__s.py

@@ -0,0 +1,35 @@
+"""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"
+    )

backend/alembic/versions/9787be927e58_agent_metric_table_renames__agent__.py

@@ -0,0 +1,25 @@
+"""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")

backend/alembic/versions/98a5008d8711_agent_tracking.py

@@ -0,0 +1,42 @@
+"""agent_tracking
+
+Revision ID: 98a5008d8711
+Revises: 027381bce97c
+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 = "027381bce97c"
+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")

backend/alembic/versions/aeda5f2df4f6_add_pinned_assistants.py

@@ -0,0 +1,27 @@
+"""add pinned assistants
+
+Revision ID: aeda5f2df4f6
+Revises: 369644546676
+Create Date: 2025-01-09 16:04:10.770636
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "aeda5f2df4f6"
+down_revision = "369644546676"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "user", sa.Column("pinned_assistants", postgresql.JSONB(), nullable=True)
+    )
+    op.execute('UPDATE "user" SET pinned_assistants = chosen_assistants')
+
+
+def downgrade() -> None:
+    op.drop_column("user", "pinned_assistants")