functioning well

- may need tweaking

.github/workflows/docker-build-push-model-server-container-on-tag.yml

@@ -118,6 +118,6 @@ jobs:
           TRIVY_DB_REPOSITORY: "public.ecr.aws/aquasecurity/trivy-db:2"
           TRIVY_JAVA_DB_REPOSITORY: "public.ecr.aws/aquasecurity/trivy-java-db:1"
         with:
-          image-ref: docker.io/onyxdotapp/onyx-model-server:${{ github.ref_name }}
+          image-ref: docker.io/${{ env.REGISTRY_IMAGE }}:${{ github.ref_name }}
           severity: "CRITICAL,HIGH"
           timeout: "10m"

.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

@@ -5,6 +5,8 @@
 # For local dev, often user Authentication is not needed
 AUTH_TYPE=disabled
 
+# Skip warm up for dev
+SKIP_WARM_UP=True
 
 # Always keep these on for Dev
 # Logs all model prompts to stdout
@@ -49,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

.vscode/launch.template.jsonc

@@ -355,5 +355,20 @@
                 "PYTHONPATH": "."
             },
         },
+        {
+            "name": "Install Python Requirements",
+            "type": "node",
+            "request": "launch",
+            "runtimeExecutable": "bash",
+            "runtimeArgs": [
+                "-c",
+                "pip install -r backend/requirements/default.txt && pip install -r backend/requirements/dev.txt && pip install -r backend/requirements/ee.txt && pip install -r backend/requirements/model_server.txt"
+            ],
+            "cwd": "${workspaceFolder}",
+            "console": "integratedTerminal",
+            "presentation": {
+                 "group": "3"
+            }
+        },
     ]
 }

CONTRIBUTING.md

@@ -12,6 +12,10 @@ As an open source project in a rapidly changing space, we welcome all contributi
 
 The [GitHub Issues](https://github.com/onyx-dot-app/onyx/issues) page is a great place to start for contribution ideas.
 
+To ensure that your contribution is aligned with the project's direction, please reach out to Hagen (or any other maintainer) on the Onyx team
+via [Slack](https://join.slack.com/t/onyx-dot-app/shared_invite/zt-2twesxdr6-5iQitKZQpgq~hYIZ~dv3KA) /
+[Discord](https://discord.gg/TDJ59cGV2X) or [email](mailto:founders@onyx.app).
+
 Issues that have been explicitly approved by the maintainers (aligned with the direction of the project)
 will be marked with the `approved by maintainers` label.
 Issues marked `good first issue` are an especially great place to start.
@@ -23,8 +27,8 @@ If you have a new/different contribution in mind, we'd love to hear about it!
 Your input is vital to making sure that Onyx moves in the right direction.
 Before starting on implementation, please raise a GitHub issue.
 
-And always feel free to message us (Chris Weaver / Yuhong Sun) on
-[Slack](https://join.slack.com/t/danswer/shared_invite/zt-1w76msxmd-HJHLe3KNFIAIzk_0dSOKaQ) /
+Also, always feel free to message the founders (Chris Weaver / Yuhong Sun) on
+[Slack](https://join.slack.com/t/onyx-dot-app/shared_invite/zt-2twesxdr6-5iQitKZQpgq~hYIZ~dv3KA) /
 [Discord](https://discord.gg/TDJ59cGV2X) directly about anything at all.
 
 ### Contributing Code
@@ -42,7 +46,7 @@ Our goal is to make contributing as easy as possible. If you run into any issues
 That way we can help future contributors and users can avoid the same issue.
 
 We also have support channels and generally interesting discussions on our
-[Slack](https://join.slack.com/t/danswer/shared_invite/zt-1w76msxmd-HJHLe3KNFIAIzk_0dSOKaQ)
+[Slack](https://join.slack.com/t/onyx-dot-app/shared_invite/zt-2twesxdr6-5iQitKZQpgq~hYIZ~dv3KA)
 and
 [Discord](https://discord.gg/TDJ59cGV2X).
 
@@ -123,7 +127,47 @@ Once the above is done, navigate to `onyx/web` run:
 npm i
 ```
 
-#### Docker containers for external software
+## Formatting and Linting
+
+### Backend
+
+For the backend, you'll need to setup pre-commit hooks (black / reorder-python-imports).
+First, install pre-commit (if you don't have it already) following the instructions
+[here](https://pre-commit.com/#installation).
+
+With the virtual environment active, install the pre-commit library with:
+
+```bash
+pip install pre-commit
+```
+
+Then, from the `onyx/backend` directory, run:
+
+```bash
+pre-commit install
+```
+
+Additionally, we use `mypy` for static type checking.
+Onyx is fully type-annotated, and we want to keep it that way!
+To run the mypy checks manually, run `python -m mypy .` from the `onyx/backend` directory.
+
+### Web
+
+We use `prettier` for formatting. The desired version (2.8.8) will be installed via a `npm i` from the `onyx/web` directory.
+To run the formatter, use `npx prettier --write .` from the `onyx/web` directory.
+Please double check that prettier passes before creating a pull request.
+
+# Running the application for development
+
+## Developing using VSCode Debugger (recommended)
+
+We highly recommend using VSCode debugger for development.
+See [CONTRIBUTING_VSCODE.md](./CONTRIBUTING_VSCODE.md) for more details.
+
+Otherwise, you can follow the instructions below to run the application for development.
+
+## Manually running the application for development
+### Docker containers for external software
 
 You will need Docker installed to run these containers.
 
@@ -135,7 +179,7 @@ docker compose -f docker-compose.dev.yml -p onyx-stack up -d index relational_db
 
 (index refers to Vespa, relational_db refers to Postgres, and cache refers to Redis)
 
-#### Running Onyx locally
+### Running Onyx locally
 
 To start the frontend, navigate to `onyx/web` and run:
 
@@ -223,35 +267,6 @@ If you want to make changes to Onyx and run those changes in Docker, you can als
 docker compose -f docker-compose.dev.yml -p onyx-stack up -d --build
 ```
 
-### Formatting and Linting
-
-#### Backend
-
-For the backend, you'll need to setup pre-commit hooks (black / reorder-python-imports).
-First, install pre-commit (if you don't have it already) following the instructions
-[here](https://pre-commit.com/#installation).
-
-With the virtual environment active, install the pre-commit library with:
-
-```bash
-pip install pre-commit
-```
-
-Then, from the `onyx/backend` directory, run:
-
-```bash
-pre-commit install
-```
-
-Additionally, we use `mypy` for static type checking.
-Onyx is fully type-annotated, and we want to keep it that way!
-To run the mypy checks manually, run `python -m mypy .` from the `onyx/backend` directory.
-
-#### Web
-
-We use `prettier` for formatting. The desired version (2.8.8) will be installed via a `npm i` from the `onyx/web` directory.
-To run the formatter, use `npx prettier --write .` from the `onyx/web` directory.
-Please double check that prettier passes before creating a pull request.
 
 ### Release Process
 

CONTRIBUTING_VSCODE.md

@@ -0,0 +1,29 @@
+# VSCode Debugging Setup
+
+This guide explains how to set up and use VSCode's debugging capabilities with this project.
+
+## Initial Setup
+
+1. **Environment Setup**:
+   - Copy `.vscode/.env.template` to `.vscode/.env`
+   - Fill in the necessary environment variables in `.vscode/.env`
+2. **launch.json**:
+   - Copy `.vscode/launch.template.jsonc` to `.vscode/launch.json`
+
+## Using the Debugger
+
+Before starting, make sure the Docker Daemon is running.
+
+1. Open the Debug view in VSCode (Cmd+Shift+D on macOS)
+2. From the dropdown at the top, select "Clear and Restart External Volumes and Containers" and press the green play button
+3. From the dropdown at the top, select "Run All Onyx Services" and press the green play button
+4. Now, you can navigate to onyx in your browser (default is http://localhost:3000) and start using the app
+5. You can set breakpoints by clicking to the left of line numbers to help debug while the app is running
+6. Use the debug toolbar to step through code, inspect variables, etc.
+
+## Features
+
+- Hot reload is enabled for the web server and API servers
+- Python debugging is configured with debugpy
+- Environment variables are loaded from `.vscode/.env`
+- Console output is organized in the integrated terminal with labeled tabs

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/369644546676_add_composite_index_for_index_attempt_.py

@@ -0,0 +1,35 @@
+"""add composite index for index attempt time updated
+
+Revision ID: 369644546676
+Revises: 2955778aa44c
+Create Date: 2025-01-08 15:38:17.224380
+
+"""
+from alembic import op
+from sqlalchemy import text
+
+# revision identifiers, used by Alembic.
+revision = "369644546676"
+down_revision = "2955778aa44c"
+branch_labels: None = None
+depends_on: None = None
+
+
+def upgrade() -> None:
+    op.create_index(
+        "ix_index_attempt_ccpair_search_settings_time_updated",
+        "index_attempt",
+        [
+            "connector_credential_pair_id",
+            "search_settings_id",
+            text("time_updated DESC"),
+        ],
+        unique=False,
+    )
+
+
+def downgrade() -> None:
+    op.drop_index(
+        "ix_index_attempt_ccpair_search_settings_time_updated",
+        table_name="index_attempt",
+    )

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")

backend/alembic/versions/bceb76d618ec_agent_table_renames__agent__.py

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

backend/alembic/versions/c0132518a25b_agent_table_changes_rename_level.py

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

backend/alembic/versions/e9cf2bd7baed_create_pro_search_persistence_tables.py

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

backend/chatt.txt

@@ -0,0 +1,536 @@
+"{\"user_message_id\": 475, \"reserved_assistant_message_id\": 476}\n"
+"{\"sub_question\": \"What\", \"level\": 0, \"level_question_nr\": 1}\n"
+"{\"sub_query\": \"ony\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 0}\n"
+"{\"sub_question\": \" is\", \"level\": 0, \"level_question_nr\": 1}\n"
+"{\"sub_query\": \"x\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 0}\n"
+"{\"sub_question\": \" On\", \"level\": 0, \"level_question_nr\": 1}\n"
+"{\"sub_query\": \" \", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 0}\n"
+"{\"sub_question\": \"yx\", \"level\": 0, \"level_question_nr\": 1}\n"
+"{\"sub_question\": \" \", \"level\": 0, \"level_question_nr\": 1}\n"
+"{\"sub_query\": \"1\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 0}\n"
+"{\"sub_question\": \"1\", \"level\": 0, \"level_question_nr\": 1}\n"
+"{\"sub_query\": \" features\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 0}\n"
+"{\"sub_query\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 0}\n"
+"{\"sub_question\": \"?\", \"level\": 0, \"level_question_nr\": 1}\n"
+"{\"sub_question\": \"  \", \"level\": 0, \"level_question_nr\": 1}\n"
+"{\"sub_question\": \"\", \"level\": 0, \"level_question_nr\": 2}\n"
+"{\"sub_question\": \"What\", \"level\": 0, \"level_question_nr\": 2}\n"
+"{\"sub_question\": \" is\", \"level\": 0, \"level_question_nr\": 2}\n"
+"{\"sub_query\": \" specifications\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 0}\n"
+"{\"sub_question\": \" On\", \"level\": 0, \"level_question_nr\": 2}\n"
+"{\"sub_query\": \"  \", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 0}\n"
+"{\"sub_query\": \"\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_question\": \"yx\", \"level\": 0, \"level_question_nr\": 2}\n"
+"{\"sub_question\": \" \", \"level\": 0, \"level_question_nr\": 2}\n"
+"{\"sub_question\": \"2\", \"level\": 0, \"level_question_nr\": 2}\n"
+"{\"sub_query\": \"ony\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_question\": \"?\", \"level\": 0, \"level_question_nr\": 2}\n"
+"{\"sub_query\": \"x\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_question\": \"  \", \"level\": 0, \"level_question_nr\": 2}\n"
+"{\"sub_query\": \" \", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_question\": \"\", \"level\": 0, \"level_question_nr\": 3}\n"
+"{\"sub_query\": \"2\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_question\": \"What\", \"level\": 0, \"level_question_nr\": 3}\n"
+"{\"sub_query\": \" applications\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_question\": \" is\", \"level\": 0, \"level_question_nr\": 3}\n"
+"{\"sub_query\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_question\": \" On\", \"level\": 0, \"level_question_nr\": 3}\n"
+"{\"sub_query\": \" use\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_question\": \"yx\", \"level\": 0, \"level_question_nr\": 3}\n"
+"{\"sub_query\": \" cases\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_question\": \" \", \"level\": 0, \"level_question_nr\": 3}\n"
+"{\"sub_question\": \"3\", \"level\": 0, \"level_question_nr\": 3}\n"
+"{\"sub_question\": \"?\", \"level\": 0, \"level_question_nr\": 3}\n"
+"{\"sub_question\": \"  \", \"level\": 0, \"level_question_nr\": 3}\n"
+"{\"sub_question\": \"\", \"level\": 0, \"level_question_nr\": 4}\n"
+"{\"sub_question\": \"What\", \"level\": 0, \"level_question_nr\": 4}\n"
+"{\"sub_question\": \" is\", \"level\": 0, \"level_question_nr\": 4}\n"
+"{\"sub_query\": \"  \", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_question\": \" On\", \"level\": 0, \"level_question_nr\": 4}\n"
+"{\"sub_query\": \"\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_question\": \"yx\", \"level\": 0, \"level_question_nr\": 4}\n"
+"{\"sub_question\": \" \", \"level\": 0, \"level_question_nr\": 4}\n"
+"{\"sub_query\": \"ony\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_question\": \"4\", \"level\": 0, \"level_question_nr\": 4}\n"
+"{\"sub_question\": \"?\", \"level\": 0, \"level_question_nr\": 4}\n"
+"{\"sub_query\": \"x\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \" \", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_question\": \"  \", \"level\": 0, \"level_question_nr\": 4}\n"
+"{\"sub_question\": \"\", \"level\": 0, \"level_question_nr\": 4}\n"
+"{\"sub_query\": \"3\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \" \", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \"4\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \" comparison\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \" differences\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \"  \", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \"\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \"On\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 0}\n"
+"{\"sub_query\": \"yx\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 0}\n"
+"{\"sub_query\": \"On\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 0}\n"
+"{\"sub_query\": \"yx\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 0}\n"
+"{\"sub_query\": \" \", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 0}\n"
+"{\"sub_query\": \"4\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 0}\n"
+"{\"sub_query\": \"On\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 0}\n"
+"{\"sub_query\": \"yx\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 0}\n"
+"{\"sub_query\": \" \", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 0}\n"
+"{\"sub_query\": \"1\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 0}\n"
+"{\"sub_query\": \"On\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 0}\n"
+"{\"sub_query\": \" \", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 0}\n"
+"{\"sub_query\": \"yx\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 0}\n"
+"{\"sub_query\": \" product\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 0}\n"
+"{\"sub_query\": \"3\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 0}\n"
+"{\"sub_query\": \" information\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 0}\n"
+"{\"sub_query\": \" software\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 0}\n"
+"{\"sub_query\": \" features\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 0}\n"
+"{\"sub_query\": \" \", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 0}\n"
+"{\"sub_query\": \"  \", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 0}\n"
+"{\"sub_query\": \" software\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 0}\n"
+"{\"sub_query\": \"2\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 0}\n"
+"{\"sub_query\": \"\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 1}\n"
+"{\"sub_query\": \" features\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 0}\n"
+"{\"sub_query\": \"On\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 1}\n"
+"{\"sub_query\": \"  \", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 0}\n"
+"{\"sub_query\": \"\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_query\": \"On\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_query\": \"yx\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 1}\n"
+"{\"sub_query\": \" \", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 1}\n"
+"{\"sub_query\": \"yx\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_query\": \" \", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_query\": \" software\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 0}\n"
+"{\"sub_query\": \" features\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 0}\n"
+"{\"sub_query\": \"4\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 1}\n"
+"{\"sub_query\": \" applications\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 1}\n"
+"{\"sub_query\": \"  \", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 0}\n"
+"{\"sub_query\": \"1\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_query\": \"\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 1}\n"
+"{\"sub_query\": \" features\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_query\": \"On\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 1}\n"
+"{\"sub_query\": \" in\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 1}\n"
+"{\"sub_query\": \" industry\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 1}\n"
+"{\"sub_query\": \"  \", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 0}\n"
+"{\"sub_query\": \"\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 1}\n"
+"{\"sub_query\": \"yx\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 1}\n"
+"{\"sub_query\": \"On\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 1}\n"
+"{\"sub_query\": \" \", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 1}\n"
+"{\"sub_query\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_query\": \"yx\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 1}\n"
+"{\"sub_query\": \" specifications\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_query\": \" \", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 1}\n"
+"{\"sub_query\": \"2\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 1}\n"
+"{\"sub_query\": \" applications\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 1}\n"
+"{\"sub_query\": \"  \", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 1}\n"
+"{\"sub_query\": \"  \", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 1}\n"
+"{\"sub_query\": \"\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 2}\n"
+"{\"sub_query\": \"\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \"On\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 2}\n"
+"{\"sub_query\": \"On\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \"3\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 1}\n"
+"{\"sub_query\": \"yx\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \" applications\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 1}\n"
+"{\"sub_query\": \" \", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \"yx\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 2}\n"
+"{\"sub_query\": \" in\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 1}\n"
+"{\"sub_query\": \" \", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 2}\n"
+"{\"sub_query\": \" industry\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 1}\n"
+"{\"sub_query\": \"1\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \" applications\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \"  \", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 1}\n"
+"{\"sub_query\": \"\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 2}\n"
+"{\"sub_query\": \"On\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 2}\n"
+"{\"sub_query\": \"4\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 2}\n"
+"{\"sub_query\": \" comparison\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 2}\n"
+"{\"sub_query\": \" with\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 2}\n"
+"{\"sub_query\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \" previous\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 2}\n"
+"{\"sub_query\": \" use\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \" versions\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 2}\n"
+"{\"sub_query\": \" in\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 1}\n"
+"{\"sub_query\": \" industry\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 1}\n"
+"{\"sub_query\": \"  \", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 2}\n"
+"{\"sub_query\": \"\", \"level\": 0, \"level_question_nr\": 3, \"query_id\": 2}\n"
+"{\"sub_query\": \"yx\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 2}\n"
+"{\"sub_query\": \"  \", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 1}\n"
+"{\"sub_query\": \"\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 2}\n"
+"{\"sub_query\": \" \", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 2}\n"
+"{\"sub_query\": \"On\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 2}\n"
+"{\"sub_query\": \"yx\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 2}\n"
+"{\"sub_query\": \" \", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 2}\n"
+"{\"sub_query\": \"3\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 2}\n"
+"{\"sub_query\": \" comparison\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 2}\n"
+"{\"sub_query\": \"2\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 2}\n"
+"{\"sub_query\": \" comparison\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 2}\n"
+"{\"sub_query\": \" with\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 2}\n"
+"{\"sub_query\": \" cases\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \" previous\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 2}\n"
+"{\"sub_query\": \"  \", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \"\", \"level\": 0, \"level_question_nr\": 0, \"query_id\": 2}\n"
+"{\"sub_query\": \" with\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 2}\n"
+"{\"sub_query\": \" other\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 2}\n"
+"{\"sub_query\": \" software\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 2}\n"
+"{\"sub_query\": \"\", \"level\": 0, \"level_question_nr\": 1, \"query_id\": 2}\n"
+"{\"sub_query\": \" versions\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 2}\n"
+"{\"sub_query\": \"  \", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 2}\n"
+"{\"sub_query\": \"\", \"level\": 0, \"level_question_nr\": 2, \"query_id\": 2}\n"
+"{\"top_documents\": [], \"rephrased_query\": \"What is Onyx 4?\", \"predicted_flow\": \"question-answer\", \"predicted_search\": \"keyword\", \"applied_source_filters\": null, \"applied_time_cutoff\": null, \"recency_bias_multiplier\": 0.5}\n"
+"{\"llm_selected_doc_indices\": []}\n"
+"{\"final_context_docs\": []}\n"
+"{\"answer_piece\": \"I\", \"level\": 0, \"level_question_nr\": 3, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" don't\", \"level\": 0, \"level_question_nr\": 3, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"On\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" know\", \"level\": 0, \"level_question_nr\": 3, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"yx\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \".\", \"level\": 0, \"level_question_nr\": 3, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" \", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"1\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" formerly\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" known\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" as\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" D\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"answer\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" is\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" an\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" AI\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" Assistant\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" that\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" connects\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" to\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" a\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" company's\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" documents\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" applications\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" personnel\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \".\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" It\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" provides\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" a\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" chat\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" interface\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" can\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" integrate\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" with\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" any\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" large\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" language\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" model\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" (\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"LL\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"M\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \")\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"top_documents\": [], \"rephrased_query\": \"What is Onyx 2?\", \"predicted_flow\": \"question-answer\", \"predicted_search\": \"keyword\", \"applied_source_filters\": null, \"applied_time_cutoff\": null, \"recency_bias_multiplier\": 0.5}\n"
+"{\"llm_selected_doc_indices\": []}\n"
+"{\"final_context_docs\": []}\n"
+"{\"answer_piece\": \" of\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" choice\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \".\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" On\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"yx\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" is\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" designed\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" to\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" be\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" modular\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" easily\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" extens\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"ible\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" allowing\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" for\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" deployment\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" on\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" various\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" platforms\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" including\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" laptops\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" on\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"-prem\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"ise\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" or\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" cloud\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" environments\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \".\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" It\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" ensures\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" that\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" user\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" data\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" chats\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" remain\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"I\", \"level\": 0, \"level_question_nr\": 1, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" under\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" don't\", \"level\": 0, \"level_question_nr\": 1, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" the\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" user's\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" know\", \"level\": 0, \"level_question_nr\": 1, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \".\", \"level\": 0, \"level_question_nr\": 1, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" control\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" as\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" the\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" deployment\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" is\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" owned\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" by\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" the\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" user\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \".\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" On\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"yx\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" is\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" MIT\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" licensed\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" comes\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" ready\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" for\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" production\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" use\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" featuring\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" user\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" authentication\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" role\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" management\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" chat\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" persistence\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" a\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" user\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" interface\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" for\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" configuring\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" AI\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" Assist\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"ants\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" their\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" prompts\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \".\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" Additionally\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" On\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"yx\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" serves\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" as\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" a\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" unified\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" search\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" tool\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" across\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" common\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" workplace\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" applications\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" like\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" Slack\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" Google\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" Drive\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" Con\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"fluence\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" enabling\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" it\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" to\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" act\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" as\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" a\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" subject\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" matter\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" expert\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" for\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" teams\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" by\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" combining\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" L\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"LM\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"s\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" with\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" team\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \"-specific\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" knowledge\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" [[1]]()\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"top_documents\": [], \"rephrased_query\": \"What is Onyx 3?\", \"predicted_flow\": \"question-answer\", \"predicted_search\": \"keyword\", \"applied_source_filters\": null, \"applied_time_cutoff\": null, \"recency_bias_multiplier\": 0.5}\n"
+"{\"llm_selected_doc_indices\": []}\n"
+"{\"final_context_docs\": []}\n"
+"{\"answer_piece\": \"I\", \"level\": 0, \"level_question_nr\": 2, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" don't\", \"level\": 0, \"level_question_nr\": 2, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \" know\", \"level\": 0, \"level_question_nr\": 2, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"answer_piece\": \".\", \"level\": 0, \"level_question_nr\": 2, \"answer_type\": \"agent_sub_answer\"}\n"
+"{\"top_documents\": [{\"document_id\": \"https://docs.onyx.app/introduction\", \"chunk_ind\": 0, \"semantic_identifier\": \"Introduction - Onyx Documentation\", \"link\": \"https://docs.onyx.app/introduction\", \"blurb\": \"Onyx Documentation home page\\nSearch...\\nNavigation\\nWelcome to Onyx\\nIntroduction\\nWelcome to Onyx\\nIntroduction\\nOnyx Overview\\n\\nWhat is Onyx\\nOnyx (Formerly Danswer) is the AI Assistant connected to your companys docs, apps, and people. Onyx provides a Chat interface and plugs into any LLM of your choice. Onyx can be deployed anywhere and for any scale - on a laptop, on-premise, or to cloud. Since you own the deployment, your user data and chats are fully in your own control. Onyx is MIT licensed and designed to be modular and easily extensible.\", \"source_type\": \"web\", \"boost\": 0, \"hidden\": false, \"metadata\": {}, \"score\": 0.6275177643886491, \"is_relevant\": null, \"relevance_explanation\": null, \"match_highlights\": [\"\", \"such as A customer wants feature X, is this already supported? or Wheres the pull request for feature Y?\\n<hi>Onyx</hi> can also be plugged into existing tools like Slack to get answers and AI chats directly in Slack.\\n\\nDemo\\n\\nMain <hi>Features</hi> \\n- Chat UI with the ability to select documents to chat with.\\n- Create custom AI Assistants\", \"\"], \"updated_at\": null, \"primary_owners\": null, \"secondary_owners\": null, \"is_internet\": false, \"db_doc_id\": 35923}], \"rephrased_query\": \"what is onyx 1, 2, 3, 4\", \"predicted_flow\": \"question-answer\", \"predicted_search\": \"keyword\", \"applied_source_filters\": null, \"applied_time_cutoff\": null, \"recency_bias_multiplier\": 0.5}\n"
+"{\"llm_selected_doc_indices\": []}\n"
+"{\"final_context_docs\": [{\"document_id\": \"https://docs.onyx.app/introduction\", \"content\": \"Onyx Documentation home page\\nSearch...\\nNavigation\\nWelcome to Onyx\\nIntroduction\\nWelcome to Onyx\\nIntroduction\\nOnyx Overview\\n\\nWhat is Onyx\\nOnyx (Formerly Danswer) is the AI Assistant connected to your companys docs, apps, and people. Onyx provides a Chat interface and plugs into any LLM of your choice. Onyx can be deployed anywhere and for any scale - on a laptop, on-premise, or to cloud. Since you own the deployment, your user data and chats are fully in your own control. Onyx is MIT licensed and designed to be modular and easily extensible. The system also comes fully ready for production usage with user authentication, role management (admin/basic users), chat persistence, and a UI for configuring Personas (AI Assistants) and their Prompts.\\nOnyx also serves as a Unified Search across all common workplace tools such as Slack, Google Drive, Confluence, etc. By combining LLMs and team specific knowledge, Onyx becomes a subject matter expert for the team. Its like ChatGPT if it had access to your teams unique knowledge! It enables questions such as A customer wants feature X, is this already supported? or Wheres the pull request for feature Y?\\nOnyx can also be plugged into existing tools like Slack to get answers and AI chats directly in Slack.\\n\\nDemo\\n\\nMain Features \\n- Chat UI with the ability to select documents to chat with.\\n- Create custom AI Assistants with different prompts and backing knowledge sets.\\n- Connect Onyx with LLM of your choice (self-host for a fully airgapped solution).\\n- Document Search + AI Answers for natural language queries.\\n- Connectors to all common workplace tools like Google Drive, Confluence, Slack, etc.\\n- Slack integration to get answers and search results directly in Slack.\\n\\nUpcoming\\n- Chat/Prompt sharing with specific teammates and user groups.\\n- Multi-modal model support, chat with images, video etc.\\n- Choosing between LLMs and parameters during chat session.\\n- Tool calling and agent configurations options.\\n- Organizational understanding and ability to locate and suggest experts from your team.\\n\\nOther Noteable Benefits of Onyx\\n- User Authentication with document level access management.\\n- Best in class Hybrid Search across all sources (BM-25 + prefix aware embedding models).\\n- Admin Dashboard to configure connectors, document-sets, access, etc.\\n- Custom deep learning models + learn from user feedback.\\n- Easy deployment and ability to host Onyx anywhere of your choosing.\\nQuickstart\", \"blurb\": \"Onyx Documentation home page\\nSearch...\\nNavigation\\nWelcome to Onyx\\nIntroduction\\nWelcome to Onyx\\nIntroduction\\nOnyx Overview\\n\\nWhat is Onyx\\nOnyx (Formerly Danswer) is the AI Assistant connected to your companys docs, apps, and people. Onyx provides a Chat interface and plugs into any LLM of your choice. Onyx can be deployed anywhere and for any scale - on a laptop, on-premise, or to cloud. Since you own the deployment, your user data and chats are fully in your own control. Onyx is MIT licensed and designed to be modular and easily extensible.\", \"semantic_identifier\": \"Introduction - Onyx Documentation\", \"source_type\": \"web\", \"metadata\": {}, \"updated_at\": null, \"link\": \"https://docs.onyx.app/introduction\", \"source_links\": {\"0\": \"https://docs.onyx.app/introduction\"}, \"match_highlights\": [\"\", \"such as A customer wants feature X, is this already supported? or Wheres the pull request for feature Y?\\n<hi>Onyx</hi> can also be plugged into existing tools like Slack to get answers and AI chats directly in Slack.\\n\\nDemo\\n\\nMain <hi>Features</hi> \\n- Chat UI with the ability to select documents to chat with.\\n- Create custom AI Assistants\", \"\"]}]}\n"
+"{\"answer_piece\": \"I\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" cannot\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" reliably\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" answer\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" the\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" question\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" about\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" On\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"yx\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" \", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"2\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" \", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"3\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" \", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"4\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" as\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" the\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" provided\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" information\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" only\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" describes\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" On\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"yx\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" \", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"1\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" which\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" is\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" an\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" AI\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" Assistant\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" formerly\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" known\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" as\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" D\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"answer\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \".\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" On\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"yx\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" \", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"1\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" connects\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" to\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" a\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" company's\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" documents\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" applications\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" personnel\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" providing\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" a\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" chat\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" interface\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" integration\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" with\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" any\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" large\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" language\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" model\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" (\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"LL\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"M\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \")\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" of\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" choice\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \".\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" It\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" is\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" designed\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" to\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" be\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" modular\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" easily\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" extens\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"ible\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" can\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" be\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" deployed\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" on\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" various\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" platforms\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" while\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" ensuring\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" user\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" data\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" control\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \".\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" It\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" also\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" serves\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" as\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" a\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" unified\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" search\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" tool\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" across\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" common\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" workplace\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" applications\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" like\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" Slack\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" Google\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" Drive\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" and\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" Con\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"fluence\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" acting\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" as\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" a\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" subject\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" matter\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" expert\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" for\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" teams\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" [[1]]()\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"{{1}}\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"There\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" is\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" no\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" information\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" available\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" regarding\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" On\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"yx\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" \", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"2\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" \", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"3\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" or\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" \", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \"4\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \",\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" so\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" I\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" cannot\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" provide\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" details\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" about\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \" them\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"answer_piece\": \".\", \"level\": 0, \"level_question_nr\": 0, \"answer_type\": \"agent_level_answer\"}\n"
+"{\"citations\": []}\n"
+"{\"message_id\": 476, \"parent_message\": 475, \"latest_child_message\": null, \"message\": \"I cannot reliably answer the question about Onyx 2, 3, and 4, as the provided information only describes Onyx 1, which is an AI Assistant formerly known as Danswer. Onyx 1 connects to a company's documents, applications, and personnel, providing a chat interface and integration with any large language model (LLM) of choice. It is designed to be modular, easily extensible, and can be deployed on various platforms while ensuring user data control. It also serves as a unified search tool across common workplace applications like Slack, Google Drive, and Confluence, acting as a subject matter expert for teams [[1]](){{1}}There is no information available regarding Onyx 2, 3, or 4, so I cannot provide details about them.\", \"rephrased_query\": \"what is onyx 1, 2, 3, 4\", \"context_docs\": {\"top_documents\": [{\"document_id\": \"https://docs.onyx.app/introduction\", \"chunk_ind\": 0, \"semantic_identifier\": \"Introduction - Onyx Documentation\", \"link\": \"https://docs.onyx.app/introduction\", \"blurb\": \"Onyx Documentation home page\\nSearch...\\nNavigation\\nWelcome to Onyx\\nIntroduction\\nWelcome to Onyx\\nIntroduction\\nOnyx Overview\\n\\nWhat is Onyx\\nOnyx (Formerly Danswer) is the AI Assistant connected to your companys docs, apps, and people. Onyx provides a Chat interface and plugs into any LLM of your choice. Onyx can be deployed anywhere and for any scale - on a laptop, on-premise, or to cloud. Since you own the deployment, your user data and chats are fully in your own control. Onyx is MIT licensed and designed to be modular and easily extensible.\", \"source_type\": \"web\", \"boost\": 0, \"hidden\": false, \"metadata\": {}, \"score\": 0.6275177643886491, \"is_relevant\": null, \"relevance_explanation\": null, \"match_highlights\": [\"\", \"such as A customer wants feature X, is this already supported? or Wheres the pull request for feature Y?\\n<hi>Onyx</hi> can also be plugged into existing tools like Slack to get answers and AI chats directly in Slack.\\n\\nDemo\\n\\nMain <hi>Features</hi> \\n- Chat UI with the ability to select documents to chat with.\\n- Create custom AI Assistants\", \"\"], \"updated_at\": null, \"primary_owners\": null, \"secondary_owners\": null, \"is_internet\": false, \"db_doc_id\": 35923}]}, \"message_type\": \"assistant\", \"time_sent\": \"2025-01-12T05:37:18.318251+00:00\", \"overridden_model\": \"gpt-4o\", \"alternate_assistant_id\": 0, \"chat_session_id\": \"40f91916-7419-48d1-9681-5882b0869d88\", \"citations\": {}, \"sub_questions\": [], \"files\": [], \"tool_call\": null}\n"

backend/ee/onyx/db/analytics.py

@@ -345,7 +345,8 @@ def fetch_assistant_unique_users_total(
 def user_can_view_assistant_stats(
     db_session: Session, user: User | None, assistant_id: int
 ) -> bool:
-    # If user is None, assume the user is an admin or auth is disabled
+    # If user is None and auth is disabled, assume the user is an admin
+
     if user is None or user.role == UserRole.ADMIN:
         return True
 

backend/ee/onyx/db/token_limit.py

@@ -7,6 +7,7 @@ from sqlalchemy import select
 from sqlalchemy.orm import aliased
 from sqlalchemy.orm import Session
 
+from onyx.configs.app_configs import DISABLE_AUTH
 from onyx.configs.constants import TokenRateLimitScope
 from onyx.db.models import TokenRateLimit
 from onyx.db.models import TokenRateLimit__UserGroup
@@ -20,10 +21,11 @@ from onyx.server.token_rate_limits.models import TokenRateLimitArgs
 def _add_user_filters(
     stmt: Select, user: User | None, get_editable: bool = True
 ) -> Select:
-    # If user is None, assume the user is an admin or auth is disabled
-    if user is None or user.role == UserRole.ADMIN:
+    # If user is None and auth is disabled, assume the user is an admin
+    if (user is None and DISABLE_AUTH) or (user and user.role == UserRole.ADMIN):
         return stmt
 
+    stmt = stmt.distinct()
     TRLimit_UG = aliased(TokenRateLimit__UserGroup)
     User__UG = aliased(User__UserGroup)
 
@@ -46,6 +48,12 @@ def _add_user_filters(
     that the user isn't a curator for
     - if we are not editing, we show all token_rate_limits in the groups the user curates
     """
+
+    # If user is None, this is an anonymous user and we should only show public token_rate_limits
+    if user is None:
+        where_clause = TokenRateLimit.scope == TokenRateLimitScope.GLOBAL
+        return stmt.where(where_clause)
+
     where_clause = User__UG.user_id == user.id
     if user.role == UserRole.CURATOR and get_editable:
         where_clause &= User__UG.is_curator == True  # noqa: E712

backend/ee/onyx/external_permissions/confluence/doc_sync.py

@@ -24,7 +24,9 @@ _REQUEST_PAGINATION_LIMIT = 5000
 def _get_server_space_permissions(
     confluence_client: OnyxConfluence, space_key: str
 ) -> ExternalAccess:
-    space_permissions = confluence_client.get_space_permissions(space_key=space_key)
+    space_permissions = confluence_client.get_all_space_permissions_server(
+        space_key=space_key
+    )
 
     viewspace_permissions = []
     for permission_category in space_permissions:
@@ -67,6 +69,13 @@ def _get_server_space_permissions(
         else:
             logger.warning(f"Email for user {user_name} not found in Confluence")
 
+    if not user_emails and not group_names:
+        logger.warning(
+            "No user emails or group names found in Confluence space permissions"
+            f"\nSpace key: {space_key}"
+            f"\nSpace permissions: {space_permissions}"
+        )
+
     return ExternalAccess(
         external_user_emails=user_emails,
         external_user_group_ids=group_names,

backend/ee/onyx/external_permissions/confluence/group_sync.py

@@ -30,6 +30,7 @@ def _build_group_member_email_map(
                 )
         if not email:
             # If we still don't have an email, skip this user
+            logger.warning(f"user result missing email field: {user_result}")
             continue
 
         for group in confluence_client.paginated_groups_by_user_retrieval(user):

backend/ee/onyx/server/analytics/api.py

@@ -228,6 +228,8 @@ def get_assistant_stats(
         datetime.datetime.utcnow() - datetime.timedelta(days=_DEFAULT_LOOKBACK_DAYS)
     )
     end = end or datetime.datetime.utcnow()
+    print("current user")
+    print(user)
 
     if not user_can_view_assistant_stats(db_session, user, assistant_id):
         raise HTTPException(

backend/ee/onyx/server/query_and_chat/chat_backend.py

@@ -179,6 +179,7 @@ def handle_simplified_chat_message(
         chunks_below=0,
         full_doc=chat_message_req.full_doc,
         structured_response_format=chat_message_req.structured_response_format,
+        use_agentic_search=chat_message_req.use_agentic_search,
     )
 
     packets = stream_chat_message_objects(
@@ -301,6 +302,7 @@ def handle_send_message_simple_with_history(
         chunks_below=0,
         full_doc=req.full_doc,
         structured_response_format=req.structured_response_format,
+        use_agentic_search=req.use_agentic_search,
     )
 
     packets = stream_chat_message_objects(

backend/ee/onyx/server/query_and_chat/models.py

@@ -57,6 +57,9 @@ class BasicCreateChatMessageRequest(ChunkContext):
     # https://platform.openai.com/docs/guides/structured-outputs/introduction
     structured_response_format: dict | None = None
 
+    # If True, uses agentic search instead of basic search
+    use_agentic_search: bool = False
+
 
 class BasicCreateChatMessageWithHistoryRequest(ChunkContext):
     # Last element is the new query. All previous elements are historical context
@@ -71,6 +74,8 @@ class BasicCreateChatMessageWithHistoryRequest(ChunkContext):
     # only works if using an OpenAI model. See the following for more details:
     # https://platform.openai.com/docs/guides/structured-outputs/introduction
     structured_response_format: dict | None = None
+    # If True, uses agentic search instead of basic search
+    use_agentic_search: bool = False
 
 
 class SimpleDoc(BaseModel):
@@ -123,6 +128,9 @@ class OneShotQARequest(ChunkContext):
     # If True, skips generative an AI response to the search query
     skip_gen_ai_answer_generation: bool = False
 
+    # If True, uses pro search instead of basic search
+    use_agentic_search: bool = False
+
     @model_validator(mode="after")
     def check_persona_fields(self) -> "OneShotQARequest":
         if self.persona_override_config is None and self.persona_id is None:

backend/ee/onyx/server/query_and_chat/query_backend.py

@@ -196,6 +196,7 @@ def get_answer_stream(
         retrieval_details=query_request.retrieval_options,
         rerank_settings=query_request.rerank_settings,
         db_session=db_session,
+        use_agentic_search=query_request.use_agentic_search,
     )
 
     packets = stream_chat_message_objects(

backend/onyx/agent_search/basic/graph_builder.py

@@ -0,0 +1,103 @@
+from typing import cast
+
+from langchain_core.callbacks.manager import dispatch_custom_event
+from langchain_core.runnables.config import RunnableConfig
+from langgraph.graph import END
+from langgraph.graph import START
+from langgraph.graph import StateGraph
+
+from onyx.agent_search.basic.states import BasicInput
+from onyx.agent_search.basic.states import BasicOutput
+from onyx.agent_search.basic.states import BasicState
+from onyx.agent_search.basic.states import BasicStateUpdate
+from onyx.agent_search.models import ProSearchConfig
+from onyx.chat.stream_processing.utils import (
+    map_document_id_order,
+)
+from onyx.tools.tool_implementations.search.search_tool import SearchTool
+
+
+def basic_graph_builder() -> StateGraph:
+    graph = StateGraph(
+        state_schema=BasicState,
+        input=BasicInput,
+        output=BasicOutput,
+    )
+
+    ### Add nodes ###
+
+    graph.add_node(
+        node="get_response",
+        action=get_response,
+    )
+
+    ### Add edges ###
+
+    graph.add_edge(start_key=START, end_key="get_response")
+
+    graph.add_conditional_edges("get_response", should_continue, ["get_response", END])
+    graph.add_edge(
+        start_key="get_response",
+        end_key=END,
+    )
+
+    return graph
+
+
+def should_continue(state: BasicState) -> str:
+    return (
+        END if state["last_llm_call"] is None or state["calls"] > 1 else "get_response"
+    )
+
+
+def get_response(state: BasicState, config: RunnableConfig) -> BasicStateUpdate:
+    pro_search_config = cast(ProSearchConfig, config["metadata"]["config"])
+    llm = pro_search_config.primary_llm
+    current_llm_call = state["last_llm_call"]
+    if current_llm_call is None:
+        raise ValueError("last_llm_call is None")
+    structured_response_format = pro_search_config.structured_response_format
+    response_handler_manager = state["response_handler_manager"]
+    # DEBUG: good breakpoint
+    stream = llm.stream(
+        # For tool calling LLMs, we want to insert the task prompt as part of this flow, this is because the LLM
+        # may choose to not call any tools and just generate the answer, in which case the task prompt is needed.
+        prompt=current_llm_call.prompt_builder.build(),
+        tools=[tool.tool_definition() for tool in current_llm_call.tools] or None,
+        tool_choice=(
+            "required"
+            if current_llm_call.tools and current_llm_call.force_use_tool.force_use
+            else None
+        ),
+        structured_response_format=structured_response_format,
+    )
+
+    for response in response_handler_manager.handle_llm_response(stream):
+        dispatch_custom_event(
+            "basic_response",
+            response,
+        )
+
+    next_call = response_handler_manager.next_llm_call(current_llm_call)
+    if next_call is not None:
+        final_search_results, displayed_search_results = SearchTool.get_search_result(
+            next_call
+        ) or ([], [])
+    else:
+        final_search_results, displayed_search_results = [], []
+
+    response_handler_manager.answer_handler.update(
+        (
+            final_search_results,
+            map_document_id_order(final_search_results),
+            map_document_id_order(displayed_search_results),
+        )
+    )
+    return BasicStateUpdate(
+        last_llm_call=next_call,
+        calls=state["calls"] + 1,
+    )
+
+
+if __name__ == "__main__":
+    pass

backend/onyx/agent_search/basic/states.py

@@ -0,0 +1,38 @@
+from typing import TypedDict
+
+from onyx.chat.llm_response_handler import LLMResponseHandlerManager
+from onyx.chat.prompt_builder.build import LLMCall
+
+## Update States
+
+
+## Graph Input State
+
+
+class BasicInput(TypedDict):
+    base_question: str
+    last_llm_call: LLMCall | None
+    response_handler_manager: LLMResponseHandlerManager
+    calls: int
+
+
+## Graph Output State
+
+
+class BasicOutput(TypedDict):
+    pass
+
+
+class BasicStateUpdate(TypedDict):
+    last_llm_call: LLMCall | None
+    calls: int
+
+
+## Graph State
+
+
+class BasicState(
+    BasicInput,
+    BasicOutput,
+):
+    pass

backend/onyx/agent_search/core_state.py

@@ -0,0 +1,20 @@
+from operator import add
+from typing import Annotated
+from typing import TypedDict
+
+
+class CoreState(TypedDict, total=False):
+    """
+    This is the core state that is shared across all subgraphs.
+    """
+
+    base_question: str
+    log_messages: Annotated[list[str], add]
+
+
+class SubgraphCoreState(TypedDict, total=False):
+    """
+    This is the core state that is shared across all subgraphs.
+    """
+
+    log_messages: Annotated[list[str], add]

backend/onyx/agent_search/db_operations.py

@@ -0,0 +1,66 @@
+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()
+    )

backend/onyx/agent_search/models.py

@@ -0,0 +1,57 @@
+from dataclasses import dataclass
+from uuid import UUID
+
+from pydantic import BaseModel
+from sqlalchemy.orm import Session
+
+from onyx.context.search.models import SearchRequest
+from onyx.llm.interfaces import LLM
+from onyx.llm.models import PreviousMessage
+from onyx.tools.tool_implementations.search.search_tool import SearchTool
+
+
+@dataclass
+class ProSearchConfig:
+    """
+    Configuration for the Pro Search feature.
+    """
+
+    # The search request that was used to generate the Pro Search
+    search_request: SearchRequest
+
+    primary_llm: LLM
+    fast_llm: LLM
+    search_tool: SearchTool
+    use_agentic_search: bool = False
+
+    # For persisting agent search data
+    chat_session_id: UUID | None = None
+
+    # The message ID of the user message that triggered the Pro Search
+    message_id: int | None = None
+
+    # Whether to persistence data for the Pro Search (turned off for testing)
+    use_persistence: bool = True
+
+    # The database session for the Pro Search
+    db_session: Session | None = None
+
+    # Whether to perform initial search to inform decomposition
+    perform_initial_search_path_decision: bool = False
+
+    # Whether to perform initial search to inform decomposition
+    perform_initial_search_decomposition: bool = False
+
+    # Whether to allow creation of refinement questions (and entity extraction, etc.)
+    allow_refinement: bool = False
+
+    # Message history for the current chat session
+    message_history: list[PreviousMessage] | None = None
+
+    structured_response_format: dict | None = None
+
+
+class AgentDocumentCitations(BaseModel):
+    document_id: str
+    document_title: str
+    link: str

backend/onyx/agent_search/pro_search_a/answer_initial_sub_question/edges.py

@@ -0,0 +1,26 @@
+from collections.abc import Hashable
+
+from langgraph.types import Send
+
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.states import (
+    AnswerQuestionInput,
+)
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import (
+    ExpandedRetrievalInput,
+)
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def send_to_expanded_retrieval(state: AnswerQuestionInput) -> Send | Hashable:
+    logger.debug("sending to expanded retrieval via edge")
+
+    return Send(
+        "initial_sub_question_expanded_retrieval",
+        ExpandedRetrievalInput(
+            question=state["question"],
+            base_search=False,
+            sub_question_id=state["question_id"],
+        ),
+    )

backend/onyx/agent_search/pro_search_a/answer_initial_sub_question/graph_builder.py

@@ -0,0 +1,125 @@
+from langgraph.graph import END
+from langgraph.graph import START
+from langgraph.graph import StateGraph
+
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.edges import (
+    send_to_expanded_retrieval,
+)
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.nodes.answer_check import (
+    answer_check,
+)
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.nodes.answer_generation import (
+    answer_generation,
+)
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.nodes.format_answer import (
+    format_answer,
+)
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.nodes.ingest_retrieval import (
+    ingest_retrieval,
+)
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.states import (
+    AnswerQuestionInput,
+)
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.states import (
+    AnswerQuestionOutput,
+)
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.states import (
+    AnswerQuestionState,
+)
+from onyx.agent_search.pro_search_a.expanded_retrieval.graph_builder import (
+    expanded_retrieval_graph_builder,
+)
+from onyx.agent_search.shared_graph_utils.utils import get_test_config
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def answer_query_graph_builder() -> StateGraph:
+    graph = StateGraph(
+        state_schema=AnswerQuestionState,
+        input=AnswerQuestionInput,
+        output=AnswerQuestionOutput,
+    )
+
+    ### Add nodes ###
+
+    expanded_retrieval = expanded_retrieval_graph_builder().compile()
+    graph.add_node(
+        node="initial_sub_question_expanded_retrieval",
+        action=expanded_retrieval,
+    )
+    graph.add_node(
+        node="answer_check",
+        action=answer_check,
+    )
+    graph.add_node(
+        node="answer_generation",
+        action=answer_generation,
+    )
+    graph.add_node(
+        node="format_answer",
+        action=format_answer,
+    )
+    graph.add_node(
+        node="ingest_retrieval",
+        action=ingest_retrieval,
+    )
+
+    ### Add edges ###
+
+    graph.add_conditional_edges(
+        source=START,
+        path=send_to_expanded_retrieval,
+        path_map=["initial_sub_question_expanded_retrieval"],
+    )
+    graph.add_edge(
+        start_key="initial_sub_question_expanded_retrieval",
+        end_key="ingest_retrieval",
+    )
+    graph.add_edge(
+        start_key="ingest_retrieval",
+        end_key="answer_generation",
+    )
+    graph.add_edge(
+        start_key="answer_generation",
+        end_key="answer_check",
+    )
+    graph.add_edge(
+        start_key="answer_check",
+        end_key="format_answer",
+    )
+    graph.add_edge(
+        start_key="format_answer",
+        end_key=END,
+    )
+
+    return graph
+
+
+if __name__ == "__main__":
+    from onyx.db.engine import get_session_context_manager
+    from onyx.llm.factory import get_default_llms
+    from onyx.context.search.models import SearchRequest
+
+    graph = answer_query_graph_builder()
+    compiled_graph = graph.compile()
+    primary_llm, fast_llm = get_default_llms()
+    search_request = SearchRequest(
+        query="what can you do with onyx or danswer?",
+    )
+    with get_session_context_manager() as db_session:
+        pro_search_config, search_tool = get_test_config(
+            db_session, primary_llm, fast_llm, search_request
+        )
+        inputs = AnswerQuestionInput(
+            question="what can you do with onyx?",
+            question_id="0_0",
+        )
+        for thing in compiled_graph.stream(
+            input=inputs,
+            config={"configurable": {"config": pro_search_config}},
+            # debug=True,
+            # subgraphs=True,
+        ):
+            logger.debug(thing)

backend/onyx/agent_search/pro_search_a/answer_initial_sub_question/models.py

@@ -0,0 +1,8 @@
+from pydantic import BaseModel
+
+
+### Models ###
+
+
+class AnswerRetrievalStats(BaseModel):
+    answer_retrieval_stats: dict[str, float | int]

backend/onyx/agent_search/pro_search_a/answer_initial_sub_question/nodes/answer_check.py

@@ -0,0 +1,45 @@
+from typing import cast
+
+from langchain_core.messages import HumanMessage
+from langchain_core.messages import merge_message_runs
+from langchain_core.runnables.config import RunnableConfig
+
+from onyx.agent_search.models import ProSearchConfig
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.states import (
+    AnswerQuestionState,
+)
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.states import (
+    QACheckUpdate,
+)
+from onyx.agent_search.shared_graph_utils.prompts import SUB_CHECK_NO
+from onyx.agent_search.shared_graph_utils.prompts import SUB_CHECK_PROMPT
+from onyx.agent_search.shared_graph_utils.prompts import UNKNOWN_ANSWER
+
+
+def answer_check(state: AnswerQuestionState, config: RunnableConfig) -> QACheckUpdate:
+    if state["answer"] == UNKNOWN_ANSWER:
+        return QACheckUpdate(
+            answer_quality=SUB_CHECK_NO,
+        )
+    msg = [
+        HumanMessage(
+            content=SUB_CHECK_PROMPT.format(
+                question=state["question"],
+                base_answer=state["answer"],
+            )
+        )
+    ]
+
+    pro_search_config = cast(ProSearchConfig, config["metadata"]["config"])
+    fast_llm = pro_search_config.fast_llm
+    response = list(
+        fast_llm.stream(
+            prompt=msg,
+        )
+    )
+
+    quality_str = merge_message_runs(response, chunk_separator="")[0].content
+
+    return QACheckUpdate(
+        answer_quality=quality_str,
+    )

backend/onyx/agent_search/pro_search_a/answer_initial_sub_question/nodes/answer_generation.py

@@ -0,0 +1,106 @@
+import datetime
+from typing import Any
+from typing import cast
+
+from langchain_core.callbacks.manager import dispatch_custom_event
+from langchain_core.messages import merge_message_runs
+from langchain_core.runnables.config import RunnableConfig
+
+from onyx.agent_search.models import ProSearchConfig
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.states import (
+    AnswerQuestionState,
+)
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.states import (
+    QAGenerationUpdate,
+)
+from onyx.agent_search.shared_graph_utils.agent_prompt_ops import (
+    build_sub_question_answer_prompt,
+)
+from onyx.agent_search.shared_graph_utils.prompts import ASSISTANT_SYSTEM_PROMPT_DEFAULT
+from onyx.agent_search.shared_graph_utils.prompts import ASSISTANT_SYSTEM_PROMPT_PERSONA
+from onyx.agent_search.shared_graph_utils.prompts import UNKNOWN_ANSWER
+from onyx.agent_search.shared_graph_utils.utils import get_persona_prompt
+from onyx.agent_search.shared_graph_utils.utils import parse_question_id
+from onyx.chat.models import AgentAnswerPiece
+from onyx.chat.models import StreamStopInfo
+from onyx.chat.models import StreamStopReason
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def answer_generation(
+    state: AnswerQuestionState, config: RunnableConfig
+) -> QAGenerationUpdate:
+    now_start = datetime.datetime.now()
+    logger.debug(f"--------{now_start}--------START ANSWER GENERATION---")
+
+    pro_search_config = cast(ProSearchConfig, config["metadata"]["config"])
+    question = state["question"]
+    docs = state["documents"]
+    level, question_nr = parse_question_id(state["question_id"])
+    persona_prompt = get_persona_prompt(pro_search_config.search_request.persona)
+
+    if len(docs) == 0:
+        answer_str = UNKNOWN_ANSWER
+        dispatch_custom_event(
+            "sub_answers",
+            AgentAnswerPiece(
+                answer_piece=answer_str,
+                level=level,
+                level_question_nr=question_nr,
+                answer_type="agent_sub_answer",
+            ),
+        )
+    else:
+        if len(persona_prompt) > 0:
+            persona_specification = ASSISTANT_SYSTEM_PROMPT_DEFAULT
+        else:
+            persona_specification = ASSISTANT_SYSTEM_PROMPT_PERSONA.format(
+                persona_prompt=persona_prompt
+            )
+
+        logger.debug(f"Number of verified retrieval docs: {len(docs)}")
+
+        fast_llm = pro_search_config.fast_llm
+        msg = build_sub_question_answer_prompt(
+            question=question,
+            original_question=pro_search_config.search_request.query,
+            docs=docs,
+            persona_specification=persona_specification,
+            config=fast_llm.config,
+        )
+
+        response: list[str | list[str | dict[str, Any]]] = []
+        for message in fast_llm.stream(
+            prompt=msg,
+        ):
+            # TODO: in principle, the answer here COULD contain images, but we don't support that yet
+            content = message.content
+            if not isinstance(content, str):
+                raise ValueError(
+                    f"Expected content to be a string, but got {type(content)}"
+                )
+            dispatch_custom_event(
+                "sub_answers",
+                AgentAnswerPiece(
+                    answer_piece=content,
+                    level=level,
+                    level_question_nr=question_nr,
+                    answer_type="agent_sub_answer",
+                ),
+            )
+            response.append(content)
+
+        answer_str = merge_message_runs(response, chunk_separator="")[0].content
+
+    stop_event = StreamStopInfo(
+        stop_reason=StreamStopReason.FINISHED,
+        level=level,
+        level_question_nr=question_nr,
+    )
+    dispatch_custom_event("sub_answer_finished", stop_event)
+
+    return QAGenerationUpdate(
+        answer=answer_str,
+    )

backend/onyx/agent_search/pro_search_a/answer_initial_sub_question/nodes/format_answer.py

@@ -0,0 +1,25 @@
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.states import (
+    AnswerQuestionOutput,
+)
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.states import (
+    AnswerQuestionState,
+)
+from onyx.agent_search.shared_graph_utils.models import (
+    QuestionAnswerResults,
+)
+
+
+def format_answer(state: AnswerQuestionState) -> AnswerQuestionOutput:
+    return AnswerQuestionOutput(
+        answer_results=[
+            QuestionAnswerResults(
+                question=state["question"],
+                question_id=state["question_id"],
+                quality=state.get("answer_quality", "No"),
+                answer=state["answer"],
+                expanded_retrieval_results=state["expanded_retrieval_results"],
+                documents=state["documents"],
+                sub_question_retrieval_stats=state["sub_question_retrieval_stats"],
+            )
+        ],
+    )

backend/onyx/agent_search/pro_search_a/answer_initial_sub_question/nodes/ingest_retrieval.py

@@ -0,0 +1,23 @@
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.states import (
+    RetrievalIngestionUpdate,
+)
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import (
+    ExpandedRetrievalOutput,
+)
+from onyx.agent_search.shared_graph_utils.models import AgentChunkStats
+
+
+def ingest_retrieval(state: ExpandedRetrievalOutput) -> RetrievalIngestionUpdate:
+    sub_question_retrieval_stats = state[
+        "expanded_retrieval_result"
+    ].sub_question_retrieval_stats
+    if sub_question_retrieval_stats is None:
+        sub_question_retrieval_stats = [AgentChunkStats()]
+
+    return RetrievalIngestionUpdate(
+        expanded_retrieval_results=state[
+            "expanded_retrieval_result"
+        ].expanded_queries_results,
+        documents=state["expanded_retrieval_result"].all_documents,
+        sub_question_retrieval_stats=sub_question_retrieval_stats,
+    )

backend/onyx/agent_search/pro_search_a/answer_initial_sub_question/states.py

@@ -0,0 +1,63 @@
+from operator import add
+from typing import Annotated
+from typing import TypedDict
+
+from onyx.agent_search.core_state import SubgraphCoreState
+from onyx.agent_search.shared_graph_utils.models import AgentChunkStats
+from onyx.agent_search.shared_graph_utils.models import QueryResult
+from onyx.agent_search.shared_graph_utils.models import (
+    QuestionAnswerResults,
+)
+from onyx.agent_search.shared_graph_utils.operators import dedup_inference_sections
+from onyx.context.search.models import InferenceSection
+
+
+## Update States
+class QACheckUpdate(TypedDict):
+    answer_quality: str
+
+
+class QAGenerationUpdate(TypedDict):
+    answer: str
+    # answer_stat: AnswerStats
+
+
+class RetrievalIngestionUpdate(TypedDict):
+    expanded_retrieval_results: list[QueryResult]
+    documents: Annotated[list[InferenceSection], dedup_inference_sections]
+    sub_question_retrieval_stats: AgentChunkStats
+
+
+## Graph Input State
+
+
+class AnswerQuestionInput(SubgraphCoreState):
+    question: str
+    question_id: str  # 0_0 is original question, everything else is <level>_<question_num>.
+    # level 0 is original question and first decomposition, level 1 is follow up, etc
+    # question_num is a unique number per original question per level.
+
+
+## Graph State
+
+
+class AnswerQuestionState(
+    AnswerQuestionInput,
+    QAGenerationUpdate,
+    QACheckUpdate,
+    RetrievalIngestionUpdate,
+):
+    pass
+
+
+## Graph Output State
+
+
+class AnswerQuestionOutput(TypedDict):
+    """
+    This is a list of results even though each call of this subgraph only returns one result.
+    This is because if we parallelize the answer query subgraph, there will be multiple
+      results in a list so the add operator is used to add them together.
+    """
+
+    answer_results: Annotated[list[QuestionAnswerResults], add]

backend/onyx/agent_search/pro_search_a/answer_refinement_sub_question/edges.py

@@ -0,0 +1,26 @@
+from collections.abc import Hashable
+
+from langgraph.types import Send
+
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.states import (
+    AnswerQuestionInput,
+)
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import (
+    ExpandedRetrievalInput,
+)
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def send_to_expanded_refined_retrieval(state: AnswerQuestionInput) -> Send | Hashable:
+    logger.debug("sending to expanded retrieval for follow up question via edge")
+
+    return Send(
+        "refined_sub_question_expanded_retrieval",
+        ExpandedRetrievalInput(
+            question=state["question"],
+            sub_question_id=state["question_id"],
+            base_search=False,
+        ),
+    )

backend/onyx/agent_search/pro_search_a/answer_refinement_sub_question/graph_builder.py

@@ -0,0 +1,122 @@
+from langgraph.graph import END
+from langgraph.graph import START
+from langgraph.graph import StateGraph
+
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.nodes.answer_check import (
+    answer_check,
+)
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.nodes.answer_generation import (
+    answer_generation,
+)
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.nodes.format_answer import (
+    format_answer,
+)
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.nodes.ingest_retrieval import (
+    ingest_retrieval,
+)
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.states import (
+    AnswerQuestionInput,
+)
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.states import (
+    AnswerQuestionOutput,
+)
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.states import (
+    AnswerQuestionState,
+)
+from onyx.agent_search.pro_search_a.answer_refinement_sub_question.edges import (
+    send_to_expanded_refined_retrieval,
+)
+from onyx.agent_search.pro_search_a.expanded_retrieval.graph_builder import (
+    expanded_retrieval_graph_builder,
+)
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def answer_refined_query_graph_builder() -> StateGraph:
+    graph = StateGraph(
+        state_schema=AnswerQuestionState,
+        input=AnswerQuestionInput,
+        output=AnswerQuestionOutput,
+    )
+
+    ### Add nodes ###
+
+    expanded_retrieval = expanded_retrieval_graph_builder().compile()
+    graph.add_node(
+        node="refined_sub_question_expanded_retrieval",
+        action=expanded_retrieval,
+    )
+    graph.add_node(
+        node="refined_sub_answer_check",
+        action=answer_check,
+    )
+    graph.add_node(
+        node="refined_sub_answer_generation",
+        action=answer_generation,
+    )
+    graph.add_node(
+        node="format_refined_sub_answer",
+        action=format_answer,
+    )
+    graph.add_node(
+        node="ingest_refined_retrieval",
+        action=ingest_retrieval,
+    )
+
+    ### Add edges ###
+
+    graph.add_conditional_edges(
+        source=START,
+        path=send_to_expanded_refined_retrieval,
+        path_map=["refined_sub_question_expanded_retrieval"],
+    )
+    graph.add_edge(
+        start_key="refined_sub_question_expanded_retrieval",
+        end_key="ingest_refined_retrieval",
+    )
+    graph.add_edge(
+        start_key="ingest_refined_retrieval",
+        end_key="refined_sub_answer_generation",
+    )
+    graph.add_edge(
+        start_key="refined_sub_answer_generation",
+        end_key="refined_sub_answer_check",
+    )
+    graph.add_edge(
+        start_key="refined_sub_answer_check",
+        end_key="format_refined_sub_answer",
+    )
+    graph.add_edge(
+        start_key="format_refined_sub_answer",
+        end_key=END,
+    )
+
+    return graph
+
+
+if __name__ == "__main__":
+    from onyx.db.engine import get_session_context_manager
+    from onyx.llm.factory import get_default_llms
+    from onyx.context.search.models import SearchRequest
+
+    graph = answer_refined_query_graph_builder()
+    compiled_graph = graph.compile()
+    primary_llm, fast_llm = get_default_llms()
+    search_request = SearchRequest(
+        query="what can you do with onyx or danswer?",
+    )
+    with get_session_context_manager() as db_session:
+        inputs = AnswerQuestionInput(
+            question="what can you do with onyx?",
+            question_id="0_0",
+        )
+        for thing in compiled_graph.stream(
+            input=inputs,
+            # debug=True,
+            # subgraphs=True,
+        ):
+            logger.debug(thing)
+        # output = compiled_graph.invoke(inputs)
+        #  logger.debug(output)

backend/onyx/agent_search/pro_search_a/answer_refinement_sub_question/models.py

@@ -0,0 +1,19 @@
+from pydantic import BaseModel
+
+from onyx.agent_search.shared_graph_utils.models import AgentChunkStats
+from onyx.context.search.models import InferenceSection
+
+### Models ###
+
+
+class AnswerRetrievalStats(BaseModel):
+    answer_retrieval_stats: dict[str, float | int]
+
+
+class QuestionAnswerResults(BaseModel):
+    question: str
+    answer: str
+    quality: str
+    # expanded_retrieval_results: list[QueryResult]
+    documents: list[InferenceSection]
+    sub_question_retrieval_stats: AgentChunkStats

backend/onyx/agent_search/pro_search_a/base_raw_search/graph_builder.py

@@ -0,0 +1,70 @@
+from langgraph.graph import END
+from langgraph.graph import START
+from langgraph.graph import StateGraph
+
+from onyx.agent_search.pro_search_a.base_raw_search.nodes.format_raw_search_results import (
+    format_raw_search_results,
+)
+from onyx.agent_search.pro_search_a.base_raw_search.nodes.generate_raw_search_data import (
+    generate_raw_search_data,
+)
+from onyx.agent_search.pro_search_a.base_raw_search.states import BaseRawSearchInput
+from onyx.agent_search.pro_search_a.base_raw_search.states import BaseRawSearchOutput
+from onyx.agent_search.pro_search_a.base_raw_search.states import BaseRawSearchState
+from onyx.agent_search.pro_search_a.expanded_retrieval.graph_builder import (
+    expanded_retrieval_graph_builder,
+)
+
+
+def base_raw_search_graph_builder() -> StateGraph:
+    graph = StateGraph(
+        state_schema=BaseRawSearchState,
+        input=BaseRawSearchInput,
+        output=BaseRawSearchOutput,
+    )
+
+    ### Add nodes ###
+
+    graph.add_node(
+        node="generate_raw_search_data",
+        action=generate_raw_search_data,
+    )
+
+    expanded_retrieval = expanded_retrieval_graph_builder().compile()
+    graph.add_node(
+        node="expanded_retrieval_base_search",
+        action=expanded_retrieval,
+    )
+    graph.add_node(
+        node="format_raw_search_results",
+        action=format_raw_search_results,
+    )
+
+    ### Add edges ###
+
+    graph.add_edge(start_key=START, end_key="generate_raw_search_data")
+
+    graph.add_edge(
+        start_key="generate_raw_search_data",
+        end_key="expanded_retrieval_base_search",
+    )
+    graph.add_edge(
+        start_key="expanded_retrieval_base_search",
+        end_key="format_raw_search_results",
+    )
+
+    # graph.add_edge(
+    #     start_key="expanded_retrieval_base_search",
+    #     end_key=END,
+    # )
+
+    graph.add_edge(
+        start_key="format_raw_search_results",
+        end_key=END,
+    )
+
+    return graph
+
+
+if __name__ == "__main__":
+    pass

backend/onyx/agent_search/pro_search_a/base_raw_search/models.py

@@ -0,0 +1,20 @@
+from pydantic import BaseModel
+
+from onyx.agent_search.shared_graph_utils.models import AgentChunkStats
+from onyx.agent_search.shared_graph_utils.models import QueryResult
+from onyx.context.search.models import InferenceSection
+
+### Models ###
+
+
+class AnswerRetrievalStats(BaseModel):
+    answer_retrieval_stats: dict[str, float | int]
+
+
+class QuestionAnswerResults(BaseModel):
+    question: str
+    answer: str
+    quality: str
+    expanded_retrieval_results: list[QueryResult]
+    documents: list[InferenceSection]
+    sub_question_retrieval_stats: list[AgentChunkStats]

backend/onyx/agent_search/pro_search_a/base_raw_search/nodes/format_raw_search_results.py

@@ -0,0 +1,16 @@
+from onyx.agent_search.pro_search_a.base_raw_search.states import BaseRawSearchOutput
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import (
+    ExpandedRetrievalOutput,
+)
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def format_raw_search_results(state: ExpandedRetrievalOutput) -> BaseRawSearchOutput:
+    logger.debug("format_raw_search_results")
+    return BaseRawSearchOutput(
+        base_expanded_retrieval_result=state["expanded_retrieval_result"],
+        # base_retrieval_results=[state["expanded_retrieval_result"]],
+        # base_search_documents=[],
+    )

backend/onyx/agent_search/pro_search_a/base_raw_search/nodes/generate_raw_search_data.py

@@ -0,0 +1,24 @@
+from typing import cast
+
+from langchain_core.runnables.config import RunnableConfig
+
+from onyx.agent_search.core_state import CoreState
+from onyx.agent_search.models import ProSearchConfig
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import (
+    ExpandedRetrievalInput,
+)
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def generate_raw_search_data(
+    state: CoreState, config: RunnableConfig
+) -> ExpandedRetrievalInput:
+    logger.debug("generate_raw_search_data")
+    pro_search_config = cast(ProSearchConfig, config["metadata"]["config"])
+    return ExpandedRetrievalInput(
+        question=pro_search_config.search_request.query,
+        base_search=True,
+        sub_question_id=None,  # This graph is always and only used for the original question
+    )

backend/onyx/agent_search/pro_search_a/base_raw_search/states.py

@@ -0,0 +1,43 @@
+from typing import TypedDict
+
+from onyx.agent_search.pro_search_a.expanded_retrieval.models import (
+    ExpandedRetrievalResult,
+)
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import (
+    ExpandedRetrievalInput,
+)
+
+
+## Update States
+
+
+## Graph Input State
+
+
+class BaseRawSearchInput(ExpandedRetrievalInput):
+    pass
+
+
+## Graph Output State
+
+
+class BaseRawSearchOutput(TypedDict):
+    """
+    This is a list of results even though each call of this subgraph only returns one result.
+    This is because if we parallelize the answer query subgraph, there will be multiple
+      results in a list so the add operator is used to add them together.
+    """
+
+    # base_search_documents: Annotated[list[InferenceSection], dedup_inference_sections]
+    # base_retrieval_results: Annotated[list[ExpandedRetrievalResult], add]
+    base_expanded_retrieval_result: ExpandedRetrievalResult
+
+
+## Graph State
+
+
+class BaseRawSearchState(
+    BaseRawSearchInput,
+    BaseRawSearchOutput,
+):
+    pass

backend/onyx/agent_search/pro_search_a/expanded_retrieval/edges.py

@@ -0,0 +1,32 @@
+from collections.abc import Hashable
+from typing import cast
+
+from langchain_core.runnables.config import RunnableConfig
+from langgraph.types import Send
+
+from onyx.agent_search.models import ProSearchConfig
+from onyx.agent_search.pro_search_a.expanded_retrieval.nodes import RetrievalInput
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import (
+    ExpandedRetrievalState,
+)
+
+
+def parallel_retrieval_edge(
+    state: ExpandedRetrievalState, config: RunnableConfig
+) -> list[Send | Hashable]:
+    pro_search_config = cast(ProSearchConfig, config["metadata"]["config"])
+    question = state.get("question", pro_search_config.search_request.query)
+
+    query_expansions = state.get("expanded_queries", []) + [question]
+    return [
+        Send(
+            "doc_retrieval",
+            RetrievalInput(
+                query_to_retrieve=query,
+                question=question,
+                base_search=False,
+                sub_question_id=state.get("sub_question_id"),
+            ),
+        )
+        for query in query_expansions
+    ]

backend/onyx/agent_search/pro_search_a/expanded_retrieval/graph_builder.py

@@ -0,0 +1,122 @@
+from langgraph.graph import END
+from langgraph.graph import START
+from langgraph.graph import StateGraph
+
+from onyx.agent_search.pro_search_a.expanded_retrieval.edges import (
+    parallel_retrieval_edge,
+)
+from onyx.agent_search.pro_search_a.expanded_retrieval.nodes import doc_reranking
+from onyx.agent_search.pro_search_a.expanded_retrieval.nodes import doc_retrieval
+from onyx.agent_search.pro_search_a.expanded_retrieval.nodes import doc_verification
+from onyx.agent_search.pro_search_a.expanded_retrieval.nodes import expand_queries
+from onyx.agent_search.pro_search_a.expanded_retrieval.nodes import format_results
+from onyx.agent_search.pro_search_a.expanded_retrieval.nodes import verification_kickoff
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import (
+    ExpandedRetrievalInput,
+)
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import (
+    ExpandedRetrievalOutput,
+)
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import (
+    ExpandedRetrievalState,
+)
+from onyx.agent_search.shared_graph_utils.utils import get_test_config
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def expanded_retrieval_graph_builder() -> StateGraph:
+    graph = StateGraph(
+        state_schema=ExpandedRetrievalState,
+        input=ExpandedRetrievalInput,
+        output=ExpandedRetrievalOutput,
+    )
+
+    ### Add nodes ###
+
+    graph.add_node(
+        node="expand_queries",
+        action=expand_queries,
+    )
+
+    graph.add_node(
+        node="doc_retrieval",
+        action=doc_retrieval,
+    )
+    graph.add_node(
+        node="verification_kickoff",
+        action=verification_kickoff,
+    )
+    graph.add_node(
+        node="doc_verification",
+        action=doc_verification,
+    )
+    graph.add_node(
+        node="doc_reranking",
+        action=doc_reranking,
+    )
+    graph.add_node(
+        node="format_results",
+        action=format_results,
+    )
+
+    ### Add edges ###
+    graph.add_edge(
+        start_key=START,
+        end_key="expand_queries",
+    )
+
+    graph.add_conditional_edges(
+        source="expand_queries",
+        path=parallel_retrieval_edge,
+        path_map=["doc_retrieval"],
+    )
+    graph.add_edge(
+        start_key="doc_retrieval",
+        end_key="verification_kickoff",
+    )
+    graph.add_edge(
+        start_key="doc_verification",
+        end_key="doc_reranking",
+    )
+    graph.add_edge(
+        start_key="doc_reranking",
+        end_key="format_results",
+    )
+    graph.add_edge(
+        start_key="format_results",
+        end_key=END,
+    )
+
+    return graph
+
+
+if __name__ == "__main__":
+    from onyx.db.engine import get_session_context_manager
+    from onyx.llm.factory import get_default_llms
+    from onyx.context.search.models import SearchRequest
+
+    graph = expanded_retrieval_graph_builder()
+    compiled_graph = graph.compile()
+    primary_llm, fast_llm = get_default_llms()
+    search_request = SearchRequest(
+        query="what can you do with onyx or danswer?",
+    )
+
+    with get_session_context_manager() as db_session:
+        pro_search_config, search_tool = get_test_config(
+            db_session, primary_llm, fast_llm, search_request
+        )
+        inputs = ExpandedRetrievalInput(
+            question="what can you do with onyx?",
+            base_search=False,
+            sub_question_id=None,
+        )
+        for thing in compiled_graph.stream(
+            input=inputs,
+            config={"configurable": {"config": pro_search_config}},
+            # debug=True,
+            subgraphs=True,
+        ):
+            logger.debug(thing)

backend/onyx/agent_search/pro_search_a/expanded_retrieval/models.py

@@ -0,0 +1,11 @@
+from pydantic import BaseModel
+
+from onyx.agent_search.shared_graph_utils.models import AgentChunkStats
+from onyx.agent_search.shared_graph_utils.models import QueryResult
+from onyx.context.search.models import InferenceSection
+
+
+class ExpandedRetrievalResult(BaseModel):
+    expanded_queries_results: list[QueryResult]
+    all_documents: list[InferenceSection]
+    sub_question_retrieval_stats: AgentChunkStats

backend/onyx/agent_search/pro_search_a/expanded_retrieval/nodes.py

@@ -0,0 +1,431 @@
+from collections import defaultdict
+from collections.abc import Callable
+from typing import cast
+from typing import Literal
+
+import numpy as np
+from langchain_core.callbacks.manager import dispatch_custom_event
+from langchain_core.messages import HumanMessage
+from langchain_core.messages import merge_message_runs
+from langchain_core.runnables.config import RunnableConfig
+from langgraph.types import Command
+from langgraph.types import Send
+
+from onyx.agent_search.models import ProSearchConfig
+from onyx.agent_search.pro_search_a.expanded_retrieval.models import (
+    ExpandedRetrievalResult,
+)
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import DocRerankingUpdate
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import DocRetrievalUpdate
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import (
+    DocVerificationInput,
+)
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import (
+    DocVerificationUpdate,
+)
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import (
+    ExpandedRetrievalInput,
+)
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import (
+    ExpandedRetrievalState,
+)
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import (
+    ExpandedRetrievalUpdate,
+)
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import (
+    QueryExpansionUpdate,
+)
+from onyx.agent_search.pro_search_a.expanded_retrieval.states import RetrievalInput
+from onyx.agent_search.shared_graph_utils.agent_prompt_ops import trim_prompt_piece
+from onyx.agent_search.shared_graph_utils.calculations import get_fit_scores
+from onyx.agent_search.shared_graph_utils.models import AgentChunkStats
+from onyx.agent_search.shared_graph_utils.models import QueryResult
+from onyx.agent_search.shared_graph_utils.models import RetrievalFitStats
+from onyx.agent_search.shared_graph_utils.prompts import REWRITE_PROMPT_MULTI_ORIGINAL
+from onyx.agent_search.shared_graph_utils.prompts import VERIFIER_PROMPT
+from onyx.agent_search.shared_graph_utils.utils import dispatch_separated
+from onyx.agent_search.shared_graph_utils.utils import parse_question_id
+from onyx.chat.models import ExtendedToolResponse
+from onyx.chat.models import SubQueryPiece
+from onyx.configs.dev_configs import AGENT_MAX_QUERY_RETRIEVAL_RESULTS
+from onyx.configs.dev_configs import AGENT_RERANKING_MAX_QUERY_RETRIEVAL_RESULTS
+from onyx.configs.dev_configs import AGENT_RERANKING_STATS
+from onyx.configs.dev_configs import AGENT_RETRIEVAL_STATS
+from onyx.context.search.models import InferenceSection
+from onyx.context.search.models import SearchRequest
+from onyx.context.search.pipeline import retrieval_preprocessing
+from onyx.context.search.postprocessing.postprocessing import rerank_sections
+from onyx.db.engine import get_session_context_manager
+from onyx.tools.models import SearchQueryInfo
+from onyx.tools.tool_implementations.search.search_tool import (
+    SEARCH_RESPONSE_SUMMARY_ID,
+)
+from onyx.tools.tool_implementations.search.search_tool import SearchResponseSummary
+from onyx.tools.tool_implementations.search.search_tool import yield_search_responses
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def dispatch_subquery(level: int, question_nr: int) -> Callable[[str, int], None]:
+    def helper(token: str, num: int) -> None:
+        dispatch_custom_event(
+            "subqueries",
+            SubQueryPiece(
+                sub_query=token,
+                level=level,
+                level_question_nr=question_nr,
+                query_id=num,
+            ),
+        )
+
+    return helper
+
+
+def expand_queries(
+    state: ExpandedRetrievalInput, config: RunnableConfig
+) -> QueryExpansionUpdate:
+    # Sometimes we want to expand the original question, sometimes we want to expand a sub-question.
+    # When we are running this node on the original question, no question is explictly passed in.
+    # Instead, we use the original question from the search request.
+    pro_search_config = cast(ProSearchConfig, config["metadata"]["config"])
+    question = state.get("question", pro_search_config.search_request.query)
+    llm = pro_search_config.fast_llm
+    chat_session_id = pro_search_config.chat_session_id
+    sub_question_id = state.get("sub_question_id")
+    if sub_question_id is None:
+        level, question_nr = 0, 0
+    else:
+        level, question_nr = parse_question_id(sub_question_id)
+
+    if chat_session_id is None:
+        raise ValueError("chat_session_id must be provided for agent search")
+
+    msg = [
+        HumanMessage(
+            content=REWRITE_PROMPT_MULTI_ORIGINAL.format(question=question),
+        )
+    ]
+
+    llm_response_list = dispatch_separated(
+        llm.stream(prompt=msg), dispatch_subquery(level, question_nr)
+    )
+
+    llm_response = merge_message_runs(llm_response_list, chunk_separator="")[0].content
+
+    rewritten_queries = llm_response.split("\n")
+
+    return QueryExpansionUpdate(
+        expanded_queries=rewritten_queries,
+    )
+
+
+def doc_retrieval(state: RetrievalInput, config: RunnableConfig) -> DocRetrievalUpdate:
+    """
+    Retrieve documents
+
+    Args:
+        state (RetrievalInput): Primary state + the query to retrieve
+        config (RunnableConfig): Configuration containing ProSearchConfig
+
+    Updates:
+        expanded_retrieval_results: list[ExpandedRetrievalResult]
+        retrieved_documents: list[InferenceSection]
+    """
+    query_to_retrieve = state["query_to_retrieve"]
+    pro_search_config = cast(ProSearchConfig, config["metadata"]["config"])
+    search_tool = pro_search_config.search_tool
+
+    retrieved_docs: list[InferenceSection] = []
+    if not query_to_retrieve.strip():
+        logger.warning("Empty query, skipping retrieval")
+        return DocRetrievalUpdate(
+            expanded_retrieval_results=[],
+            retrieved_documents=[],
+        )
+
+    query_info = None
+    # new db session to avoid concurrency issues
+    with get_session_context_manager() as db_session:
+        for tool_response in search_tool.run(
+            query=query_to_retrieve,
+            force_no_rerank=True,
+            alternate_db_session=db_session,
+        ):
+            # get retrieved docs to send to the rest of the graph
+            if tool_response.id == SEARCH_RESPONSE_SUMMARY_ID:
+                response = cast(SearchResponseSummary, tool_response.response)
+                retrieved_docs = response.top_sections
+                query_info = SearchQueryInfo(
+                    predicted_search=response.predicted_search,
+                    final_filters=response.final_filters,
+                    recency_bias_multiplier=response.recency_bias_multiplier,
+                )
+                break
+
+    retrieved_docs = retrieved_docs[:AGENT_MAX_QUERY_RETRIEVAL_RESULTS]
+    pre_rerank_docs = retrieved_docs
+    if search_tool.search_pipeline is not None:
+        pre_rerank_docs = (
+            search_tool.search_pipeline._retrieved_sections or retrieved_docs
+        )
+
+    if AGENT_RETRIEVAL_STATS:
+        fit_scores = get_fit_scores(
+            pre_rerank_docs,
+            retrieved_docs,
+        )
+    else:
+        fit_scores = None
+
+    expanded_retrieval_result = QueryResult(
+        query=query_to_retrieve,
+        search_results=retrieved_docs,
+        stats=fit_scores,
+        query_info=query_info,
+    )
+    return DocRetrievalUpdate(
+        expanded_retrieval_results=[expanded_retrieval_result],
+        retrieved_documents=retrieved_docs,
+    )
+
+
+def verification_kickoff(
+    state: ExpandedRetrievalState,
+    config: RunnableConfig,
+) -> Command[Literal["doc_verification"]]:
+    documents = state["retrieved_documents"]
+    pro_search_config = cast(ProSearchConfig, config["metadata"]["config"])
+    verification_question = state.get(
+        "question", pro_search_config.search_request.query
+    )
+    sub_question_id = state.get("sub_question_id")
+    return Command(
+        update={},
+        goto=[
+            Send(
+                node="doc_verification",
+                arg=DocVerificationInput(
+                    doc_to_verify=doc,
+                    question=verification_question,
+                    base_search=False,
+                    sub_question_id=sub_question_id,
+                ),
+            )
+            for doc in documents
+        ],
+    )
+
+
+def doc_verification(
+    state: DocVerificationInput, config: RunnableConfig
+) -> DocVerificationUpdate:
+    """
+    Check whether the document is relevant for the original user question
+
+    Args:
+        state (DocVerificationInput): The current state
+        config (RunnableConfig): Configuration containing ProSearchConfig
+
+    Updates:
+        verified_documents: list[InferenceSection]
+    """
+
+    question = state["question"]
+    doc_to_verify = state["doc_to_verify"]
+    document_content = doc_to_verify.combined_content
+
+    pro_search_config = cast(ProSearchConfig, config["metadata"]["config"])
+    fast_llm = pro_search_config.fast_llm
+
+    document_content = trim_prompt_piece(
+        fast_llm.config, document_content, VERIFIER_PROMPT + question
+    )
+
+    msg = [
+        HumanMessage(
+            content=VERIFIER_PROMPT.format(
+                question=question, document_content=document_content
+            )
+        )
+    ]
+
+    response = fast_llm.invoke(msg)
+
+    verified_documents = []
+    if isinstance(response.content, str) and "yes" in response.content.lower():
+        verified_documents.append(doc_to_verify)
+
+    return DocVerificationUpdate(
+        verified_documents=verified_documents,
+    )
+
+
+def doc_reranking(
+    state: ExpandedRetrievalState, config: RunnableConfig
+) -> DocRerankingUpdate:
+    verified_documents = state["verified_documents"]
+
+    # Rerank post retrieval and verification. First, create a search query
+    # then create the list of reranked sections
+
+    pro_search_config = cast(ProSearchConfig, config["metadata"]["config"])
+    question = state.get("question", pro_search_config.search_request.query)
+    with get_session_context_manager() as db_session:
+        _search_query = retrieval_preprocessing(
+            search_request=SearchRequest(query=question),
+            user=pro_search_config.search_tool.user,  # bit of a hack
+            llm=pro_search_config.fast_llm,
+            db_session=db_session,
+        )
+
+    # skip section filtering
+
+    if (
+        _search_query.rerank_settings
+        and _search_query.rerank_settings.rerank_model_name
+        and _search_query.rerank_settings.num_rerank > 0
+    ):
+        reranked_documents = rerank_sections(
+            _search_query,
+            verified_documents,
+        )
+    else:
+        logger.warning("No reranking settings found, using unranked documents")
+        reranked_documents = verified_documents
+
+    if AGENT_RERANKING_STATS:
+        fit_scores = get_fit_scores(verified_documents, reranked_documents)
+    else:
+        fit_scores = RetrievalFitStats(fit_score_lift=0, rerank_effect=0, fit_scores={})
+
+    # TODO: stream deduped docs here, or decide to use search tool ranking/verification
+
+    return DocRerankingUpdate(
+        reranked_documents=[
+            doc for doc in reranked_documents if type(doc) == InferenceSection
+        ][:AGENT_RERANKING_MAX_QUERY_RETRIEVAL_RESULTS],
+        sub_question_retrieval_stats=fit_scores,
+    )
+
+
+def _calculate_sub_question_retrieval_stats(
+    verified_documents: list[InferenceSection],
+    expanded_retrieval_results: list[QueryResult],
+) -> AgentChunkStats:
+    chunk_scores: dict[str, dict[str, list[int | float]]] = defaultdict(
+        lambda: defaultdict(list)
+    )
+
+    for expanded_retrieval_result in expanded_retrieval_results:
+        for doc in expanded_retrieval_result.search_results:
+            doc_chunk_id = f"{doc.center_chunk.document_id}_{doc.center_chunk.chunk_id}"
+            if doc.center_chunk.score is not None:
+                chunk_scores[doc_chunk_id]["score"].append(doc.center_chunk.score)
+
+    verified_doc_chunk_ids = [
+        f"{verified_document.center_chunk.document_id}_{verified_document.center_chunk.chunk_id}"
+        for verified_document in verified_documents
+    ]
+    dismissed_doc_chunk_ids = []
+
+    raw_chunk_stats_counts: dict[str, int] = defaultdict(int)
+    raw_chunk_stats_scores: dict[str, float] = defaultdict(float)
+    for doc_chunk_id, chunk_data in chunk_scores.items():
+        if doc_chunk_id in verified_doc_chunk_ids:
+            raw_chunk_stats_counts["verified_count"] += 1
+
+            valid_chunk_scores = [
+                score for score in chunk_data["score"] if score is not None
+            ]
+            raw_chunk_stats_scores["verified_scores"] += float(
+                np.mean(valid_chunk_scores)
+            )
+        else:
+            raw_chunk_stats_counts["rejected_count"] += 1
+            valid_chunk_scores = [
+                score for score in chunk_data["score"] if score is not None
+            ]
+            raw_chunk_stats_scores["rejected_scores"] += float(
+                np.mean(valid_chunk_scores)
+            )
+            dismissed_doc_chunk_ids.append(doc_chunk_id)
+
+    if raw_chunk_stats_counts["verified_count"] == 0:
+        verified_avg_scores = 0.0
+    else:
+        verified_avg_scores = raw_chunk_stats_scores["verified_scores"] / float(
+            raw_chunk_stats_counts["verified_count"]
+        )
+
+    rejected_scores = raw_chunk_stats_scores.get("rejected_scores", None)
+    if rejected_scores is not None:
+        rejected_avg_scores = rejected_scores / float(
+            raw_chunk_stats_counts["rejected_count"]
+        )
+    else:
+        rejected_avg_scores = None
+
+    chunk_stats = AgentChunkStats(
+        verified_count=raw_chunk_stats_counts["verified_count"],
+        verified_avg_scores=verified_avg_scores,
+        rejected_count=raw_chunk_stats_counts["rejected_count"],
+        rejected_avg_scores=rejected_avg_scores,
+        verified_doc_chunk_ids=verified_doc_chunk_ids,
+        dismissed_doc_chunk_ids=dismissed_doc_chunk_ids,
+    )
+
+    return chunk_stats
+
+
+def format_results(
+    state: ExpandedRetrievalState, config: RunnableConfig
+) -> ExpandedRetrievalUpdate:
+    level, question_nr = parse_question_id(state.get("sub_question_id") or "0_0")
+    query_infos = [
+        result.query_info
+        for result in state["expanded_retrieval_results"]
+        if result.query_info is not None
+    ]
+    if len(query_infos) == 0:
+        raise ValueError("No query info found")
+
+    pro_search_config = cast(ProSearchConfig, config["metadata"]["config"])
+    # main question docs will be sent later after aggregation and deduping with sub-question docs
+    if not (level == 0 and question_nr == 0):
+        for tool_response in yield_search_responses(
+            query=state["question"],
+            reranked_sections=state[
+                "retrieved_documents"
+            ],  # TODO: rename params. this one is supposed to be the sections pre-merging
+            final_context_sections=state["reranked_documents"],
+            search_query_info=query_infos[0],  # TODO: handle differing query infos?
+            get_section_relevance=lambda: None,  # TODO: add relevance
+            search_tool=pro_search_config.search_tool,
+        ):
+            dispatch_custom_event(
+                "tool_response",
+                ExtendedToolResponse(
+                    id=tool_response.id,
+                    response=tool_response.response,
+                    level=level,
+                    level_question_nr=question_nr,
+                ),
+            )
+    sub_question_retrieval_stats = _calculate_sub_question_retrieval_stats(
+        verified_documents=state["verified_documents"],
+        expanded_retrieval_results=state["expanded_retrieval_results"],
+    )
+
+    if sub_question_retrieval_stats is None:
+        sub_question_retrieval_stats = AgentChunkStats()
+    # else:
+    #    sub_question_retrieval_stats = [sub_question_retrieval_stats]
+
+    return ExpandedRetrievalUpdate(
+        expanded_retrieval_result=ExpandedRetrievalResult(
+            expanded_queries_results=state["expanded_retrieval_results"],
+            all_documents=state["reranked_documents"],
+            sub_question_retrieval_stats=sub_question_retrieval_stats,
+        ),
+    )

backend/onyx/agent_search/pro_search_a/expanded_retrieval/states.py

@@ -0,0 +1,82 @@
+from operator import add
+from typing import Annotated
+from typing import TypedDict
+
+from onyx.agent_search.core_state import SubgraphCoreState
+from onyx.agent_search.pro_search_a.expanded_retrieval.models import (
+    ExpandedRetrievalResult,
+)
+from onyx.agent_search.shared_graph_utils.models import QueryResult
+from onyx.agent_search.shared_graph_utils.models import RetrievalFitStats
+from onyx.agent_search.shared_graph_utils.operators import dedup_inference_sections
+from onyx.context.search.models import InferenceSection
+
+
+### States ###
+
+## Graph Input State
+
+
+class ExpandedRetrievalInput(SubgraphCoreState):
+    question: str
+    base_search: bool
+    sub_question_id: str | None
+
+
+## Update/Return States
+
+
+class QueryExpansionUpdate(TypedDict):
+    expanded_queries: list[str]
+
+
+class DocVerificationUpdate(TypedDict):
+    verified_documents: Annotated[list[InferenceSection], dedup_inference_sections]
+
+
+class DocRetrievalUpdate(TypedDict):
+    expanded_retrieval_results: Annotated[list[QueryResult], add]
+    retrieved_documents: Annotated[list[InferenceSection], dedup_inference_sections]
+
+
+class DocRerankingUpdate(TypedDict):
+    reranked_documents: Annotated[list[InferenceSection], dedup_inference_sections]
+    sub_question_retrieval_stats: RetrievalFitStats | None
+
+
+class ExpandedRetrievalUpdate(TypedDict):
+    expanded_retrieval_result: ExpandedRetrievalResult
+
+
+## Graph Output State
+
+
+class ExpandedRetrievalOutput(TypedDict):
+    expanded_retrieval_result: ExpandedRetrievalResult
+    base_expanded_retrieval_result: ExpandedRetrievalResult
+
+
+## Graph State
+
+
+class ExpandedRetrievalState(
+    # This includes the core state
+    ExpandedRetrievalInput,
+    QueryExpansionUpdate,
+    DocRetrievalUpdate,
+    DocVerificationUpdate,
+    DocRerankingUpdate,
+    ExpandedRetrievalOutput,
+):
+    pass
+
+
+## Conditional Input States
+
+
+class DocVerificationInput(ExpandedRetrievalInput):
+    doc_to_verify: InferenceSection
+
+
+class RetrievalInput(ExpandedRetrievalInput):
+    query_to_retrieve: str

backend/onyx/agent_search/pro_search_a/main/edges.py

@@ -0,0 +1,89 @@
+from collections.abc import Hashable
+from typing import Literal
+
+from langgraph.types import Send
+
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.states import (
+    AnswerQuestionInput,
+)
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.states import (
+    AnswerQuestionOutput,
+)
+from onyx.agent_search.pro_search_a.main.states import MainState
+from onyx.agent_search.pro_search_a.main.states import RequireRefinedAnswerUpdate
+from onyx.agent_search.shared_graph_utils.utils import make_question_id
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def parallelize_initial_sub_question_answering(
+    state: MainState,
+) -> list[Send | Hashable]:
+    if len(state["initial_decomp_questions"]) > 0:
+        # sub_question_record_ids = [subq_record.id for subq_record in state["sub_question_records"]]
+        # if len(state["sub_question_records"]) == 0:
+        #     if state["config"].use_persistence:
+        #         raise ValueError("No sub-questions found for initial decompozed questions")
+        #     else:
+        #         # in this case, we are doing retrieval on the original question.
+        #         # to make all the logic consistent, we create a new sub-question
+        #         # with the same content as the original question
+        #         sub_question_record_ids = [1] * len(state["initial_decomp_questions"])
+
+        return [
+            Send(
+                "answer_query_subgraph",
+                AnswerQuestionInput(
+                    question=question,
+                    question_id=make_question_id(0, question_nr + 1),
+                ),
+            )
+            for question_nr, question in enumerate(state["initial_decomp_questions"])
+        ]
+
+    else:
+        return [
+            Send(
+                "ingest_answers",
+                AnswerQuestionOutput(
+                    answer_results=[],
+                ),
+            )
+        ]
+
+
+# Define the function that determines whether to continue or not
+def continue_to_refined_answer_or_end(
+    state: RequireRefinedAnswerUpdate,
+) -> Literal["refined_sub_question_creation", "logging_node"]:
+    if state["require_refined_answer"]:
+        return "refined_sub_question_creation"
+    else:
+        return "logging_node"
+
+
+def parallelize_refined_sub_question_answering(
+    state: MainState,
+) -> list[Send | Hashable]:
+    if len(state["refined_sub_questions"]) > 0:
+        return [
+            Send(
+                "answer_refined_question",
+                AnswerQuestionInput(
+                    question=question_data.sub_question,
+                    question_id=make_question_id(1, question_nr),
+                ),
+            )
+            for question_nr, question_data in state["refined_sub_questions"].items()
+        ]
+
+    else:
+        return [
+            Send(
+                "ingest_refined_sub_answers",
+                AnswerQuestionOutput(
+                    answer_results=[],
+                ),
+            )
+        ]

backend/onyx/agent_search/pro_search_a/main/graph_builder.py

@@ -0,0 +1,309 @@
+from langgraph.graph import END
+from langgraph.graph import START
+from langgraph.graph import StateGraph
+
+from onyx.agent_search.pro_search_a.answer_initial_sub_question.graph_builder import (
+    answer_query_graph_builder,
+)
+from onyx.agent_search.pro_search_a.answer_refinement_sub_question.graph_builder import (
+    answer_refined_query_graph_builder,
+)
+from onyx.agent_search.pro_search_a.base_raw_search.graph_builder import (
+    base_raw_search_graph_builder,
+)
+from onyx.agent_search.pro_search_a.main.edges import continue_to_refined_answer_or_end
+from onyx.agent_search.pro_search_a.main.edges import (
+    parallelize_initial_sub_question_answering,
+)
+from onyx.agent_search.pro_search_a.main.edges import (
+    parallelize_refined_sub_question_answering,
+)
+from onyx.agent_search.pro_search_a.main.nodes import agent_logging
+from onyx.agent_search.pro_search_a.main.nodes import agent_path_decision
+from onyx.agent_search.pro_search_a.main.nodes import agent_path_routing
+from onyx.agent_search.pro_search_a.main.nodes import agent_search_start
+from onyx.agent_search.pro_search_a.main.nodes import direct_llm_handling
+from onyx.agent_search.pro_search_a.main.nodes import entity_term_extraction_llm
+from onyx.agent_search.pro_search_a.main.nodes import generate_initial_answer
+from onyx.agent_search.pro_search_a.main.nodes import generate_refined_answer
+from onyx.agent_search.pro_search_a.main.nodes import ingest_initial_base_retrieval
+from onyx.agent_search.pro_search_a.main.nodes import (
+    ingest_initial_sub_question_answers,
+)
+from onyx.agent_search.pro_search_a.main.nodes import ingest_refined_answers
+from onyx.agent_search.pro_search_a.main.nodes import initial_answer_quality_check
+from onyx.agent_search.pro_search_a.main.nodes import initial_sub_question_creation
+from onyx.agent_search.pro_search_a.main.nodes import refined_answer_decision
+from onyx.agent_search.pro_search_a.main.nodes import refined_sub_question_creation
+from onyx.agent_search.pro_search_a.main.states import MainInput
+from onyx.agent_search.pro_search_a.main.states import MainState
+from onyx.agent_search.shared_graph_utils.utils import get_test_config
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+test_mode = False
+
+
+def main_graph_builder(test_mode: bool = False) -> StateGraph:
+    graph = StateGraph(
+        state_schema=MainState,
+        input=MainInput,
+    )
+
+    graph.add_node(
+        node="agent_path_decision",
+        action=agent_path_decision,
+    )
+
+    graph.add_node(
+        node="agent_path_routing",
+        action=agent_path_routing,
+    )
+
+    graph.add_node(
+        node="LLM",
+        action=direct_llm_handling,
+    )
+
+    graph.add_node(
+        node="agent_search_start",
+        action=agent_search_start,
+    )
+
+    graph.add_node(
+        node="initial_sub_question_creation",
+        action=initial_sub_question_creation,
+    )
+    answer_query_subgraph = answer_query_graph_builder().compile()
+    graph.add_node(
+        node="answer_query_subgraph",
+        action=answer_query_subgraph,
+    )
+
+    base_raw_search_subgraph = base_raw_search_graph_builder().compile()
+    graph.add_node(
+        node="base_raw_search_subgraph",
+        action=base_raw_search_subgraph,
+    )
+
+    # refined_answer_subgraph = refined_answers_graph_builder().compile()
+    # graph.add_node(
+    #     node="refined_answer_subgraph",
+    #     action=refined_answer_subgraph,
+    # )
+
+    graph.add_node(
+        node="refined_sub_question_creation",
+        action=refined_sub_question_creation,
+    )
+
+    answer_refined_question = answer_refined_query_graph_builder().compile()
+    graph.add_node(
+        node="answer_refined_question",
+        action=answer_refined_question,
+    )
+
+    graph.add_node(
+        node="ingest_refined_answers",
+        action=ingest_refined_answers,
+    )
+
+    graph.add_node(
+        node="generate_refined_answer",
+        action=generate_refined_answer,
+    )
+
+    # graph.add_node(
+    #     node="check_refined_answer",
+    #     action=check_refined_answer,
+    # )
+
+    graph.add_node(
+        node="ingest_initial_retrieval",
+        action=ingest_initial_base_retrieval,
+    )
+    graph.add_node(
+        node="ingest_initial_sub_question_answers",
+        action=ingest_initial_sub_question_answers,
+    )
+    graph.add_node(
+        node="generate_initial_answer",
+        action=generate_initial_answer,
+    )
+
+    graph.add_node(
+        node="initial_answer_quality_check",
+        action=initial_answer_quality_check,
+    )
+
+    graph.add_node(
+        node="entity_term_extraction_llm",
+        action=entity_term_extraction_llm,
+    )
+    graph.add_node(
+        node="refined_answer_decision",
+        action=refined_answer_decision,
+    )
+
+    graph.add_node(
+        node="logging_node",
+        action=agent_logging,
+    )
+    # if test_mode:
+    #     graph.add_node(
+    #         node="generate_initial_base_answer",
+    #         action=generate_initial_base_answer,
+    #     )
+
+    ### Add edges ###
+
+    # raph.add_edge(start_key=START, end_key="base_raw_search_subgraph")
+
+    graph.add_edge(
+        start_key=START,
+        end_key="agent_path_decision",
+    )
+
+    graph.add_edge(
+        start_key="agent_path_decision",
+        end_key="agent_path_routing",
+    )
+
+    graph.add_edge(
+        start_key="agent_search_start",
+        end_key="base_raw_search_subgraph",
+    )
+
+    graph.add_edge(
+        start_key="agent_search_start",
+        end_key="initial_sub_question_creation",
+    )
+
+    graph.add_edge(
+        start_key="base_raw_search_subgraph",
+        end_key="ingest_initial_retrieval",
+    )
+
+    graph.add_edge(
+        start_key="LLM",
+        end_key=END,
+    )
+
+    # graph.add_edge(
+    #     start_key=START,
+    #     end_key="initial_sub_question_creation",
+    # )
+
+    graph.add_conditional_edges(
+        source="initial_sub_question_creation",
+        path=parallelize_initial_sub_question_answering,
+        path_map=["answer_query_subgraph"],
+    )
+    graph.add_edge(
+        start_key="answer_query_subgraph",
+        end_key="ingest_initial_sub_question_answers",
+    )
+
+    graph.add_edge(
+        start_key=["ingest_initial_sub_question_answers", "ingest_initial_retrieval"],
+        end_key="generate_initial_answer",
+    )
+
+    graph.add_edge(
+        start_key="generate_initial_answer",
+        end_key="entity_term_extraction_llm",
+    )
+
+    graph.add_edge(
+        start_key="generate_initial_answer",
+        end_key="initial_answer_quality_check",
+    )
+
+    graph.add_edge(
+        start_key=["initial_answer_quality_check", "entity_term_extraction_llm"],
+        end_key="refined_answer_decision",
+    )
+
+    graph.add_conditional_edges(
+        source="refined_answer_decision",
+        path=continue_to_refined_answer_or_end,
+        path_map=["refined_sub_question_creation", "logging_node"],
+    )
+
+    graph.add_conditional_edges(
+        source="refined_sub_question_creation",  # DONE
+        path=parallelize_refined_sub_question_answering,
+        path_map=["answer_refined_question"],
+    )
+    graph.add_edge(
+        start_key="answer_refined_question",  # HERE
+        end_key="ingest_refined_answers",
+    )
+
+    graph.add_edge(
+        start_key="ingest_refined_answers",
+        end_key="generate_refined_answer",
+    )
+
+    # graph.add_conditional_edges(
+    #     source="refined_answer_decision",
+    #     path=continue_to_refined_answer_or_end,
+    #     path_map=["refined_answer_subgraph", END],
+    # )
+
+    # graph.add_edge(
+    #     start_key="refined_answer_subgraph",
+    #     end_key="generate_refined_answer",
+    # )
+
+    graph.add_edge(
+        start_key="generate_refined_answer",
+        end_key="logging_node",
+    )
+
+    graph.add_edge(
+        start_key="logging_node",
+        end_key=END,
+    )
+
+    # graph.add_edge(
+    #     start_key="generate_refined_answer",
+    #     end_key="check_refined_answer",
+    # )
+
+    # graph.add_edge(
+    #     start_key="check_refined_answer",
+    #     end_key=END,
+    # )
+
+    return graph
+
+
+if __name__ == "__main__":
+    pass
+
+    from onyx.db.engine import get_session_context_manager
+    from onyx.llm.factory import get_default_llms
+    from onyx.context.search.models import SearchRequest
+
+    graph = main_graph_builder()
+    compiled_graph = graph.compile()
+    primary_llm, fast_llm = get_default_llms()
+
+    with get_session_context_manager() as db_session:
+        search_request = SearchRequest(query="Who created Excel?")
+        pro_search_config, search_tool = get_test_config(
+            db_session, primary_llm, fast_llm, search_request
+        )
+
+        inputs = MainInput()
+
+        for thing in compiled_graph.stream(
+            input=inputs,
+            config={"configurable": {"config": pro_search_config}},
+            # stream_mode="debug",
+            # debug=True,
+            subgraphs=True,
+        ):
+            logger.debug(thing)

backend/onyx/agent_search/pro_search_a/main/models.py

@@ -0,0 +1,36 @@
+from pydantic import BaseModel
+
+
+class FollowUpSubQuestion(BaseModel):
+    sub_question: str
+    sub_question_id: str
+    verified: bool
+    answered: bool
+    answer: str
+
+
+class AgentTimings(BaseModel):
+    base_duration__s: float | None
+    refined_duration__s: float | None
+    full_duration__s: float | None
+
+
+class AgentBaseMetrics(BaseModel):
+    num_verified_documents_total: int | None
+    num_verified_documents_core: int | None
+    verified_avg_score_core: float | None
+    num_verified_documents_base: int | float | None
+    verified_avg_score_base: float | None
+    base_doc_boost_factor: float | None
+    support_boost_factor: float | None
+    duration__s: float | None
+
+
+class AgentRefinedMetrics(BaseModel):
+    refined_doc_boost_factor: float | None
+    refined_question_boost_factor: float | None
+    duration__s: float | None
+
+
+class AgentAdditionalMetrics(BaseModel):
+    pass

backend/onyx/agent_search/pro_search_a/main/states.py

@@ -0,0 +1,165 @@
+from datetime import datetime
+from operator import add
+from typing import Annotated
+from typing import TypedDict
+
+from onyx.agent_search.core_state import CoreState
+from onyx.agent_search.pro_search_a.expanded_retrieval.models import (
+    ExpandedRetrievalResult,
+)
+from onyx.agent_search.pro_search_a.main.models import AgentBaseMetrics
+from onyx.agent_search.pro_search_a.main.models import AgentRefinedMetrics
+from onyx.agent_search.pro_search_a.main.models import FollowUpSubQuestion
+from onyx.agent_search.shared_graph_utils.models import AgentChunkStats
+from onyx.agent_search.shared_graph_utils.models import EntityRelationshipTermExtraction
+from onyx.agent_search.shared_graph_utils.models import InitialAgentResultStats
+from onyx.agent_search.shared_graph_utils.models import QueryResult
+from onyx.agent_search.shared_graph_utils.models import (
+    QuestionAnswerResults,
+)
+from onyx.agent_search.shared_graph_utils.models import RefinedAgentStats
+from onyx.agent_search.shared_graph_utils.operators import dedup_inference_sections
+from onyx.agent_search.shared_graph_utils.operators import dedup_question_answer_results
+from onyx.context.search.models import InferenceSection
+
+
+### States ###
+
+## Update States
+
+
+class RefinedAgentStartStats(TypedDict):
+    agent_refined_start_time: datetime | None
+
+
+class RefinedAgentEndStats(TypedDict):
+    agent_refined_end_time: datetime | None
+    agent_refined_metrics: AgentRefinedMetrics
+
+
+class BaseDecompUpdateBase(TypedDict):
+    agent_start_time: datetime
+    initial_decomp_questions: list[str]
+
+
+class RoutingDecisionBase(TypedDict):
+    routing: str
+    sample_doc_str: str
+
+
+class RoutingDecision(RoutingDecisionBase):
+    log_messages: list[str]
+
+
+class BaseDecompUpdate(
+    RefinedAgentStartStats, RefinedAgentEndStats, BaseDecompUpdateBase
+):
+    pass
+
+
+class InitialAnswerBASEUpdate(TypedDict):
+    initial_base_answer: str
+
+
+class InitialAnswerUpdateBase(TypedDict):
+    initial_answer: str
+    initial_agent_stats: InitialAgentResultStats | None
+    generated_sub_questions: list[str]
+    agent_base_end_time: datetime
+    agent_base_metrics: AgentBaseMetrics | None
+
+
+class InitialAnswerUpdate(InitialAnswerUpdateBase):
+    log_messages: list[str]
+
+
+class RefinedAnswerUpdateBase(TypedDict):
+    refined_answer: str
+    refined_agent_stats: RefinedAgentStats | None
+    refined_answer_quality: bool
+
+
+class RefinedAnswerUpdate(RefinedAgentEndStats, RefinedAnswerUpdateBase):
+    pass
+
+
+class InitialAnswerQualityUpdate(TypedDict):
+    initial_answer_quality: bool
+
+
+class RequireRefinedAnswerUpdate(TypedDict):
+    require_refined_answer: bool
+
+
+class DecompAnswersUpdate(TypedDict):
+    documents: Annotated[list[InferenceSection], dedup_inference_sections]
+    decomp_answer_results: Annotated[
+        list[QuestionAnswerResults], dedup_question_answer_results
+    ]
+
+
+class FollowUpDecompAnswersUpdate(TypedDict):
+    refined_documents: Annotated[list[InferenceSection], dedup_inference_sections]
+    refined_decomp_answer_results: Annotated[list[QuestionAnswerResults], add]
+
+
+class ExpandedRetrievalUpdate(TypedDict):
+    all_original_question_documents: Annotated[
+        list[InferenceSection], dedup_inference_sections
+    ]
+    original_question_retrieval_results: list[QueryResult]
+    original_question_retrieval_stats: AgentChunkStats
+
+
+class EntityTermExtractionUpdate(TypedDict):
+    entity_retlation_term_extractions: EntityRelationshipTermExtraction
+
+
+class FollowUpSubQuestionsUpdateBase(TypedDict):
+    refined_sub_questions: dict[int, FollowUpSubQuestion]
+
+
+class FollowUpSubQuestionsUpdate(
+    RefinedAgentStartStats, FollowUpSubQuestionsUpdateBase
+):
+    pass
+
+
+## Graph Input State
+## Graph Input State
+
+
+class MainInput(CoreState):
+    pass
+
+
+## Graph State
+
+
+class MainState(
+    # This includes the core state
+    MainInput,
+    BaseDecompUpdateBase,
+    InitialAnswerUpdateBase,
+    InitialAnswerBASEUpdate,
+    DecompAnswersUpdate,
+    ExpandedRetrievalUpdate,
+    EntityTermExtractionUpdate,
+    InitialAnswerQualityUpdate,
+    RequireRefinedAnswerUpdate,
+    FollowUpSubQuestionsUpdateBase,
+    FollowUpDecompAnswersUpdate,
+    RefinedAnswerUpdateBase,
+    RefinedAgentStartStats,
+    RefinedAgentEndStats,
+    RoutingDecisionBase,
+):
+    # expanded_retrieval_result: Annotated[list[ExpandedRetrievalResult], add]
+    base_raw_search_result: Annotated[list[ExpandedRetrievalResult], add]
+
+
+## Graph Output State - presently not used
+
+
+class MainOutput(TypedDict):
+    pass

backend/onyx/agent_search/pro_search_b/answer_initial_sub_question/edges.py

@@ -0,0 +1,26 @@
+from collections.abc import Hashable
+
+from langgraph.types import Send
+
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.states import (
+    AnswerQuestionInput,
+)
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import (
+    ExpandedRetrievalInput,
+)
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def send_to_expanded_retrieval(state: AnswerQuestionInput) -> Send | Hashable:
+    logger.debug("sending to expanded retrieval via edge")
+
+    return Send(
+        "initial_sub_question_expanded_retrieval",
+        ExpandedRetrievalInput(
+            question=state["question"],
+            base_search=False,
+            sub_question_id=state["question_id"],
+        ),
+    )

backend/onyx/agent_search/pro_search_b/answer_initial_sub_question/graph_builder.py

@@ -0,0 +1,129 @@
+from langgraph.graph import END
+from langgraph.graph import START
+from langgraph.graph import StateGraph
+
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.edges import (
+    send_to_expanded_retrieval,
+)
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.nodes.answer_check import (
+    answer_check,
+)
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.nodes.answer_generation import (
+    answer_generation,
+)
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.nodes.format_answer import (
+    format_answer,
+)
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.nodes.ingest_retrieval import (
+    ingest_retrieval,
+)
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.states import (
+    AnswerQuestionInput,
+)
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.states import (
+    AnswerQuestionOutput,
+)
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.states import (
+    AnswerQuestionState,
+)
+from onyx.agent_search.pro_search_b.expanded_retrieval.graph_builder import (
+    expanded_retrieval_graph_builder,
+)
+from onyx.agent_search.shared_graph_utils.utils import get_test_config
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def answer_query_graph_builder() -> StateGraph:
+    graph = StateGraph(
+        state_schema=AnswerQuestionState,
+        input=AnswerQuestionInput,
+        output=AnswerQuestionOutput,
+    )
+
+    ### Add nodes ###
+
+    expanded_retrieval = expanded_retrieval_graph_builder().compile()
+    graph.add_node(
+        node="initial_sub_question_expanded_retrieval",
+        action=expanded_retrieval,
+    )
+    graph.add_node(
+        node="answer_check",
+        action=answer_check,
+    )
+    graph.add_node(
+        node="answer_generation",
+        action=answer_generation,
+    )
+    graph.add_node(
+        node="format_answer",
+        action=format_answer,
+    )
+    graph.add_node(
+        node="ingest_retrieval",
+        action=ingest_retrieval,
+    )
+
+    ### Add edges ###
+
+    graph.add_conditional_edges(
+        source=START,
+        path=send_to_expanded_retrieval,
+        path_map=["initial_sub_question_expanded_retrieval"],
+    )
+    graph.add_edge(
+        start_key="initial_sub_question_expanded_retrieval",
+        end_key="ingest_retrieval",
+    )
+    graph.add_edge(
+        start_key="ingest_retrieval",
+        end_key="answer_generation",
+    )
+    graph.add_edge(
+        start_key="answer_generation",
+        end_key="answer_check",
+    )
+    graph.add_edge(
+        start_key="answer_check",
+        end_key="format_answer",
+    )
+    graph.add_edge(
+        start_key="format_answer",
+        end_key=END,
+    )
+
+    return graph
+
+
+if __name__ == "__main__":
+    from onyx.db.engine import get_session_context_manager
+    from onyx.llm.factory import get_default_llms
+    from onyx.context.search.models import SearchRequest
+
+    graph = answer_query_graph_builder()
+    compiled_graph = graph.compile()
+    primary_llm, fast_llm = get_default_llms()
+    search_request = SearchRequest(
+        query="what can you do with onyx or danswer?",
+    )
+    with get_session_context_manager() as db_session:
+        pro_search_config, search_tool = get_test_config(
+            db_session, primary_llm, fast_llm, search_request
+        )
+        inputs = AnswerQuestionInput(
+            question="what can you do with onyx?",
+            subgraph_fast_llm=fast_llm,
+            subgraph_primary_llm=primary_llm,
+            subgraph_config=pro_search_config,
+            subgraph_search_tool=search_tool,
+            subgraph_db_session=db_session,
+            question_id="0_0",
+        )
+        for thing in compiled_graph.stream(
+            input=inputs,
+            # debug=True,
+            # subgraphs=True,
+        ):
+            logger.debug(thing)

backend/onyx/agent_search/pro_search_b/answer_initial_sub_question/models.py

@@ -0,0 +1,8 @@
+from pydantic import BaseModel
+
+
+### Models ###
+
+
+class AnswerRetrievalStats(BaseModel):
+    answer_retrieval_stats: dict[str, float | int]

backend/onyx/agent_search/pro_search_b/answer_initial_sub_question/nodes/answer_check.py

@@ -0,0 +1,14 @@
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.states import (
+    AnswerQuestionState,
+)
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.states import (
+    QACheckUpdate,
+)
+
+
+def answer_check(state: AnswerQuestionState) -> QACheckUpdate:
+    quality_str = "yes"
+
+    return QACheckUpdate(
+        answer_quality=quality_str,
+    )

backend/onyx/agent_search/pro_search_b/answer_initial_sub_question/nodes/answer_generation.py

@@ -0,0 +1,41 @@
+import datetime
+
+from langchain_core.callbacks.manager import dispatch_custom_event
+
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.states import (
+    AnswerQuestionState,
+)
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.states import (
+    QAGenerationUpdate,
+)
+from onyx.agent_search.shared_graph_utils.utils import get_persona_prompt
+from onyx.agent_search.shared_graph_utils.utils import parse_question_id
+from onyx.chat.models import AgentAnswerPiece
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def answer_generation(state: AnswerQuestionState) -> QAGenerationUpdate:
+    now_start = datetime.datetime.now()
+    logger.debug(f"--------{now_start}--------START ANSWER GENERATION---")
+
+    state["question"]
+    state["documents"]
+    level, question_nr = parse_question_id(state["question_id"])
+    get_persona_prompt(state["subgraph_config"].search_request.persona)
+
+    dispatch_custom_event(
+        "sub_answers",
+        AgentAnswerPiece(
+            answer_piece="",
+            level=level,
+            level_question_nr=question_nr,
+            answer_type="agent_sub_answer",
+        ),
+    )
+    answer_str = ""
+
+    return QAGenerationUpdate(
+        answer=answer_str,
+    )

backend/onyx/agent_search/pro_search_b/answer_initial_sub_question/nodes/format_answer.py

@@ -0,0 +1,25 @@
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.states import (
+    AnswerQuestionOutput,
+)
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.states import (
+    AnswerQuestionState,
+)
+from onyx.agent_search.shared_graph_utils.models import (
+    QuestionAnswerResults,
+)
+
+
+def format_answer(state: AnswerQuestionState) -> AnswerQuestionOutput:
+    return AnswerQuestionOutput(
+        answer_results=[
+            QuestionAnswerResults(
+                question=state["question"],
+                question_id=state["question_id"],
+                quality=state.get("answer_quality", "No"),
+                answer=state["answer"],
+                expanded_retrieval_results=state["expanded_retrieval_results"],
+                documents=state["documents"],
+                sub_question_retrieval_stats=state["sub_question_retrieval_stats"],
+            )
+        ],
+    )

backend/onyx/agent_search/pro_search_b/answer_initial_sub_question/nodes/ingest_retrieval.py

@@ -0,0 +1,23 @@
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.states import (
+    RetrievalIngestionUpdate,
+)
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import (
+    ExpandedRetrievalOutput,
+)
+from onyx.agent_search.shared_graph_utils.models import AgentChunkStats
+
+
+def ingest_retrieval(state: ExpandedRetrievalOutput) -> RetrievalIngestionUpdate:
+    sub_question_retrieval_stats = state[
+        "expanded_retrieval_result"
+    ].sub_question_retrieval_stats
+    if sub_question_retrieval_stats is None:
+        sub_question_retrieval_stats = [AgentChunkStats()]
+
+    return RetrievalIngestionUpdate(
+        expanded_retrieval_results=state[
+            "expanded_retrieval_result"
+        ].expanded_queries_results,
+        documents=state["expanded_retrieval_result"].all_documents,
+        sub_question_retrieval_stats=sub_question_retrieval_stats,
+    )

backend/onyx/agent_search/pro_search_b/answer_initial_sub_question/states.py

@@ -0,0 +1,63 @@
+from operator import add
+from typing import Annotated
+from typing import TypedDict
+
+from onyx.agent_search.core_state import SubgraphCoreState
+from onyx.agent_search.pro_search_b.expanded_retrieval.models import QueryResult
+from onyx.agent_search.shared_graph_utils.models import AgentChunkStats
+from onyx.agent_search.shared_graph_utils.models import (
+    QuestionAnswerResults,
+)
+from onyx.agent_search.shared_graph_utils.operators import dedup_inference_sections
+from onyx.context.search.models import InferenceSection
+
+
+## Update States
+class QACheckUpdate(TypedDict):
+    answer_quality: str
+
+
+class QAGenerationUpdate(TypedDict):
+    answer: str
+    # answer_stat: AnswerStats
+
+
+class RetrievalIngestionUpdate(TypedDict):
+    expanded_retrieval_results: list[QueryResult]
+    documents: Annotated[list[InferenceSection], dedup_inference_sections]
+    sub_question_retrieval_stats: AgentChunkStats
+
+
+## Graph Input State
+
+
+class AnswerQuestionInput(SubgraphCoreState):
+    question: str
+    question_id: str  # 0_0 is original question, everything else is <level>_<question_num>.
+    # level 0 is original question and first decomposition, level 1 is follow up, etc
+    # question_num is a unique number per original question per level.
+
+
+## Graph State
+
+
+class AnswerQuestionState(
+    AnswerQuestionInput,
+    QAGenerationUpdate,
+    QACheckUpdate,
+    RetrievalIngestionUpdate,
+):
+    pass
+
+
+## Graph Output State
+
+
+class AnswerQuestionOutput(TypedDict):
+    """
+    This is a list of results even though each call of this subgraph only returns one result.
+    This is because if we parallelize the answer query subgraph, there will be multiple
+      results in a list so the add operator is used to add them together.
+    """
+
+    answer_results: Annotated[list[QuestionAnswerResults], add]

backend/onyx/agent_search/pro_search_b/answer_refinement_sub_question/edges.py

@@ -0,0 +1,26 @@
+from collections.abc import Hashable
+
+from langgraph.types import Send
+
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.states import (
+    AnswerQuestionInput,
+)
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import (
+    ExpandedRetrievalInput,
+)
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def send_to_expanded_refined_retrieval(state: AnswerQuestionInput) -> Send | Hashable:
+    logger.debug("sending to expanded retrieval for follow up question via edge")
+
+    return Send(
+        "refined_sub_question_expanded_retrieval",
+        ExpandedRetrievalInput(
+            question=state["question"],
+            sub_question_id=state["question_id"],
+            base_search=False,
+        ),
+    )

backend/onyx/agent_search/pro_search_b/answer_refinement_sub_question/graph_builder.py

@@ -0,0 +1,122 @@
+from langgraph.graph import END
+from langgraph.graph import START
+from langgraph.graph import StateGraph
+
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.nodes.answer_check import (
+    answer_check,
+)
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.nodes.answer_generation import (
+    answer_generation,
+)
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.nodes.format_answer import (
+    format_answer,
+)
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.nodes.ingest_retrieval import (
+    ingest_retrieval,
+)
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.states import (
+    AnswerQuestionInput,
+)
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.states import (
+    AnswerQuestionOutput,
+)
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.states import (
+    AnswerQuestionState,
+)
+from onyx.agent_search.pro_search_b.answer_refinement_sub_question.edges import (
+    send_to_expanded_refined_retrieval,
+)
+from onyx.agent_search.pro_search_b.expanded_retrieval.graph_builder import (
+    expanded_retrieval_graph_builder,
+)
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def answer_refined_query_graph_builder() -> StateGraph:
+    graph = StateGraph(
+        state_schema=AnswerQuestionState,
+        input=AnswerQuestionInput,
+        output=AnswerQuestionOutput,
+    )
+
+    ### Add nodes ###
+
+    expanded_retrieval = expanded_retrieval_graph_builder().compile()
+    graph.add_node(
+        node="refined_sub_question_expanded_retrieval",
+        action=expanded_retrieval,
+    )
+    graph.add_node(
+        node="refined_sub_answer_check",
+        action=answer_check,
+    )
+    graph.add_node(
+        node="refined_sub_answer_generation",
+        action=answer_generation,
+    )
+    graph.add_node(
+        node="format_refined_sub_answer",
+        action=format_answer,
+    )
+    graph.add_node(
+        node="ingest_refined_retrieval",
+        action=ingest_retrieval,
+    )
+
+    ### Add edges ###
+
+    graph.add_conditional_edges(
+        source=START,
+        path=send_to_expanded_refined_retrieval,
+        path_map=["refined_sub_question_expanded_retrieval"],
+    )
+    graph.add_edge(
+        start_key="refined_sub_question_expanded_retrieval",
+        end_key="ingest_refined_retrieval",
+    )
+    graph.add_edge(
+        start_key="ingest_refined_retrieval",
+        end_key="refined_sub_answer_generation",
+    )
+    graph.add_edge(
+        start_key="refined_sub_answer_generation",
+        end_key="refined_sub_answer_check",
+    )
+    graph.add_edge(
+        start_key="refined_sub_answer_check",
+        end_key="format_refined_sub_answer",
+    )
+    graph.add_edge(
+        start_key="format_refined_sub_answer",
+        end_key=END,
+    )
+
+    return graph
+
+
+if __name__ == "__main__":
+    from onyx.db.engine import get_session_context_manager
+    from onyx.llm.factory import get_default_llms
+    from onyx.context.search.models import SearchRequest
+
+    graph = answer_refined_query_graph_builder()
+    compiled_graph = graph.compile()
+    primary_llm, fast_llm = get_default_llms()
+    search_request = SearchRequest(
+        query="what can you do with onyx or danswer?",
+    )
+    with get_session_context_manager() as db_session:
+        inputs = AnswerQuestionInput(
+            question="what can you do with onyx?",
+            question_id="0_0",
+        )
+        for thing in compiled_graph.stream(
+            input=inputs,
+            # debug=True,
+            # subgraphs=True,
+        ):
+            logger.debug(thing)
+        # output = compiled_graph.invoke(inputs)
+        #  logger.debug(output)

backend/onyx/agent_search/pro_search_b/answer_refinement_sub_question/models.py

@@ -0,0 +1,19 @@
+from pydantic import BaseModel
+
+from onyx.agent_search.shared_graph_utils.models import AgentChunkStats
+from onyx.context.search.models import InferenceSection
+
+### Models ###
+
+
+class AnswerRetrievalStats(BaseModel):
+    answer_retrieval_stats: dict[str, float | int]
+
+
+class QuestionAnswerResults(BaseModel):
+    question: str
+    answer: str
+    quality: str
+    # expanded_retrieval_results: list[QueryResult]
+    documents: list[InferenceSection]
+    sub_question_retrieval_stats: AgentChunkStats

backend/onyx/agent_search/pro_search_b/base_raw_search/graph_builder.py

@@ -0,0 +1,70 @@
+from langgraph.graph import END
+from langgraph.graph import START
+from langgraph.graph import StateGraph
+
+from onyx.agent_search.pro_search_b.base_raw_search.nodes.format_raw_search_results import (
+    format_raw_search_results,
+)
+from onyx.agent_search.pro_search_b.base_raw_search.nodes.generate_raw_search_data import (
+    generate_raw_search_data,
+)
+from onyx.agent_search.pro_search_b.base_raw_search.states import BaseRawSearchInput
+from onyx.agent_search.pro_search_b.base_raw_search.states import BaseRawSearchOutput
+from onyx.agent_search.pro_search_b.base_raw_search.states import BaseRawSearchState
+from onyx.agent_search.pro_search_b.expanded_retrieval.graph_builder import (
+    expanded_retrieval_graph_builder,
+)
+
+
+def base_raw_search_graph_builder() -> StateGraph:
+    graph = StateGraph(
+        state_schema=BaseRawSearchState,
+        input=BaseRawSearchInput,
+        output=BaseRawSearchOutput,
+    )
+
+    ### Add nodes ###
+
+    expanded_retrieval = expanded_retrieval_graph_builder().compile()
+    graph.add_node(
+        node="generate_raw_search_data",
+        action=generate_raw_search_data,
+    )
+
+    graph.add_node(
+        node="expanded_retrieval_base_search",
+        action=expanded_retrieval,
+    )
+    graph.add_node(
+        node="format_raw_search_results",
+        action=format_raw_search_results,
+    )
+
+    ### Add edges ###
+
+    graph.add_edge(start_key=START, end_key="generate_raw_search_data")
+
+    graph.add_edge(
+        start_key="generate_raw_search_data",
+        end_key="expanded_retrieval_base_search",
+    )
+    graph.add_edge(
+        start_key="expanded_retrieval_base_search",
+        end_key="format_raw_search_results",
+    )
+
+    # graph.add_edge(
+    #     start_key="expanded_retrieval_base_search",
+    #     end_key=END,
+    # )
+
+    graph.add_edge(
+        start_key="format_raw_search_results",
+        end_key=END,
+    )
+
+    return graph
+
+
+if __name__ == "__main__":
+    pass

backend/onyx/agent_search/pro_search_b/base_raw_search/models.py

@@ -0,0 +1,20 @@
+from pydantic import BaseModel
+
+from onyx.agent_search.pro_search_b.expanded_retrieval.models import QueryResult
+from onyx.agent_search.shared_graph_utils.models import AgentChunkStats
+from onyx.context.search.models import InferenceSection
+
+### Models ###
+
+
+class AnswerRetrievalStats(BaseModel):
+    answer_retrieval_stats: dict[str, float | int]
+
+
+class QuestionAnswerResults(BaseModel):
+    question: str
+    answer: str
+    quality: str
+    expanded_retrieval_results: list[QueryResult]
+    documents: list[InferenceSection]
+    sub_question_retrieval_stats: list[AgentChunkStats]

backend/onyx/agent_search/pro_search_b/base_raw_search/nodes/format_raw_search_results.py

@@ -0,0 +1,16 @@
+from onyx.agent_search.pro_search_b.base_raw_search.states import BaseRawSearchOutput
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import (
+    ExpandedRetrievalOutput,
+)
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def format_raw_search_results(state: ExpandedRetrievalOutput) -> BaseRawSearchOutput:
+    logger.debug("format_raw_search_results")
+    return BaseRawSearchOutput(
+        base_expanded_retrieval_result=state["expanded_retrieval_result"],
+        # base_retrieval_results=[state["expanded_retrieval_result"]],
+        # base_search_documents=[],
+    )

backend/onyx/agent_search/pro_search_b/base_raw_search/nodes/generate_raw_search_data.py

@@ -0,0 +1,16 @@
+from onyx.agent_search.core_state import CoreState
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import (
+    ExpandedRetrievalInput,
+)
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def generate_raw_search_data(state: CoreState) -> ExpandedRetrievalInput:
+    logger.debug("generate_raw_search_data")
+    return ExpandedRetrievalInput(
+        question=state["base_question"],
+        base_search=True,
+        sub_question_id=None,  # This graph is always and only used for the original question
+    )

backend/onyx/agent_search/pro_search_b/base_raw_search/states.py

@@ -0,0 +1,42 @@
+from typing import TypedDict
+
+from onyx.agent_search.core_state import CoreState
+from onyx.agent_search.core_state import SubgraphCoreState
+from onyx.agent_search.pro_search_b.expanded_retrieval.models import (
+    ExpandedRetrievalResult,
+)
+
+
+## Update States
+
+
+## Graph Input State
+
+
+class BaseRawSearchInput(CoreState, SubgraphCoreState):
+    pass
+
+
+## Graph Output State
+
+
+class BaseRawSearchOutput(TypedDict):
+    """
+    This is a list of results even though each call of this subgraph only returns one result.
+    This is because if we parallelize the answer query subgraph, there will be multiple
+      results in a list so the add operator is used to add them together.
+    """
+
+    # base_search_documents: Annotated[list[InferenceSection], dedup_inference_sections]
+    # base_retrieval_results: Annotated[list[ExpandedRetrievalResult], add]
+    base_expanded_retrieval_result: ExpandedRetrievalResult
+
+
+## Graph State
+
+
+class BaseRawSearchState(
+    BaseRawSearchInput,
+    BaseRawSearchOutput,
+):
+    pass

backend/onyx/agent_search/pro_search_b/expanded_retrieval/edges.py

@@ -0,0 +1,26 @@
+from collections.abc import Hashable
+
+from langgraph.types import Send
+
+from onyx.agent_search.pro_search_b.expanded_retrieval.nodes import RetrievalInput
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import (
+    ExpandedRetrievalState,
+)
+
+
+def parallel_retrieval_edge(state: ExpandedRetrievalState) -> list[Send | Hashable]:
+    question = state.get("question", state["subgraph_config"].search_request.query)
+
+    query_expansions = state.get("expanded_queries", []) + [question]
+    return [
+        Send(
+            "doc_retrieval",
+            RetrievalInput(
+                query_to_retrieve=query,
+                question=question,
+                base_search=False,
+                sub_question_id=state.get("sub_question_id"),
+            ),
+        )
+        for query in query_expansions
+    ]

backend/onyx/agent_search/pro_search_b/expanded_retrieval/graph_builder.py

@@ -0,0 +1,122 @@
+from langgraph.graph import END
+from langgraph.graph import START
+from langgraph.graph import StateGraph
+
+from onyx.agent_search.pro_search_b.expanded_retrieval.edges import (
+    parallel_retrieval_edge,
+)
+from onyx.agent_search.pro_search_b.expanded_retrieval.nodes import doc_reranking
+from onyx.agent_search.pro_search_b.expanded_retrieval.nodes import doc_retrieval
+from onyx.agent_search.pro_search_b.expanded_retrieval.nodes import doc_verification
+from onyx.agent_search.pro_search_b.expanded_retrieval.nodes import expand_queries
+from onyx.agent_search.pro_search_b.expanded_retrieval.nodes import format_results
+from onyx.agent_search.pro_search_b.expanded_retrieval.nodes import verification_kickoff
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import (
+    ExpandedRetrievalInput,
+)
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import (
+    ExpandedRetrievalOutput,
+)
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import (
+    ExpandedRetrievalState,
+)
+from onyx.agent_search.shared_graph_utils.utils import get_test_config
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def expanded_retrieval_graph_builder() -> StateGraph:
+    graph = StateGraph(
+        state_schema=ExpandedRetrievalState,
+        input=ExpandedRetrievalInput,
+        output=ExpandedRetrievalOutput,
+    )
+
+    ### Add nodes ###
+
+    graph.add_node(
+        node="expand_queries",
+        action=expand_queries,
+    )
+
+    graph.add_node(
+        node="doc_retrieval",
+        action=doc_retrieval,
+    )
+    graph.add_node(
+        node="verification_kickoff",
+        action=verification_kickoff,
+    )
+    graph.add_node(
+        node="doc_verification",
+        action=doc_verification,
+    )
+    graph.add_node(
+        node="doc_reranking",
+        action=doc_reranking,
+    )
+    graph.add_node(
+        node="format_results",
+        action=format_results,
+    )
+
+    ### Add edges ###
+    graph.add_edge(
+        start_key=START,
+        end_key="expand_queries",
+    )
+
+    graph.add_conditional_edges(
+        source="expand_queries",
+        path=parallel_retrieval_edge,
+        path_map=["doc_retrieval"],
+    )
+    graph.add_edge(
+        start_key="doc_retrieval",
+        end_key="verification_kickoff",
+    )
+    graph.add_edge(
+        start_key="doc_verification",
+        end_key="doc_reranking",
+    )
+    graph.add_edge(
+        start_key="doc_reranking",
+        end_key="format_results",
+    )
+    graph.add_edge(
+        start_key="format_results",
+        end_key=END,
+    )
+
+    return graph
+
+
+if __name__ == "__main__":
+    from onyx.db.engine import get_session_context_manager
+    from onyx.llm.factory import get_default_llms
+    from onyx.context.search.models import SearchRequest
+
+    graph = expanded_retrieval_graph_builder()
+    compiled_graph = graph.compile()
+    primary_llm, fast_llm = get_default_llms()
+    search_request = SearchRequest(
+        query="what can you do with onyx or danswer?",
+    )
+
+    with get_session_context_manager() as db_session:
+        pro_search_config, search_tool = get_test_config(
+            db_session, primary_llm, fast_llm, search_request
+        )
+        inputs = ExpandedRetrievalInput(
+            question="what can you do with onyx?",
+            base_search=False,
+            sub_question_id=None,
+        )
+        for thing in compiled_graph.stream(
+            input=inputs,
+            config={"metadata": {"config": pro_search_config}},
+            # debug=True,
+            subgraphs=True,
+        ):
+            logger.debug(thing)

backend/onyx/agent_search/pro_search_b/expanded_retrieval/models.py

@@ -0,0 +1,13 @@
+from pydantic import BaseModel
+
+from onyx.agent_search.shared_graph_utils.models import AgentChunkStats
+from onyx.agent_search.shared_graph_utils.models import QueryResult
+from onyx.context.search.models import InferenceSection
+
+### Models ###
+
+
+class ExpandedRetrievalResult(BaseModel):
+    expanded_queries_results: list[QueryResult]
+    all_documents: list[InferenceSection]
+    sub_question_retrieval_stats: AgentChunkStats

backend/onyx/agent_search/pro_search_b/expanded_retrieval/nodes.py

@@ -0,0 +1,408 @@
+from collections import defaultdict
+from collections.abc import Callable
+from typing import cast
+from typing import Literal
+
+import numpy as np
+from langchain_core.callbacks.manager import dispatch_custom_event
+from langchain_core.messages import HumanMessage
+from langchain_core.messages import merge_message_runs
+from langgraph.types import Command
+from langgraph.types import Send
+
+from onyx.agent_search.pro_search_b.expanded_retrieval.models import (
+    ExpandedRetrievalResult,
+)
+from onyx.agent_search.pro_search_b.expanded_retrieval.models import QueryResult
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import DocRerankingUpdate
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import DocRetrievalUpdate
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import (
+    DocVerificationInput,
+)
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import (
+    DocVerificationUpdate,
+)
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import (
+    ExpandedRetrievalInput,
+)
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import (
+    ExpandedRetrievalState,
+)
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import (
+    ExpandedRetrievalUpdate,
+)
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import InferenceSection
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import (
+    QueryExpansionUpdate,
+)
+from onyx.agent_search.pro_search_b.expanded_retrieval.states import RetrievalInput
+from onyx.agent_search.shared_graph_utils.calculations import get_fit_scores
+from onyx.agent_search.shared_graph_utils.models import AgentChunkStats
+from onyx.agent_search.shared_graph_utils.models import RetrievalFitStats
+from onyx.agent_search.shared_graph_utils.prompts import REWRITE_PROMPT_MULTI_ORIGINAL
+from onyx.agent_search.shared_graph_utils.prompts import VERIFIER_PROMPT
+from onyx.agent_search.shared_graph_utils.utils import dispatch_separated
+from onyx.agent_search.shared_graph_utils.utils import parse_question_id
+from onyx.chat.models import ExtendedToolResponse
+from onyx.chat.models import SubQueryPiece
+from onyx.configs.dev_configs import AGENT_MAX_QUERY_RETRIEVAL_RESULTS
+from onyx.configs.dev_configs import AGENT_RERANKING_MAX_QUERY_RETRIEVAL_RESULTS
+from onyx.configs.dev_configs import AGENT_RERANKING_STATS
+from onyx.configs.dev_configs import AGENT_RETRIEVAL_STATS
+from onyx.context.search.models import SearchRequest
+from onyx.context.search.pipeline import retrieval_preprocessing
+from onyx.context.search.postprocessing.postprocessing import rerank_sections
+from onyx.db.engine import get_session_context_manager
+from onyx.llm.interfaces import LLM
+from onyx.tools.models import SearchQueryInfo
+from onyx.tools.tool_implementations.search.search_tool import (
+    SEARCH_RESPONSE_SUMMARY_ID,
+)
+from onyx.tools.tool_implementations.search.search_tool import SearchResponseSummary
+from onyx.tools.tool_implementations.search.search_tool import yield_search_responses
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def dispatch_subquery(level: int, question_nr: int) -> Callable[[str, int], None]:
+    def helper(token: str, num: int) -> None:
+        dispatch_custom_event(
+            "subqueries",
+            SubQueryPiece(
+                sub_query=token,
+                level=level,
+                level_question_nr=question_nr,
+                query_id=num,
+            ),
+        )
+
+    return helper
+
+
+def expand_queries(state: ExpandedRetrievalInput) -> QueryExpansionUpdate:
+    # Sometimes we want to expand the original question, sometimes we want to expand a sub-question.
+    # When we are running this node on the original question, no question is explictly passed in.
+    # Instead, we use the original question from the search request.
+    question = state.get("question", state["subgraph_config"].search_request.query)
+    llm: LLM = state["subgraph_fast_llm"]
+    state["subgraph_db_session"]
+    chat_session_id = state["subgraph_config"].chat_session_id
+    sub_question_id = state.get("sub_question_id")
+    if sub_question_id is None:
+        level, question_nr = 0, 0
+    else:
+        level, question_nr = parse_question_id(sub_question_id)
+
+    if chat_session_id is None:
+        raise ValueError("chat_session_id must be provided for agent search")
+
+    msg = [
+        HumanMessage(
+            content=REWRITE_PROMPT_MULTI_ORIGINAL.format(question=question),
+        )
+    ]
+
+    llm_response_list = dispatch_separated(
+        llm.stream(prompt=msg), dispatch_subquery(level, question_nr)
+    )
+
+    llm_response = merge_message_runs(llm_response_list, chunk_separator="")[0].content
+
+    rewritten_queries = llm_response.split("\n")
+
+    return QueryExpansionUpdate(
+        expanded_queries=rewritten_queries,
+    )
+
+
+def doc_retrieval(state: RetrievalInput) -> DocRetrievalUpdate:
+    """
+    Retrieve documents
+
+    Args:
+        state (RetrievalInput): Primary state + the query to retrieve
+
+    Updates:
+        expanded_retrieval_results: list[ExpandedRetrievalResult]
+        retrieved_documents: list[InferenceSection]
+    """
+    query_to_retrieve = state["query_to_retrieve"]
+    search_tool = state["subgraph_search_tool"]
+
+    retrieved_docs: list[InferenceSection] = []
+    if not query_to_retrieve.strip():
+        logger.warning("Empty query, skipping retrieval")
+        return DocRetrievalUpdate(
+            expanded_retrieval_results=[],
+            retrieved_documents=[],
+        )
+
+    query_info = None
+    # new db session to avoid concurrency issues
+    with get_session_context_manager() as db_session:
+        for tool_response in search_tool.run(
+            query=query_to_retrieve,
+            force_no_rerank=True,
+            alternate_db_session=db_session,
+        ):
+            # get retrieved docs to send to the rest of the graph
+            if tool_response.id == SEARCH_RESPONSE_SUMMARY_ID:
+                response = cast(SearchResponseSummary, tool_response.response)
+                retrieved_docs = response.top_sections
+                query_info = SearchQueryInfo(
+                    predicted_search=response.predicted_search,
+                    final_filters=response.final_filters,
+                    recency_bias_multiplier=response.recency_bias_multiplier,
+                )
+
+    retrieved_docs = retrieved_docs[:AGENT_MAX_QUERY_RETRIEVAL_RESULTS]
+    pre_rerank_docs = retrieved_docs
+    if search_tool.search_pipeline is not None:
+        pre_rerank_docs = (
+            search_tool.search_pipeline._retrieved_sections or retrieved_docs
+        )
+
+    if AGENT_RETRIEVAL_STATS:
+        fit_scores = get_fit_scores(
+            pre_rerank_docs,
+            retrieved_docs,
+        )
+    else:
+        fit_scores = None
+
+    expanded_retrieval_result = QueryResult(
+        query=query_to_retrieve,
+        search_results=retrieved_docs,
+        stats=fit_scores,
+        query_info=query_info,
+    )
+    return DocRetrievalUpdate(
+        expanded_retrieval_results=[expanded_retrieval_result],
+        retrieved_documents=retrieved_docs,
+    )
+
+
+def verification_kickoff(
+    state: ExpandedRetrievalState,
+) -> Command[Literal["doc_verification"]]:
+    documents = state["retrieved_documents"]
+    verification_question = state.get(
+        "question", state["subgraph_config"].search_request.query
+    )
+    sub_question_id = state.get("sub_question_id")
+    return Command(
+        update={},
+        goto=[
+            Send(
+                node="doc_verification",
+                arg=DocVerificationInput(
+                    doc_to_verify=doc,
+                    question=verification_question,
+                    base_search=False,
+                    sub_question_id=sub_question_id,
+                ),
+            )
+            for doc in documents
+        ],
+    )
+
+
+def doc_verification(state: DocVerificationInput) -> DocVerificationUpdate:
+    """
+    Check whether the document is relevant for the original user question
+
+    Args:
+        state (DocVerificationInput): The current state
+
+    Updates:
+        verified_documents: list[InferenceSection]
+    """
+
+    question = state["question"]
+    doc_to_verify = state["doc_to_verify"]
+    document_content = doc_to_verify.combined_content
+
+    msg = [
+        HumanMessage(
+            content=VERIFIER_PROMPT.format(
+                question=question, document_content=document_content
+            )
+        )
+    ]
+
+    fast_llm = state["subgraph_fast_llm"]
+
+    response = fast_llm.invoke(msg)
+
+    verified_documents = []
+    if isinstance(response.content, str) and "yes" in response.content.lower():
+        verified_documents.append(doc_to_verify)
+
+    return DocVerificationUpdate(
+        verified_documents=verified_documents,
+    )
+
+
+def doc_reranking(state: ExpandedRetrievalState) -> DocRerankingUpdate:
+    verified_documents = state["verified_documents"]
+
+    # Rerank post retrieval and verification. First, create a search query
+    # then create the list of reranked sections
+
+    question = state.get("question", state["subgraph_config"].search_request.query)
+    with get_session_context_manager() as db_session:
+        _search_query = retrieval_preprocessing(
+            search_request=SearchRequest(query=question),
+            user=state["subgraph_search_tool"].user,  # bit of a hack
+            llm=state["subgraph_fast_llm"],
+            db_session=db_session,
+        )
+
+    # skip section filtering
+
+    if (
+        _search_query.rerank_settings
+        and _search_query.rerank_settings.rerank_model_name
+        and _search_query.rerank_settings.num_rerank > 0
+    ):
+        reranked_documents = rerank_sections(
+            _search_query,
+            verified_documents,
+        )
+    else:
+        logger.warning("No reranking settings found, using unranked documents")
+        reranked_documents = verified_documents
+
+    if AGENT_RERANKING_STATS:
+        fit_scores = get_fit_scores(verified_documents, reranked_documents)
+    else:
+        fit_scores = RetrievalFitStats(fit_score_lift=0, rerank_effect=0, fit_scores={})
+
+    # TODO: stream deduped docs here, or decide to use search tool ranking/verification
+
+    return DocRerankingUpdate(
+        reranked_documents=[
+            doc for doc in reranked_documents if type(doc) == InferenceSection
+        ][:AGENT_RERANKING_MAX_QUERY_RETRIEVAL_RESULTS],
+        sub_question_retrieval_stats=fit_scores,
+    )
+
+
+def _calculate_sub_question_retrieval_stats(
+    verified_documents: list[InferenceSection],
+    expanded_retrieval_results: list[QueryResult],
+) -> AgentChunkStats:
+    chunk_scores: dict[str, dict[str, list[int | float]]] = defaultdict(
+        lambda: defaultdict(list)
+    )
+
+    for expanded_retrieval_result in expanded_retrieval_results:
+        for doc in expanded_retrieval_result.search_results:
+            doc_chunk_id = f"{doc.center_chunk.document_id}_{doc.center_chunk.chunk_id}"
+            if doc.center_chunk.score is not None:
+                chunk_scores[doc_chunk_id]["score"].append(doc.center_chunk.score)
+
+    verified_doc_chunk_ids = [
+        f"{verified_document.center_chunk.document_id}_{verified_document.center_chunk.chunk_id}"
+        for verified_document in verified_documents
+    ]
+    dismissed_doc_chunk_ids = []
+
+    raw_chunk_stats_counts: dict[str, int] = defaultdict(int)
+    raw_chunk_stats_scores: dict[str, float] = defaultdict(float)
+    for doc_chunk_id, chunk_data in chunk_scores.items():
+        if doc_chunk_id in verified_doc_chunk_ids:
+            raw_chunk_stats_counts["verified_count"] += 1
+
+            valid_chunk_scores = [
+                score for score in chunk_data["score"] if score is not None
+            ]
+            raw_chunk_stats_scores["verified_scores"] += float(
+                np.mean(valid_chunk_scores)
+            )
+        else:
+            raw_chunk_stats_counts["rejected_count"] += 1
+            valid_chunk_scores = [
+                score for score in chunk_data["score"] if score is not None
+            ]
+            raw_chunk_stats_scores["rejected_scores"] += float(
+                np.mean(valid_chunk_scores)
+            )
+            dismissed_doc_chunk_ids.append(doc_chunk_id)
+
+    if raw_chunk_stats_counts["verified_count"] == 0:
+        verified_avg_scores = 0.0
+    else:
+        verified_avg_scores = raw_chunk_stats_scores["verified_scores"] / float(
+            raw_chunk_stats_counts["verified_count"]
+        )
+
+    rejected_scores = raw_chunk_stats_scores.get("rejected_scores", None)
+    if rejected_scores is not None:
+        rejected_avg_scores = rejected_scores / float(
+            raw_chunk_stats_counts["rejected_count"]
+        )
+    else:
+        rejected_avg_scores = None
+
+    chunk_stats = AgentChunkStats(
+        verified_count=raw_chunk_stats_counts["verified_count"],
+        verified_avg_scores=verified_avg_scores,
+        rejected_count=raw_chunk_stats_counts["rejected_count"],
+        rejected_avg_scores=rejected_avg_scores,
+        verified_doc_chunk_ids=verified_doc_chunk_ids,
+        dismissed_doc_chunk_ids=dismissed_doc_chunk_ids,
+    )
+
+    return chunk_stats
+
+
+def format_results(state: ExpandedRetrievalState) -> ExpandedRetrievalUpdate:
+    level, question_nr = parse_question_id(state.get("sub_question_id") or "0_0")
+    query_infos = [
+        result.query_info
+        for result in state["expanded_retrieval_results"]
+        if result.query_info is not None
+    ]
+    if len(query_infos) == 0:
+        raise ValueError("No query info found")
+
+    # main question docs will be sent later after aggregation and deduping with sub-question docs
+    if not (level == 0 and question_nr == 0):
+        for tool_response in yield_search_responses(
+            query=state["question"],
+            reranked_sections=state[
+                "retrieved_documents"
+            ],  # TODO: rename params. this one is supposed to be the sections pre-merging
+            final_context_sections=state["reranked_documents"],
+            search_query_info=query_infos[0],  # TODO: handle differing query infos?
+            get_section_relevance=lambda: None,  # TODO: add relevance
+            search_tool=state["subgraph_search_tool"],
+        ):
+            dispatch_custom_event(
+                "tool_response",
+                ExtendedToolResponse(
+                    id=tool_response.id,
+                    response=tool_response.response,
+                    level=level,
+                    level_question_nr=question_nr,
+                ),
+            )
+    sub_question_retrieval_stats = _calculate_sub_question_retrieval_stats(
+        verified_documents=state["verified_documents"],
+        expanded_retrieval_results=state["expanded_retrieval_results"],
+    )
+
+    if sub_question_retrieval_stats is None:
+        sub_question_retrieval_stats = AgentChunkStats()
+    # else:
+    #    sub_question_retrieval_stats = [sub_question_retrieval_stats]
+
+    return ExpandedRetrievalUpdate(
+        expanded_retrieval_result=ExpandedRetrievalResult(
+            expanded_queries_results=state["expanded_retrieval_results"],
+            all_documents=state["reranked_documents"],
+            sub_question_retrieval_stats=sub_question_retrieval_stats,
+        ),
+    )

backend/onyx/agent_search/pro_search_b/expanded_retrieval/states.py

@@ -0,0 +1,82 @@
+from operator import add
+from typing import Annotated
+from typing import TypedDict
+
+from onyx.agent_search.core_state import SubgraphCoreState
+from onyx.agent_search.pro_search_b.expanded_retrieval.models import (
+    ExpandedRetrievalResult,
+)
+from onyx.agent_search.pro_search_b.expanded_retrieval.models import QueryResult
+from onyx.agent_search.shared_graph_utils.models import RetrievalFitStats
+from onyx.agent_search.shared_graph_utils.operators import dedup_inference_sections
+from onyx.context.search.models import InferenceSection
+
+
+### States ###
+
+## Graph Input State
+
+
+class ExpandedRetrievalInput(SubgraphCoreState):
+    question: str
+    base_search: bool
+    sub_question_id: str | None
+
+
+## Update/Return States
+
+
+class QueryExpansionUpdate(TypedDict):
+    expanded_queries: list[str]
+
+
+class DocVerificationUpdate(TypedDict):
+    verified_documents: Annotated[list[InferenceSection], dedup_inference_sections]
+
+
+class DocRetrievalUpdate(TypedDict):
+    expanded_retrieval_results: Annotated[list[QueryResult], add]
+    retrieved_documents: Annotated[list[InferenceSection], dedup_inference_sections]
+
+
+class DocRerankingUpdate(TypedDict):
+    reranked_documents: Annotated[list[InferenceSection], dedup_inference_sections]
+    sub_question_retrieval_stats: RetrievalFitStats | None
+
+
+class ExpandedRetrievalUpdate(TypedDict):
+    expanded_retrieval_result: ExpandedRetrievalResult
+
+
+## Graph Output State
+
+
+class ExpandedRetrievalOutput(TypedDict):
+    expanded_retrieval_result: ExpandedRetrievalResult
+    base_expanded_retrieval_result: ExpandedRetrievalResult
+
+
+## Graph State
+
+
+class ExpandedRetrievalState(
+    # This includes the core state
+    ExpandedRetrievalInput,
+    QueryExpansionUpdate,
+    DocRetrievalUpdate,
+    DocVerificationUpdate,
+    DocRerankingUpdate,
+    ExpandedRetrievalOutput,
+):
+    pass
+
+
+## Conditional Input States
+
+
+class DocVerificationInput(ExpandedRetrievalInput):
+    doc_to_verify: InferenceSection
+
+
+class RetrievalInput(ExpandedRetrievalInput):
+    query_to_retrieve: str

backend/onyx/agent_search/pro_search_b/main/edges.py

@@ -0,0 +1,89 @@
+from collections.abc import Hashable
+from typing import Literal
+
+from langgraph.types import Send
+
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.states import (
+    AnswerQuestionInput,
+)
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.states import (
+    AnswerQuestionOutput,
+)
+from onyx.agent_search.pro_search_b.main.states import MainState
+from onyx.agent_search.pro_search_b.main.states import RequireRefinedAnswerUpdate
+from onyx.agent_search.shared_graph_utils.utils import make_question_id
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def parallelize_initial_sub_question_answering(
+    state: MainState,
+) -> list[Send | Hashable]:
+    if len(state["initial_decomp_questions"]) > 0:
+        # sub_question_record_ids = [subq_record.id for subq_record in state["sub_question_records"]]
+        # if len(state["sub_question_records"]) == 0:
+        #     if state["config"].use_persistence:
+        #         raise ValueError("No sub-questions found for initial decompozed questions")
+        #     else:
+        #         # in this case, we are doing retrieval on the original question.
+        #         # to make all the logic consistent, we create a new sub-question
+        #         # with the same content as the original question
+        #         sub_question_record_ids = [1] * len(state["initial_decomp_questions"])
+
+        return [
+            Send(
+                "answer_query_subgraph",
+                AnswerQuestionInput(
+                    question=question,
+                    question_id=make_question_id(0, question_nr),
+                ),
+            )
+            for question_nr, question in enumerate(state["initial_decomp_questions"])
+        ]
+
+    else:
+        return [
+            Send(
+                "ingest_answers",
+                AnswerQuestionOutput(
+                    answer_results=[],
+                ),
+            )
+        ]
+
+
+# Define the function that determines whether to continue or not
+def continue_to_refined_answer_or_end(
+    state: RequireRefinedAnswerUpdate,
+) -> Literal["refined_decompose", "logging_node"]:
+    if state["require_refined_answer"]:
+        return "refined_decompose"
+    else:
+        return "logging_node"
+
+
+def parallelize_refined_sub_question_answering(
+    state: MainState,
+) -> list[Send | Hashable]:
+    if len(state["refined_sub_questions"]) > 0:
+        return [
+            Send(
+                "answer_refinement_sub_question",
+                AnswerQuestionInput(
+                    question=question_data.sub_question,
+                    question_id=make_question_id(1, question_nr),
+                ),
+            )
+            for question_nr, question_data in state["refined_sub_questions"].items()
+        ]
+
+    else:
+        return [
+            Send(
+                "ingest_refined_sub_answers",
+                AnswerQuestionOutput(
+                    answer_results=[],
+                ),
+            )
+        ]

backend/onyx/agent_search/pro_search_b/main/graph_builder.py

@@ -0,0 +1,264 @@
+from langgraph.graph import END
+from langgraph.graph import START
+from langgraph.graph import StateGraph
+
+from onyx.agent_search.pro_search_b.answer_initial_sub_question.graph_builder import (
+    answer_query_graph_builder,
+)
+from onyx.agent_search.pro_search_b.answer_refinement_sub_question.graph_builder import (
+    answer_refined_query_graph_builder,
+)
+from onyx.agent_search.pro_search_b.base_raw_search.graph_builder import (
+    base_raw_search_graph_builder,
+)
+from onyx.agent_search.pro_search_b.main.edges import continue_to_refined_answer_or_end
+from onyx.agent_search.pro_search_b.main.edges import (
+    parallelize_initial_sub_question_answering,
+)
+from onyx.agent_search.pro_search_b.main.edges import (
+    parallelize_refined_sub_question_answering,
+)
+from onyx.agent_search.pro_search_b.main.nodes import agent_logging
+from onyx.agent_search.pro_search_b.main.nodes import entity_term_extraction_llm
+from onyx.agent_search.pro_search_b.main.nodes import generate_initial_answer
+from onyx.agent_search.pro_search_b.main.nodes import generate_refined_answer
+from onyx.agent_search.pro_search_b.main.nodes import ingest_initial_base_retrieval
+from onyx.agent_search.pro_search_b.main.nodes import (
+    ingest_initial_sub_question_answers,
+)
+from onyx.agent_search.pro_search_b.main.nodes import ingest_refined_answers
+from onyx.agent_search.pro_search_b.main.nodes import initial_answer_quality_check
+from onyx.agent_search.pro_search_b.main.nodes import initial_sub_question_creation
+from onyx.agent_search.pro_search_b.main.nodes import refined_answer_decision
+from onyx.agent_search.pro_search_b.main.nodes import refined_sub_question_creation
+from onyx.agent_search.pro_search_b.main.states import MainInput
+from onyx.agent_search.pro_search_b.main.states import MainState
+from onyx.agent_search.shared_graph_utils.utils import get_test_config
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+test_mode = False
+
+
+def main_graph_builder(test_mode: bool = False) -> StateGraph:
+    graph = StateGraph(
+        state_schema=MainState,
+        input=MainInput,
+    )
+
+    graph.add_node(
+        node="initial_sub_question_creation",
+        action=initial_sub_question_creation,
+    )
+    answer_query_subgraph = answer_query_graph_builder().compile()
+    graph.add_node(
+        node="answer_query_subgraph",
+        action=answer_query_subgraph,
+    )
+
+    base_raw_search_subgraph = base_raw_search_graph_builder().compile()
+    graph.add_node(
+        node="base_raw_search_subgraph",
+        action=base_raw_search_subgraph,
+    )
+
+    # refined_answer_subgraph = refined_answers_graph_builder().compile()
+    # graph.add_node(
+    #     node="refined_answer_subgraph",
+    #     action=refined_answer_subgraph,
+    # )
+
+    graph.add_node(
+        node="refined_sub_question_creation",
+        action=refined_sub_question_creation,
+    )
+
+    answer_refined_question = answer_refined_query_graph_builder().compile()
+    graph.add_node(
+        node="answer_refined_question",
+        action=answer_refined_question,
+    )
+
+    graph.add_node(
+        node="ingest_refined_answers",
+        action=ingest_refined_answers,
+    )
+
+    graph.add_node(
+        node="generate_refined_answer",
+        action=generate_refined_answer,
+    )
+
+    # graph.add_node(
+    #     node="check_refined_answer",
+    #     action=check_refined_answer,
+    # )
+
+    graph.add_node(
+        node="ingest_initial_retrieval",
+        action=ingest_initial_base_retrieval,
+    )
+    graph.add_node(
+        node="ingest_initial_sub_question_answers",
+        action=ingest_initial_sub_question_answers,
+    )
+    graph.add_node(
+        node="generate_initial_answer",
+        action=generate_initial_answer,
+    )
+
+    graph.add_node(
+        node="initial_answer_quality_check",
+        action=initial_answer_quality_check,
+    )
+
+    graph.add_node(
+        node="entity_term_extraction_llm",
+        action=entity_term_extraction_llm,
+    )
+    graph.add_node(
+        node="refined_answer_decision",
+        action=refined_answer_decision,
+    )
+
+    graph.add_node(
+        node="logging_node",
+        action=agent_logging,
+    )
+    # if test_mode:
+    #     graph.add_node(
+    #         node="generate_initial_base_answer",
+    #         action=generate_initial_base_answer,
+    #     )
+
+    ### Add edges ###
+
+    graph.add_edge(start_key=START, end_key="base_raw_search_subgraph")
+
+    graph.add_edge(
+        start_key="base_raw_search_subgraph",
+        end_key="ingest_initial_retrieval",
+    )
+
+    graph.add_edge(
+        start_key=START,
+        end_key="initial_sub_question_creation",
+    )
+    graph.add_conditional_edges(
+        source="initial_sub_question_creation",
+        path=parallelize_initial_sub_question_answering,
+        path_map=["answer_query_subgraph"],
+    )
+    graph.add_edge(
+        start_key="answer_query_subgraph",
+        end_key="ingest_initial_sub_question_answers",
+    )
+
+    graph.add_edge(
+        start_key=["ingest_initial_sub_question_answers", "ingest_initial_retrieval"],
+        end_key="generate_initial_answer",
+    )
+
+    graph.add_edge(
+        start_key=["ingest_initial_sub_question_answers", "ingest_initial_retrieval"],
+        end_key="entity_term_extraction_llm",
+    )
+
+    graph.add_edge(
+        start_key="generate_initial_answer",
+        end_key="initial_answer_quality_check",
+    )
+
+    graph.add_edge(
+        start_key=["initial_answer_quality_check", "entity_term_extraction_llm"],
+        end_key="refined_answer_decision",
+    )
+
+    graph.add_conditional_edges(
+        source="refined_answer_decision",
+        path=continue_to_refined_answer_or_end,
+        path_map=["refined_sub_question_creation", "logging_node"],
+    )
+
+    graph.add_conditional_edges(
+        source="refined_sub_question_creation",
+        path=parallelize_refined_sub_question_answering,
+        path_map=["answer_refined_question"],
+    )
+    graph.add_edge(
+        start_key="answer_refined_question",
+        end_key="ingest_refined_answers",
+    )
+
+    graph.add_edge(
+        start_key="ingest_refined_answers",
+        end_key="generate_refined_answer",
+    )
+
+    # graph.add_conditional_edges(
+    #     source="refined_answer_decision",
+    #     path=continue_to_refined_answer_or_end,
+    #     path_map=["refined_answer_subgraph", END],
+    # )
+
+    # graph.add_edge(
+    #     start_key="refined_answer_subgraph",
+    #     end_key="generate_refined_answer",
+    # )
+
+    graph.add_edge(
+        start_key="generate_refined_answer",
+        end_key="logging_node",
+    )
+
+    graph.add_edge(
+        start_key="logging_node",
+        end_key=END,
+    )
+
+    # graph.add_edge(
+    #     start_key="generate_refined_answer",
+    #     end_key="check_refined_answer",
+    # )
+
+    # graph.add_edge(
+    #     start_key="check_refined_answer",
+    #     end_key=END,
+    # )
+
+    return graph
+
+
+if __name__ == "__main__":
+    pass
+
+    from onyx.db.engine import get_session_context_manager
+    from onyx.llm.factory import get_default_llms
+    from onyx.context.search.models import SearchRequest
+
+    graph = main_graph_builder()
+    compiled_graph = graph.compile()
+    primary_llm, fast_llm = get_default_llms()
+
+    with get_session_context_manager() as db_session:
+        search_request = SearchRequest(query="Who created Excel?")
+        pro_search_config, search_tool = get_test_config(
+            db_session, primary_llm, fast_llm, search_request
+        )
+
+        inputs = MainInput(
+            primary_llm=primary_llm,
+            fast_llm=fast_llm,
+            db_session=db_session,
+            config=pro_search_config,
+            search_tool=search_tool,
+        )
+
+        for thing in compiled_graph.stream(
+            input=inputs,
+            # stream_mode="debug",
+            # debug=True,
+            subgraphs=True,
+        ):
+            logger.debug(thing)

backend/onyx/agent_search/pro_search_b/main/models.py

@@ -0,0 +1,43 @@
+from pydantic import BaseModel
+
+
+class FollowUpSubQuestion(BaseModel):
+    sub_question: str
+    sub_question_id: str
+    verified: bool
+    answered: bool
+    answer: str
+
+
+class AgentTimings(BaseModel):
+    base_duration__s: float | None
+    refined_duration__s: float | None
+    full_duration__s: float | None
+
+
+class AgentBaseMetrics(BaseModel):
+    num_verified_documents_total: int | None
+    num_verified_documents_core: int | None
+    verified_avg_score_core: float | None
+    num_verified_documents_base: int | float | None
+    verified_avg_score_base: float | None
+    base_doc_boost_factor: float | None
+    support_boost_factor: float | None
+    duration__s: float | None
+
+
+class AgentRefinedMetrics(BaseModel):
+    refined_doc_boost_factor: float | None
+    refined_question_boost_factor: float | None
+    duration__s: float | None
+
+
+class AgentAdditionalMetrics(BaseModel):
+    pass
+
+
+class CombinedAgentMetrics(BaseModel):
+    timings: AgentTimings
+    base_metrics: AgentBaseMetrics
+    refined_metrics: AgentRefinedMetrics
+    additional_metrics: AgentAdditionalMetrics

backend/onyx/agent_search/pro_search_b/main/states.py

@@ -0,0 +1,151 @@
+from datetime import datetime
+from operator import add
+from typing import Annotated
+from typing import TypedDict
+
+from onyx.agent_search.core_state import CoreState
+from onyx.agent_search.pro_search_b.expanded_retrieval.models import (
+    ExpandedRetrievalResult,
+)
+from onyx.agent_search.pro_search_b.expanded_retrieval.models import QueryResult
+from onyx.agent_search.pro_search_b.main.models import FollowUpSubQuestion
+from onyx.agent_search.shared_graph_utils.models import AgentBaseMetrics
+from onyx.agent_search.shared_graph_utils.models import AgentChunkStats
+from onyx.agent_search.shared_graph_utils.models import AgentRefinedMetrics
+from onyx.agent_search.shared_graph_utils.models import EntityRelationshipTermExtraction
+from onyx.agent_search.shared_graph_utils.models import InitialAgentResultStats
+from onyx.agent_search.shared_graph_utils.models import (
+    QuestionAnswerResults,
+)
+from onyx.agent_search.shared_graph_utils.models import RefinedAgentStats
+from onyx.agent_search.shared_graph_utils.operators import dedup_inference_sections
+from onyx.agent_search.shared_graph_utils.operators import dedup_question_answer_results
+from onyx.context.search.models import InferenceSection
+
+
+### States ###
+
+## Update States
+
+
+class RefinedAgentStartStats(TypedDict):
+    agent_refined_start_time: datetime | None
+
+
+class RefinedAgentEndStats(TypedDict):
+    agent_refined_end_time: datetime | None
+    agent_refined_metrics: AgentRefinedMetrics
+
+
+class BaseDecompUpdateBase(TypedDict):
+    agent_start_time: datetime
+    initial_decomp_questions: list[str]
+
+
+class BaseDecompUpdate(
+    RefinedAgentStartStats, RefinedAgentEndStats, BaseDecompUpdateBase
+):
+    pass
+
+
+class InitialAnswerBASEUpdate(TypedDict):
+    initial_base_answer: str
+
+
+class InitialAnswerUpdate(TypedDict):
+    initial_answer: str
+    initial_agent_stats: InitialAgentResultStats | None
+    generated_sub_questions: list[str]
+    agent_base_end_time: datetime
+    agent_base_metrics: AgentBaseMetrics
+
+
+class RefinedAnswerUpdateBase(TypedDict):
+    refined_answer: str
+    refined_agent_stats: RefinedAgentStats | None
+    refined_answer_quality: bool
+
+
+class RefinedAnswerUpdate(RefinedAgentEndStats, RefinedAnswerUpdateBase):
+    pass
+
+
+class InitialAnswerQualityUpdate(TypedDict):
+    initial_answer_quality: bool
+
+
+class RequireRefinedAnswerUpdate(TypedDict):
+    require_refined_answer: bool
+
+
+class DecompAnswersUpdate(TypedDict):
+    documents: Annotated[list[InferenceSection], dedup_inference_sections]
+    decomp_answer_results: Annotated[
+        list[QuestionAnswerResults], dedup_question_answer_results
+    ]
+
+
+class FollowUpDecompAnswersUpdate(TypedDict):
+    refined_documents: Annotated[list[InferenceSection], dedup_inference_sections]
+    refined_decomp_answer_results: Annotated[list[QuestionAnswerResults], add]
+
+
+class ExpandedRetrievalUpdate(TypedDict):
+    all_original_question_documents: Annotated[
+        list[InferenceSection], dedup_inference_sections
+    ]
+    original_question_retrieval_results: list[QueryResult]
+    original_question_retrieval_stats: AgentChunkStats
+
+
+class EntityTermExtractionUpdate(TypedDict):
+    entity_retlation_term_extractions: EntityRelationshipTermExtraction
+
+
+class FollowUpSubQuestionsUpdateBase(TypedDict):
+    refined_sub_questions: dict[int, FollowUpSubQuestion]
+
+
+class FollowUpSubQuestionsUpdate(
+    RefinedAgentStartStats, FollowUpSubQuestionsUpdateBase
+):
+    pass
+
+
+## Graph Input State
+## Graph Input State
+
+
+class MainInput(CoreState):
+    pass
+
+
+## Graph State
+
+
+class MainState(
+    # This includes the core state
+    MainInput,
+    BaseDecompUpdateBase,
+    InitialAnswerUpdate,
+    InitialAnswerBASEUpdate,
+    DecompAnswersUpdate,
+    ExpandedRetrievalUpdate,
+    EntityTermExtractionUpdate,
+    InitialAnswerQualityUpdate,
+    RequireRefinedAnswerUpdate,
+    FollowUpSubQuestionsUpdateBase,
+    FollowUpDecompAnswersUpdate,
+    RefinedAnswerUpdateBase,
+    RefinedAgentStartStats,
+    RefinedAgentEndStats,
+):
+    # expanded_retrieval_result: Annotated[list[ExpandedRetrievalResult], add]
+    base_raw_search_result: Annotated[list[ExpandedRetrievalResult], add]
+
+
+## Graph Output State - presently not used
+
+
+class MainOutput(TypedDict):
+    pass

backend/onyx/agent_search/run_graph.py

@@ -0,0 +1,284 @@
+import asyncio
+from asyncio import AbstractEventLoop
+from collections.abc import AsyncIterable
+from collections.abc import Iterable
+from datetime import datetime
+from typing import cast
+
+from langchain_core.runnables.schema import StreamEvent
+from langgraph.graph.state import CompiledStateGraph
+
+from onyx.agent_search.basic.graph_builder import basic_graph_builder
+from onyx.agent_search.basic.states import BasicInput
+from onyx.agent_search.models import ProSearchConfig
+from onyx.agent_search.pro_search_a.main.graph_builder import (
+    main_graph_builder as main_graph_builder_a,
+)
+from onyx.agent_search.pro_search_a.main.states import MainInput as MainInput_a
+from onyx.agent_search.pro_search_b.main.graph_builder import (
+    main_graph_builder as main_graph_builder_b,
+)
+from onyx.agent_search.pro_search_b.main.states import MainInput as MainInput_b
+from onyx.agent_search.shared_graph_utils.utils import get_test_config
+from onyx.chat.llm_response_handler import LLMResponseHandlerManager
+from onyx.chat.models import AgentAnswerPiece
+from onyx.chat.models import AnswerPacket
+from onyx.chat.models import AnswerStream
+from onyx.chat.models import ExtendedToolResponse
+from onyx.chat.models import StreamStopInfo
+from onyx.chat.models import SubQueryPiece
+from onyx.chat.models import SubQuestionPiece
+from onyx.chat.models import ToolResponse
+from onyx.chat.prompt_builder.build import LLMCall
+from onyx.configs.dev_configs import GRAPH_NAME
+from onyx.context.search.models import SearchRequest
+from onyx.db.engine import get_session_context_manager
+from onyx.tools.tool_runner import ToolCallKickoff
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+_COMPILED_GRAPH: CompiledStateGraph | None = None
+
+
+def _set_combined_token_value(
+    combined_token: str, parsed_object: AgentAnswerPiece
+) -> AgentAnswerPiece:
+    parsed_object.answer_piece = combined_token
+
+    return parsed_object
+
+
+def _parse_agent_event(
+    event: StreamEvent,
+) -> AnswerPacket | None:
+    """
+    Parse the event into a typed object.
+    Return None if we are not interested in the event.
+    """
+
+    event_type = event["event"]
+
+    # We always just yield the event data, but this piece is useful for two development reasons:
+    # 1. It's a list of the names of every place we dispatch a custom event
+    # 2. We maintain the intended types yielded by each event
+    if event_type == "on_custom_event":
+        # TODO: different AnswerStream types for different events
+        if event["name"] == "decomp_qs":
+            return cast(SubQuestionPiece, event["data"])
+        elif event["name"] == "subqueries":
+            return cast(SubQueryPiece, event["data"])
+        elif event["name"] == "sub_answers":
+            return cast(AgentAnswerPiece, event["data"])
+        elif event["name"] == "sub_answer_finished":
+            return cast(StreamStopInfo, event["data"])
+        elif event["name"] == "initial_agent_answer":
+            return cast(AgentAnswerPiece, event["data"])
+        elif event["name"] == "refined_agent_answer":
+            return cast(AgentAnswerPiece, event["data"])
+        elif event["name"] == "start_refined_answer_creation":
+            return cast(ToolCallKickoff, event["data"])
+        elif event["name"] == "tool_response":
+            return cast(ToolResponse, event["data"])
+        elif event["name"] == "basic_response":
+            return cast(AnswerPacket, event["data"])
+    return None
+
+
+async def tear_down(event_loop: AbstractEventLoop) -> None:
+    # Collect all tasks and cancel those that are not 'done'.
+    tasks = asyncio.all_tasks(event_loop)
+    for task in tasks:
+        task.cancel()
+
+    # Wait for all tasks to complete, ignoring any CancelledErrors
+    try:
+        await asyncio.wait(tasks)
+    except asyncio.exceptions.CancelledError:
+        pass
+
+
+def _manage_async_event_streaming(
+    compiled_graph: CompiledStateGraph,
+    config: ProSearchConfig | None,
+    graph_input: MainInput_a | MainInput_b | BasicInput,
+) -> Iterable[StreamEvent]:
+    async def _run_async_event_stream(
+        loop: AbstractEventLoop,
+    ) -> AsyncIterable[StreamEvent]:
+        try:
+            message_id = config.message_id if config else None
+            async for event in compiled_graph.astream_events(
+                input=graph_input,
+                config={"metadata": {"config": config, "thread_id": str(message_id)}},
+                # debug=True,
+                # indicating v2 here deserves further scrutiny
+                version="v2",
+            ):
+                yield event
+        finally:
+            await tear_down(loop)
+
+    # This might be able to be simplified
+    def _yield_async_to_sync() -> Iterable[StreamEvent]:
+        loop = asyncio.new_event_loop()
+        try:
+            # Get the async generator
+            async_gen = _run_async_event_stream(loop)
+            # Convert to AsyncIterator
+            async_iter = async_gen.__aiter__()
+            while True:
+                try:
+                    # Create a coroutine by calling anext with the async iterator
+                    next_coro = anext(async_iter)
+                    # Run the coroutine to get the next event
+                    event = loop.run_until_complete(next_coro)
+                    yield event
+                except StopAsyncIteration:
+                    break
+        finally:
+            loop.close()
+
+    return _yield_async_to_sync()
+
+
+def run_graph(
+    compiled_graph: CompiledStateGraph,
+    config: ProSearchConfig,
+    input: BasicInput | MainInput_a | MainInput_b,
+) -> AnswerStream:
+    input["base_question"] = config.search_request.query if config else ""
+
+    config.perform_initial_search_path_decision = True
+    config.perform_initial_search_decomposition = True
+
+    for event in _manage_async_event_streaming(
+        compiled_graph=compiled_graph, config=config, graph_input=input
+    ):
+        if not (parsed_object := _parse_agent_event(event)):
+            continue
+
+        yield parsed_object
+
+
+# TODO: call this once on startup, TBD where and if it should be gated based
+# on dev mode or not
+def load_compiled_graph(graph_name: str) -> CompiledStateGraph:
+    main_graph_builder = (
+        main_graph_builder_a if graph_name == "a" else main_graph_builder_b
+    )
+    global _COMPILED_GRAPH
+    if _COMPILED_GRAPH is None:
+        graph = main_graph_builder()
+        _COMPILED_GRAPH = graph.compile()
+    return _COMPILED_GRAPH
+
+
+def run_main_graph(
+    config: ProSearchConfig,
+    graph_name: str = "a",
+) -> AnswerStream:
+    compiled_graph = load_compiled_graph(graph_name)
+    if graph_name == "a":
+        input = MainInput_a()
+    else:
+        input = MainInput_b()
+
+    # Agent search is not a Tool per se, but this is helpful for the frontend
+    yield ToolCallKickoff(
+        tool_name="agent_search_0",
+        tool_args={"query": config.search_request.query},
+    )
+    yield from run_graph(compiled_graph, config, input)
+
+
+# TODO: unify input types, especially prosearchconfig
+def run_basic_graph(
+    config: ProSearchConfig,
+    last_llm_call: LLMCall | None,
+    response_handler_manager: LLMResponseHandlerManager,
+) -> AnswerStream:
+    graph = basic_graph_builder()
+    compiled_graph = graph.compile()
+    # TODO: unify basic input
+    input = BasicInput(
+        base_question="",
+        last_llm_call=last_llm_call,
+        response_handler_manager=response_handler_manager,
+        calls=0,
+    )
+    return run_graph(compiled_graph, config, input)
+
+
+if __name__ == "__main__":
+    from onyx.llm.factory import get_default_llms
+
+    now_start = datetime.now()
+    logger.debug(f"Start at {now_start}")
+
+    if GRAPH_NAME == "a":
+        graph = main_graph_builder_a()
+    else:
+        graph = main_graph_builder_b()
+    compiled_graph = graph.compile()
+    now_end = datetime.now()
+    logger.debug(f"Graph compiled in {now_end - now_start} seconds")
+    primary_llm, fast_llm = get_default_llms()
+    search_request = SearchRequest(
+        # query="what can you do with gitlab?",
+        # query="What are the guiding principles behind the development of cockroachDB",
+        # query="What are the temperatures in Munich, Hawaii, and New York?",
+        # query="When was Washington born?",
+        query="What is Onyx?",
+    )
+    # Joachim custom persona
+
+    with get_session_context_manager() as db_session:
+        config, search_tool = get_test_config(
+            db_session, primary_llm, fast_llm, search_request
+        )
+        # search_request.persona = get_persona_by_id(1, None, db_session)
+        config.use_persistence = True
+        config.perform_initial_search_path_decision = True
+        config.perform_initial_search_decomposition = True
+        if GRAPH_NAME == "a":
+            input = MainInput_a()
+        else:
+            input = MainInput_b()
+        # with open("output.txt", "w") as f:
+        tool_responses: list = []
+        for output in run_graph(compiled_graph, config, input):
+            # pass
+
+            if isinstance(output, ToolCallKickoff):
+                pass
+            elif isinstance(output, ExtendedToolResponse):
+                tool_responses.append(output.response)
+                logger.info(
+                    f"   ---- ET {output.level} - {output.level_question_nr} |  "
+                )
+            elif isinstance(output, SubQueryPiece):
+                logger.info(
+                    f"Sq {output.level} - {output.level_question_nr} - {output.sub_query} | "
+                )
+            elif isinstance(output, SubQuestionPiece):
+                logger.info(
+                    f"SQ {output.level} - {output.level_question_nr} - {output.sub_question} | "
+                )
+            elif (
+                isinstance(output, AgentAnswerPiece)
+                and output.answer_type == "agent_sub_answer"
+            ):
+                logger.info(
+                    f"   ---- SA {output.level} - {output.level_question_nr} {output.answer_piece} | "
+                )
+            elif (
+                isinstance(output, AgentAnswerPiece)
+                and output.answer_type == "agent_level_answer"
+            ):
+                logger.info(
+                    f"   ---------- FA {output.level} - {output.level_question_nr}  {output.answer_piece} | "
+                )
+
+        # for tool_response in tool_responses:
+        #    logger.debug(tool_response)

backend/onyx/agent_search/shared_graph_utils/agent_prompt_ops.py

@@ -0,0 +1,62 @@
+from langchain.schema import AIMessage
+from langchain.schema import HumanMessage
+from langchain.schema import SystemMessage
+from langchain_core.messages.tool import ToolMessage
+
+from onyx.llm.interfaces import LLMConfig
+from onyx.agent_search.shared_graph_utils.prompts import BASE_RAG_PROMPT_v2
+from onyx.context.search.models import InferenceSection
+from onyx.natural_language_processing.utils import tokenizer_trim_content
+from onyx.natural_language_processing.utils import get_tokenizer
+from onyx.llm.utils import get_max_input_tokens
+
+def build_sub_question_answer_prompt(
+    question: str,
+    original_question: str,
+    docs: list[InferenceSection],
+    persona_specification: str,
+    config: LLMConfig,
+) -> list[SystemMessage | HumanMessage | AIMessage | ToolMessage]:
+    system_message = SystemMessage(
+        content=persona_specification,
+    )
+
+    docs_format_list = [
+        f"""Document Number: [D{doc_nr + 1}]\n
+                             Content: {doc.combined_content}\n\n"""
+        for doc_nr, doc in enumerate(docs)
+    ]
+
+    docs_str = "\n\n".join(docs_format_list)
+
+    docs_str = trim_prompt_piece(config, docs_str, BASE_RAG_PROMPT_v2 + question + original_question)
+    human_message = HumanMessage(
+        content=BASE_RAG_PROMPT_v2.format(
+            question=question, original_question=original_question, context=docs_str
+        )
+    )
+
+    return [system_message, human_message]
+
+
+def trim_prompt_piece(config: LLMConfig, prompt_piece: str, reserved_str: str) -> str:
+    # TODO: this truncating might add latency. We could do a rougher + faster check
+    # first to determine whether truncation is needed
+
+    #TODO: maybe save the tokenizer and max input tokens if this is getting called multiple times?
+    llm_tokenizer = get_tokenizer(
+        provider_type=config.model_provider,
+        model_name=config.model_name,
+    )
+
+    max_tokens = get_max_input_tokens(
+        model_provider=config.model_provider,
+        model_name=config.model_name,
+    )
+
+    # slightly conservative trimming
+    return tokenizer_trim_content(
+        content=prompt_piece,
+        desired_length=max_tokens - len(llm_tokenizer.encode(reserved_str)),
+        tokenizer=llm_tokenizer,
+    )

backend/onyx/agent_search/shared_graph_utils/calculations.py

@@ -0,0 +1,98 @@
+import numpy as np
+
+from onyx.agent_search.shared_graph_utils.models import RetrievalFitScoreMetrics
+from onyx.agent_search.shared_graph_utils.models import RetrievalFitStats
+from onyx.chat.models import SectionRelevancePiece
+from onyx.context.search.models import InferenceSection
+from onyx.utils.logger import setup_logger
+
+logger = setup_logger()
+
+
+def unique_chunk_id(doc: InferenceSection) -> str:
+    return f"{doc.center_chunk.document_id}_{doc.center_chunk.chunk_id}"
+
+
+def calculate_rank_shift(list1: list, list2: list, top_n: int = 20) -> float:
+    shift = 0
+    for rank_first, doc_id in enumerate(list1[:top_n], 1):
+        try:
+            rank_second = list2.index(doc_id) + 1
+        except ValueError:
+            rank_second = len(list2)  # Document not found in second list
+
+        shift += np.abs(rank_first - rank_second) / np.log(1 + rank_first * rank_second)
+
+    return shift / top_n
+
+
+def get_fit_scores(
+    pre_reranked_results: list[InferenceSection],
+    post_reranked_results: list[InferenceSection] | list[SectionRelevancePiece],
+) -> RetrievalFitStats | None:
+    """
+    Calculate retrieval metrics for search purposes
+    """
+
+    if len(pre_reranked_results) == 0 or len(post_reranked_results) == 0:
+        return None
+
+    ranked_sections = {
+        "initial": pre_reranked_results,
+        "reranked": post_reranked_results,
+    }
+
+    fit_eval: RetrievalFitStats = RetrievalFitStats(
+        fit_score_lift=0,
+        rerank_effect=0,
+        fit_scores={
+            "initial": RetrievalFitScoreMetrics(scores={}, chunk_ids=[]),
+            "reranked": RetrievalFitScoreMetrics(scores={}, chunk_ids=[]),
+        },
+    )
+
+    for rank_type, docs in ranked_sections.items():
+        logger.debug(f"rank_type: {rank_type}")
+
+        for i in [1, 5, 10]:
+            fit_eval.fit_scores[rank_type].scores[str(i)] = (
+                sum(
+                    [
+                        float(doc.center_chunk.score)
+                        for doc in docs[:i]
+                        if type(doc) == InferenceSection
+                        and doc.center_chunk.score is not None
+                    ]
+                )
+                / i
+            )
+
+        fit_eval.fit_scores[rank_type].scores["fit_score"] = (
+            1
+            / 3
+            * (
+                fit_eval.fit_scores[rank_type].scores["1"]
+                + fit_eval.fit_scores[rank_type].scores["5"]
+                + fit_eval.fit_scores[rank_type].scores["10"]
+            )
+        )
+
+        fit_eval.fit_scores[rank_type].scores["fit_score"] = fit_eval.fit_scores[
+            rank_type
+        ].scores["1"]
+
+        fit_eval.fit_scores[rank_type].chunk_ids = [
+            unique_chunk_id(doc) for doc in docs if type(doc) == InferenceSection
+        ]
+
+    fit_eval.fit_score_lift = (
+        fit_eval.fit_scores["reranked"].scores["fit_score"]
+        / fit_eval.fit_scores["initial"].scores["fit_score"]
+    )
+
+    fit_eval.rerank_effect = calculate_rank_shift(
+        fit_eval.fit_scores["initial"].chunk_ids,
+        fit_eval.fit_scores["reranked"].chunk_ids,
+    )
+
+    return fit_eval

backend/onyx/agent_search/shared_graph_utils/models.py

@@ -0,0 +1,112 @@
+from typing import Literal
+
+from pydantic import BaseModel
+
+from onyx.agent_search.pro_search_a.main.models import AgentAdditionalMetrics
+from onyx.agent_search.pro_search_a.main.models import AgentBaseMetrics
+from onyx.agent_search.pro_search_a.main.models import AgentRefinedMetrics
+from onyx.agent_search.pro_search_a.main.models import AgentTimings
+from onyx.context.search.models import InferenceSection
+from onyx.tools.models import SearchQueryInfo
+
+
+# Pydantic models for structured outputs
+class RewrittenQueries(BaseModel):
+    rewritten_queries: list[str]
+
+
+class BinaryDecision(BaseModel):
+    decision: Literal["yes", "no"]
+
+
+class BinaryDecisionWithReasoning(BaseModel):
+    reasoning: str
+    decision: Literal["yes", "no"]
+
+
+class RetrievalFitScoreMetrics(BaseModel):
+    scores: dict[str, float]
+    chunk_ids: list[str]
+
+
+class RetrievalFitStats(BaseModel):
+    fit_score_lift: float
+    rerank_effect: float
+    fit_scores: dict[str, RetrievalFitScoreMetrics]
+
+
+class AgentChunkScores(BaseModel):
+    scores: dict[str, dict[str, list[int | float]]]
+
+
+class AgentChunkStats(BaseModel):
+    verified_count: int | None
+    verified_avg_scores: float | None
+    rejected_count: int | None
+    rejected_avg_scores: float | None
+    verified_doc_chunk_ids: list[str]
+    dismissed_doc_chunk_ids: list[str]
+
+
+class InitialAgentResultStats(BaseModel):
+    sub_questions: dict[str, float | int | None]
+    original_question: dict[str, float | int | None]
+    agent_effectiveness: dict[str, float | int | None]
+
+
+class RefinedAgentStats(BaseModel):
+    revision_doc_efficiency: float | None
+    revision_question_efficiency: float | None
+
+
+class Term(BaseModel):
+    term_name: str
+    term_type: str
+    term_similar_to: list[str]
+
+
+### Models ###
+
+
+class Entity(BaseModel):
+    entity_name: str
+    entity_type: str
+
+
+class Relationship(BaseModel):
+    relationship_name: str
+    relationship_type: str
+    relationship_entities: list[str]
+
+
+class EntityRelationshipTermExtraction(BaseModel):
+    entities: list[Entity]
+    relationships: list[Relationship]
+    terms: list[Term]
+
+
+### Models ###
+
+
+class QueryResult(BaseModel):
+    query: str
+    search_results: list[InferenceSection]
+    stats: RetrievalFitStats | None
+    query_info: SearchQueryInfo | None
+
+
+class QuestionAnswerResults(BaseModel):
+    question: str
+    question_id: str
+    answer: str
+    quality: str
+    expanded_retrieval_results: list[QueryResult]
+    documents: list[InferenceSection]
+    sub_question_retrieval_stats: AgentChunkStats
+
+
+class CombinedAgentMetrics(BaseModel):
+    timings: AgentTimings
+    base_metrics: AgentBaseMetrics | None
+    refined_metrics: AgentRefinedMetrics
+    additional_metrics: AgentAdditionalMetrics

backend/onyx/agent_search/shared_graph_utils/operators.py

@@ -0,0 +1,31 @@
+from onyx.agent_search.shared_graph_utils.models import (
+    QuestionAnswerResults,
+)
+from onyx.chat.prune_and_merge import _merge_sections
+from onyx.context.search.models import InferenceSection
+
+
+def dedup_inference_sections(
+    list1: list[InferenceSection], list2: list[InferenceSection]
+) -> list[InferenceSection]:
+    deduped = _merge_sections(list1 + list2)
+    return deduped
+
+
+def dedup_question_answer_results(
+    question_answer_results_1: list[QuestionAnswerResults],
+    question_answer_results_2: list[QuestionAnswerResults],
+) -> list[QuestionAnswerResults]:
+    deduped_question_answer_results: list[
+        QuestionAnswerResults
+    ] = question_answer_results_1
+    utilized_question_ids: set[str] = set(
+        [x.question_id for x in question_answer_results_1]
+    )
+
+    for question_answer_result in question_answer_results_2:
+        if question_answer_result.question_id not in utilized_question_ids:
+            deduped_question_answer_results.append(question_answer_result)
+            utilized_question_ids.add(question_answer_result.question_id)
+
+    return deduped_question_answer_results

backend/onyx/agent_search/shared_graph_utils/prompts.py

@@ -0,0 +1,912 @@
+UNKNOWN_ANSWER = "I do not have enough information to answer this question"
+
+NO_RECOVERED_DOCS = "No relevant documents recovered"
+
+REWRITE_PROMPT_MULTI_ORIGINAL = """ \n
+    Please convert an initial user question into a 2-3 more appropriate short and pointed search queries for retrievel from a
+    document store. Particularly, try to think about resolving ambiguities and make the search queries more specific,
+    enabling the system to search more broadly.
+    Also, try to make the search queries not redundant, i.e. not too similar! \n\n
+    Here is the initial question:
+    \n ------- \n
+    {question}
+    \n ------- \n
+    Formulate the queries separated by newlines (Do not say 'Query 1: ...', just write the querytext) as follows:
+<query 1>
+<query 2>
+...
+    queries: """
+
+REWRITE_PROMPT_MULTI = """ \n
+    Please create a list of 2-3 sample documents that could answer an original question. Each document
+    should be about as long as the original question. \n
+    Here is the initial question:
+    \n ------- \n
+    {question}
+    \n ------- \n
+    Formulate the sample documents separated by '--' (Do not say 'Document 1: ...', just write the text): """
+
+# The prompt is only used if there is no persona prompt, so the placeholder is ''
+BASE_RAG_PROMPT = (
+    """ \n
+    {persona_specification}
+    Use the context provided below - and only the
+    provided context - to answer the given question. (Note that the answer is in service of anserwing a broader
+    question, given below as 'motivation'.)
+
+    Again, only use the provided context and do not use your internal knowledge! If you cannot answer the
+    question based on the context, say """
+    + f'"{UNKNOWN_ANSWER}"'
+    + """. It is a matter of life and death that you do NOT
+    use your internal knowledge, just the provided information!
+
+    Make sure that you keep all relevant information, specifically as it concerns to the ultimate goal.
+    (But keep other details as well.)
+
+    \nContext:\n {context} \n
+
+    Motivation:\n {original_question} \n\n
+    \n\n
+    And here is the question I want you to answer based on the context above (with the motivation in mind):
+    \n--\n {question} \n--\n
+    """
+)
+
+BASE_RAG_PROMPT_v2 = (
+    """ \n
+    Use the context provided below - and only the
+    provided context - to answer the given question. (Note that the answer is in service of answering a broader
+    question, given below as 'motivation'.)
+
+    Again, only use the provided context and do not use your internal knowledge! If you cannot answer the
+    question based on the context, say """
+    + f'"{UNKNOWN_ANSWER}"'
+    + """. It is a matter of life and death that you do NOT
+    use your internal knowledge, just the provided information!
+
+    Make sure that you keep all relevant information, specifically as it concerns to the ultimate goal.
+    (But keep other details as well.)
+
+    Please remember to provide inline citations in the format [[D1]](), [[D2]](), [[D3]](), etc.
+    Proper citations are very important to the user!\n\n\n
+
+    For your general information, here is the ultimate motivation:
+    \n--\n {original_question} \n--\n
+    \n\n
+    And here is the actual question I want you to answer based on the context above (with the motivation in mind):
+    \n--\n {question} \n--\n
+
+    Here is the context:
+    \n\n\n--\n {context} \n--\n
+    """
+)
+
+SUB_CHECK_YES = "yes"
+SUB_CHECK_NO = "no"
+
+SUB_CHECK_PROMPT = (
+    """
+    Your task is to see whether a given answer addresses a given question.
+    Please do not use any internal knowledge you may have - just focus on whether the answer
+    as given seems to largely address the question as given, or at least addresses part of the question.
+    Here is the question:
+    \n ------- \n
+    {question}
+    \n ------- \n
+    Here is the suggested answer:
+    \n ------- \n
+    {base_answer}
+    \n ------- \n
+    Does the suggested answer address the question? Please answer with """
+    + f'"{SUB_CHECK_YES}" or "{SUB_CHECK_NO}".'
+)
+
+
+BASE_CHECK_PROMPT = """ \n
+    Please check whether 1) the suggested answer seems to fully address the original question AND 2)the
+    original question requests a simple, factual answer, and there are no ambiguities, judgements,
+    aggregations, or any other complications that may require extra context. (I.e., if the question is
+    somewhat addressed, but the answer would benefit from more context, then answer with 'no'.)
+
+    Please only answer with 'yes' or 'no' \n
+    Here is the initial question:
+    \n ------- \n
+    {question}
+    \n ------- \n
+    Here is the proposed answer:
+    \n ------- \n
+    {initial_answer}
+    \n ------- \n
+    Please answer with yes or no:"""
+
+VERIFIER_PROMPT = """
+You are supposed to judge whether a document text contains data or information that is potentially relevant for a question.
+
+Here is a document text that you can take as a fact:
+--
+DOCUMENT INFORMATION:
+{document_content}
+--
+
+Do you think that this information is useful and relevant to answer the following question?
+(Other documents may supply additional information, so do not worry if the provided information
+is not enough to answer the question, but it needs to be relevant to the question.)
+--
+QUESTION:
+{question}
+--
+
+Please answer with 'yes' or 'no':
+
+Answer:
+
+"""
+
+INITIAL_DECOMPOSITION_PROMPT_BASIC = """ \n
+If you think it is helpful, please decompose an initial user question into not more
+than 4 appropriate sub-questions that help to answer the original question.
+The purpose for this decomposition is to isolate individulal entities
+(i.e., 'compare sales of company A and company B' -> 'what are sales for company A' + 'what are sales
+for company B'), split ambiguous terms (i.e., 'what is our success with company A' -> 'what are our
+sales with company A' + 'what is our market share with company A' + 'is company A a reference customer
+ for us'), etc. Each sub-question should be realistically be answerable by a good RAG system.
+
+Importantly, if you think it is not needed or helpful, please just return an empty list. That is ok too.
+
+Here is the initial question:
+\n ------- \n
+{question}
+\n ------- \n
+
+Please formulate your answer as a list of subquestions:
+
+Answer:
+"""
+
+REWRITE_PROMPT_SINGLE = """ \n
+    Please convert an initial user question into a more appropriate search query for retrievel from a
+    document store. \n
+    Here is the initial question:
+    \n ------- \n
+    {question}
+    \n ------- \n
+
+    Formulate the query: """
+
+MODIFIED_RAG_PROMPT = (
+    """You are an assistant for question-answering tasks. Use the context provided below
+    - and only this context - to answer the question. It is a matter of life and death that you do NOT
+    use your internal knowledge, just the provided information!
+    If you don't have enough infortmation to generate an answer, just say """
+    + f'"{UNKNOWN_ANSWER}"'
+    + """.
+    Use three sentences maximum and keep the answer concise.
+    Pay also particular attention to the sub-questions and their answers, at least it may enrich the answer.
+    Again, only use the provided context and do not use your internal knowledge!
+
+    \nQuestion: {question}
+    \nContext: {combined_context} \n
+
+    Answer:"""
+)
+
+ORIG_DEEP_DECOMPOSE_PROMPT = """ \n
+    An initial user question needs to be answered. An initial answer has been provided but it wasn't quite
+    good enough. Also, some sub-questions had been answered and this information has been used to provide
+    the initial answer. Some other subquestions may have been suggested based on little knowledge, but they
+    were not directly answerable. Also, some entities, relationships and terms are givenm to you so that
+    you have an idea of how the avaiolable data looks like.
+
+    Your role is to generate 3-5 new sub-questions that would help to answer the initial question,
+    considering:
+
+    1) The initial question
+    2) The initial answer that was found to be unsatisfactory
+    3) The sub-questions that were answered
+    4) The sub-questions that were suggested but not answered
+    5) The entities, relationships and terms that were extracted from the context
+
+    The individual questions should be answerable by a good RAG system.
+    So a good idea would be to use the sub-questions to resolve ambiguities and/or to separate the
+    question for different entities that may be involved in the original question, but in a way that does
+    not duplicate questions that were already tried.
+
+    Additional Guidelines:
+    - The sub-questions should be specific to the question and provide richer context for the question,
+    resolve ambiguities, or address shortcoming of the initial answer
+    - Each sub-question - when answered - should be relevant for the answer to the original question
+    - The sub-questions should be free from comparisions, ambiguities,judgements, aggregations, or any
+    other complications that may require extra context.
+    - The sub-questions MUST have the full context of the original question so that it can be executed by
+    a RAG system independently without the original question available
+      (Example:
+        - initial question: "What is the capital of France?"
+        - bad sub-question: "What is the name of the river there?"
+        - good sub-question: "What is the name of the river that flows through Paris?"
+    - For each sub-question, please provide a short explanation for why it is a good sub-question. So
+    generate a list of dictionaries with the following format:
+      [{{"sub_question": <sub-question>, "explanation": <explanation>, "search_term": <rewrite the
+      sub-question using as a search phrase for the document store>}}, ...]
+
+    \n\n
+    Here is the initial question:
+    \n ------- \n
+    {question}
+    \n ------- \n
+
+    Here is the initial sub-optimal answer:
+    \n ------- \n
+    {base_answer}
+    \n ------- \n
+
+    Here are the sub-questions that were answered:
+    \n ------- \n
+    {answered_sub_questions}
+    \n ------- \n
+
+    Here are the sub-questions that were suggested but not answered:
+    \n ------- \n
+    {failed_sub_questions}
+    \n ------- \n
+
+    And here are the entities, relationships and terms extracted from the context:
+    \n ------- \n
+    {entity_term_extraction_str}
+    \n ------- \n
+
+   Please generate the list of good, fully contextualized sub-questions that would help to address the
+   main question. Again, please find questions that are NOT overlapping too much with the already answered
+   sub-questions or those that already were suggested and failed.
+   In other words - what can we try in addition to what has been tried so far?
+
+   Please think through it step by step and then generate the list of json dictionaries with the following
+   format:
+
+   {{"sub_questions": [{{"sub_question": <sub-question>,
+        "explanation": <explanation>,
+        "search_term": <rewrite the sub-question using as a search phrase for the document store>}},
+        ...]}} """
+
+DEEP_DECOMPOSE_PROMPT = """ \n
+    An initial user question needs to be answered. An initial answer has been provided but it wasn't quite
+    good enough. Also, some sub-questions had been answered and this information has been used to provide
+    the initial answer. Some other subquestions may have been suggested based on little knowledge, but they
+    were not directly answerable. Also, some entities, relationships and terms are givenm to you so that
+    you have an idea of how the avaiolable data looks like.
+
+    Your role is to generate 2-4 new sub-questions that would help to answer the initial question,
+    considering:
+
+    1) The initial question
+    2) The initial answer that was found to be unsatisfactory
+    3) The sub-questions that were answered
+    4) The sub-questions that were suggested but not answered
+    5) The entities, relationships and terms that were extracted from the context
+
+    The individual questions should be answerable by a good RAG system.
+    So a good idea would be to use the sub-questions to resolve ambiguities and/or to separate the
+    question for different entities that may be involved in the original question, but in a way that does
+    not duplicate questions that were already tried.
+
+    Additional Guidelines:
+    - The sub-questions should be specific to the question and provide richer context for the question,
+    resolve ambiguities, or address shortcoming of the initial answer
+    - Each sub-question - when answered - should be relevant for the answer to the original question
+    - The sub-questions should be free from comparisions, ambiguities,judgements, aggregations, or any
+    other complications that may require extra context.
+    - The sub-questions MUST have the full context of the original question so that it can be executed by
+    a RAG system independently without the original question available
+      (Example:
+        - initial question: "What is the capital of France?"
+        - bad sub-question: "What is the name of the river there?"
+        - good sub-question: "What is the name of the river that flows through Paris?"
+    - For each sub-question, please also provide a search term that can be used to retrieve relevant
+    documents from a document store.
+    - Consider specifically the sub-questions that were suggested but not answered. This is a sign that they are not
+    answerable with the available context, and you should not ask similar questions.
+    \n\n
+    Here is the initial question:
+    \n ------- \n
+    {question}
+    \n ------- \n
+
+    Here is the initial sub-optimal answer:
+    \n ------- \n
+    {base_answer}
+    \n ------- \n
+
+    Here are the sub-questions that were answered:
+    \n ------- \n
+    {answered_sub_questions}
+    \n ------- \n
+
+    Here are the sub-questions that were suggested but not answered:
+    \n ------- \n
+    {failed_sub_questions}
+    \n ------- \n
+
+    And here are the entities, relationships and terms extracted from the context:
+    \n ------- \n
+    {entity_term_extraction_str}
+    \n ------- \n
+
+   Please generate the list of good, fully contextualized sub-questions that would help to address the
+   main question.
+
+   Specifically pay attention also to the entities, relationships and terms extracted, as these indicate what type of
+   objects/relationships/terms you can ask about! Do not ask about entities, terms or relationships that are not
+   mentioned in the 'entities, relationships and terms' section.
+
+   Again, please find questions that are NOT overlapping too much with the already answered
+   sub-questions or those that already were suggested and failed.
+   In other words - what can we try in addition to what has been tried so far?
+
+   Generate the list of questions separated by one new line like this:
+<sub-question 1>
+<sub-question 2>
+<sub-question 3>
+   ...
+   """
+
+DECOMPOSE_PROMPT = """ \n
+    For an initial user question, please generate at 5-10 individual sub-questions whose answers would help
+    \n to answer the initial question. The individual questions should be answerable by a good RAG system.
+    So a good idea would be to \n use the sub-questions to resolve ambiguities and/or to separate the
+    question for different entities that may be involved in the original question.
+
+    In order to arrive at meaningful sub-questions, please also consider the context retrieved from the
+    document store, expressed as entities, relationships and terms. You can also think about the types
+    mentioned in brackets
+
+    Guidelines:
+    - The sub-questions should be specific to the question and provide richer context for the question,
+    and or resolve ambiguities
+    - Each sub-question - when answered - should be relevant for the answer to the original question
+    - The sub-questions should be free from comparisions, ambiguities,judgements, aggregations, or any
+    other complications that may require extra context.
+    - The sub-questions MUST have the full context of the original question so that it can be executed by
+    a RAG system independently without the original question available
+      (Example:
+        - initial question: "What is the capital of France?"
+        - bad sub-question: "What is the name of the river there?"
+        - good sub-question: "What is the name of the river that flows through Paris?"
+    - For each sub-question, please provide a short explanation for why it is a good sub-question. So
+    generate a list of dictionaries with the following format:
+      [{{"sub_question": <sub-question>, "explanation": <explanation>}}, ...]
+
+    \n\n
+    Here is the initial question:
+    \n ------- \n
+    {question}
+    \n ------- \n
+
+    And here are the entities, relationships and terms extracted from the context:
+    \n ------- \n
+    {entity_term_extraction_str}
+    \n ------- \n
+
+   Please generate the list of good, fully contextualized sub-questions that would help to address the
+   main question. Don't be too specific unless the original question is specific.
+   Please think through it step by step and then generate the list of json dictionaries with the following
+   format:
+   {{"sub_questions": [{{"sub_question": <sub-question>,
+        "explanation": <explanation>,
+        "search_term": <rewrite the sub-question using as a search phrase for the document store>}},
+        ...]}} """
+
+#### Consolidations
+COMBINED_CONTEXT = """-------
+    Below you will find useful information to answer the original question. First, you see a number of
+    sub-questions with their answers. This information should be considered to be more focussed and
+    somewhat more specific to the original question as it tries to contextualized facts.
+    After that will see the documents that were considered to be relevant to answer the original question.
+
+    Here are the sub-questions and their answers:
+    \n\n {deep_answer_context} \n\n
+    \n\n Here are the documents that were considered to be relevant to answer the original question:
+    \n\n {formated_docs} \n\n
+    ----------------
+    """
+
+SUB_QUESTION_EXPLANATION_RANKER_PROMPT = """-------
+    Below you will find a question that we ultimately want to answer (the original question) and a list of
+    motivations in arbitrary order for generated sub-questions that are supposed to help us answering the
+    original question. The motivations are formatted as <motivation number>:  <motivation explanation>.
+    (Again, the numbering is arbitrary and does not necessarily mean that 1 is the most relevant
+    motivation and 2 is less relevant.)
+
+    Please rank the motivations in order of relevance for answering the original question. Also, try to
+    ensure that the top questions do not duplicate too much, i.e. that they are not too similar.
+    Ultimately, create a list with the motivation numbers where the number of the most relevant
+    motivations comes first.
+
+    Here is the original question:
+    \n\n {original_question} \n\n
+    \n\n Here is the list of sub-question motivations:
+    \n\n {sub_question_explanations} \n\n
+    ----------------
+
+    Please think step by step and then generate the ranked list of motivations.
+
+    Please format your answer as a json object in the following format:
+    {{"reasonning": <explain your reasoning for the ranking>,
+      "ranked_motivations": <ranked list of motivation numbers>}}
+    """
+
+
+INITIAL_DECOMPOSITION_PROMPT_QUESTIONS = """
+If you think it is helpful, please decompose an initial user question into no more than 3 appropriate sub-questions that help to
+answer the original question. The purpose for this decomposition may be to
+  1) isolate individual entities (i.e., 'compare sales of company A and company B' -> ['what are sales for company A',
+     'what are sales for company B')]
+  2) clarify or disambiguate ambiguous terms (i.e., 'what is our success with company A' -> ['what are our sales with company A',
+      'what is our market share with company A', 'is company A a reference customer for us', etc.])
+  3) if a term or a metric is essentially clear, but it could relate to various components of an entity and you are generally
+    familiar with the entity, then you can decompose the question into sub-questions that are more specific to components
+     (i.e., 'what do we do to improve scalability of product X', 'what do we to to improve scalability of product X',
+     'what do we do to improve stability of product X', ...])
+  4) research an area that could really help to answer the question. (But clarifications or disambiguations are more important.)
+
+If you think that a decomposition is not needed or helpful, please just return an empty string. That is ok too.
+
+Here is the initial question:
+-------
+{question}
+-------
+Please formulate your answer as a newline-separated list of questions like so:
+ <sub-question>
+ <sub-question>
+ <sub-question>
+
+Answer:"""
+
+INITIAL_DECOMPOSITION_PROMPT_QUESTIONS_AFTER_SEARCH = """
+If you think it is helpful, please decompose an initial user question into no more than 3 appropriate sub-questions that help to
+answer the original question. The purpose for this decomposition may be to
+  1) isolate individual entities (i.e., 'compare sales of company A and company B' -> ['what are sales for company A',
+     'what are sales for company B')]
+  2) clarify or disambiguate ambiguous terms (i.e., 'what is our success with company A' -> ['what are our sales with company A',
+      'what is our market share with company A', 'is company A a reference customer for us', etc.])
+  3) if a term or a metric is essentially clear, but it could relate to various components of an entity and you are generally
+    familiar with the entity, then you can decompose the question into sub-questions that are more specific to components
+     (i.e., 'what do we do to improve scalability of product X', 'what do we to to improve scalability of product X',
+     'what do we do to improve stability of product X', ...])
+  4) research an area that could really help to answer the question. (But clarifications or disambiguations are more important.)
+
+Here are some other ruleds:
+
+1) To give you some context, you will see below also some documents that relate to the question. Please only
+use this information to learn what the question is approximately asking about, but do not focus on the details
+to construct the sub-questions.
+2) If you think that a decomposition is not needed or helpful, please just return an empty string. That is very muchok too.
+
+Here are the sampple docs to give you some context:
+-------
+{sample_doc_str}
+-------
+
+And here is the initial question that you should think about decomposing:
+-------
+{question}
+-------
+
+
+Please formulate your answer as a newline-separated list of questions like so:
+ <sub-question>
+ <sub-question>
+ <sub-question>
+
+Answer:"""
+
+INITIAL_DECOMPOSITION_PROMPT = """ \n
+    Please decompose an initial user question into 2 or 3 appropriate sub-questions that help to
+    answer the original question. The purpose for this decomposition is to isolate individulal entities
+    (i.e., 'compare sales of company A and company B' -> 'what are sales for company A' + 'what are sales
+    for company B'), split ambiguous terms (i.e., 'what is our success with company A' -> 'what are our
+    sales with company A' + 'what is our market share with company A' + 'is company A a reference customer
+    for us'), etc. Each sub-question should be realistically be answerable by a good RAG system. \n
+
+    For each sub-question, please also create one search term that can be used to retrieve relevant
+    documents from a document store.
+
+    Here is the initial question:
+    \n ------- \n
+    {question}
+    \n ------- \n
+
+    Please formulate your answer as a list of json objects with the following format:
+
+   [{{"sub_question": <sub-question>, "search_term": <search term>}}, ...]
+
+    Answer:
+    """
+
+INITIAL_RAG_BASE_PROMPT = (
+    """ \n
+You are an assistant for question-answering tasks. Use the information provided below - and only the
+provided information - to answer the provided question.
+
+The information provided below consists ofa number of documents that were deemed relevant for the question.
+
+IMPORTANT RULES:
+- If you cannot reliably answer the question solely using the provided information, say that you cannot reliably answer.
+You may give some additional facts you learned, but do not try to invent an answer.
+- If the information is empty or irrelevant, just say """
+    + f'"{UNKNOWN_ANSWER}"'
+    + """.
+- If the information is relevant but not fully conclusive, specify that the information is not conclusive and say why.
+
+Try to keep your answer concise.
+
+Here is the contextual information from the document store:
+\n ------- \n
+{context} \n\n\n
+\n ------- \n
+And here is the question I want you to answer based on the context above (with the motivation in mind):
+\n--\n {question} \n--\n
+Answer:"""
+)
+
+
+AGENT_DECISION_PROMPT = """
+You are an large language model assistant helping users address their information needs. You are tasked with deciding
+whether to use a thorough agent search ('research') of a document store to answer a question or request, or whether you want to
+address the question or request yourself as an LLM.
+
+Here are some rules:
+- If you think that a thorough search through a document store will help answer the question
+or address the request, you should choose the 'research' option.
+- If the question asks you do do somethng ('please create...', 'write for me...', etc.), you should choose the 'LLM' option.
+- If you think the question is very general and does not refer to a contents of a document store, you should choose
+the 'LLM' option.
+- Otherwise, you should choose the 'research' option.
+
+Here is the initial question:
+-------
+{question}
+-------
+
+Please decide whether to use the agent search or the LLM to answer the question. Choose from two choices,
+'research' or 'LLM'.
+
+Answer:"""
+
+AGENT_DECISION_PROMPT_AFTER_SEARCH = """
+You are an large language model assistant helping users address their information needs.  You are given an initial question
+or request and very few sample of documents that a preliminary and fast search from a document store returned.
+You are tasked with deciding whether to use a thorough agent search ('research') of the document store to answer a question
+or request, or whether you want to address the question or request yourself as an LLM.
+
+Here are some rules:
+- If based on the retrieved documents you think there may be useful information in the document
+store to answer or materially help with the request, you should choose the 'research' option.
+- If you think that the retrieved document do not help to answer the question or do not help with the request, AND
+you know the answer/can handle the request, you should choose the 'LLM' option.
+- If the question asks you do do somethng ('please create...', 'write for me...', etc.), you should choose the 'LLM' option.
+- If in doubt, choose the 'research' option.
+
+Here is the initial question:
+-------
+{question}
+-------
+
+Here is the sample of documents that were retrieved from a document store:
+-------
+{sample_doc_str}
+-------
+
+Please decide whether to use the agent search ('research') or the LLM to answer the question. Choose from two choices,
+'research' or 'LLM'.
+
+Answer:"""
+
+### ANSWER GENERATION PROMPTS
+
+# Persona specification
+ASSISTANT_SYSTEM_PROMPT_DEFAULT = """
+You are an assistant for question-answering tasks."""
+
+ASSISTANT_SYSTEM_PROMPT_PERSONA = """
+You are an assistant for question-answering tasks. Here is more information about you:
+\n ------- \n
+{persona_prompt}
+\n ------- \n
+"""
+
+SUB_QUESTION_ANSWER_TEMPLATE = """
+    Sub-Question: Q{sub_question_nr}\n  Sub-Question:\n  - \n{sub_question}\n  --\nAnswer:\n  -\n {sub_answer}\n\n
+    """
+
+SUB_QUESTION_ANSWER_TEMPLATE_REVISED = """
+    Sub-Question: Q{sub_question_nr}\n  Type: {level_type}\n Sub-Question:\n
+- \n{sub_question}\n  --\nAnswer:\n  -\n {sub_answer}\n\n
+    """
+
+SUB_QUESTION_SEARCH_RESULTS_TEMPLATE = """
+    Sub-Question: Q{sub_question_nr}\n  Sub-Question:\n  - \n{sub_question}\n  --\nRelevant Documents:\n
+    -\n {formatted_sub_question_docs}\n\n
+    """
+
+INITIAL_RAG_PROMPT_SUB_QUESTION_SEARCH = (
+    """ \n
+{persona_specification}
+
+Use the information provided below - and only the
+provided information - to answer the provided question.
+
+The information provided below consists of:
+    1) a number of sub-questions and supporting document information that would help answer them.
+    2) a broader collection of documents that were deemed relevant for the question. These documents contain informattion
+    that was also provided in the sub-questions and often more.
+
+IMPORTANT RULES:
+ - If you cannot reliably answer the question solely using the provided information, say that you cannot reliably answer.
+ You may give some additional facts you learned, but do not try to invent an answer.
+ - If the information is empty or irrelevant, just say """
+    + f'"{UNKNOWN_ANSWER}"'
+    + """.
+ - If the information is relevant but not fully conclusive, specify that the information is not conclusive and say why.
+
+Please provide inline citations of documentsin the format [[D1]](), [[D2]](), [[D3]](), etc.,  If you have multiple citations,
+please cite for example as [[D1]]()[[D3]](), or [[D2]]()[[D4]](), etc. Feel free to cite documents in addition
+to the sub-questions! Proper citations are important for the final answer to be verifiable! \n\n\n
+
+Again, you should be sure that the answer is supported by the information provided!
+
+Try to keep your answer concise. But also highlight uncertainties you may have should there be substantial ones,
+or assumptions you made.
+
+Here is the contextual information:
+\n-------\n
+*Answered Sub-questions (these should really matter!):
+{answered_sub_questions}
+
+And here are relevant document information that support the sub-question answers, or that are relevant for the actual question:\n
+
+{relevant_docs}
+
+\n-------\n
+\n
+And here is the question I want you to answer based on the information above:
+\n--\n
+{question}
+\n--\n\n
+Answer:"""
+)
+
+
+DIRECT_LLM_PROMPT = """ \n
+{persona_specification}
+
+Please answer the following question/address the request:
+\n--\n
+{question}
+\n--\n\n
+Answer:"""
+
+INITIAL_RAG_PROMPT = (
+    """ \n
+{persona_specification}
+
+Use the information provided below - and only the
+provided information - to answer the provided question.
+
+The information provided below consists of:
+    1) a number of answered sub-questions - these are very important(!) and definitely should be
+    considered to answer the question.
+    2) a number of documents that were also deemed relevant for the question.
+
+IMPORTANT RULES:
+ - If you cannot reliably answer the question solely using the provided information, say that you cannot reliably answer.
+ You may give some additional facts you learned, but do not try to invent an answer.
+ - If the information is empty or irrelevant, just say """
+    + f'"{UNKNOWN_ANSWER}"'
+    + """.
+ - If the information is relevant but not fully conclusive, specify that the information is not conclusive and say why.
+
+Remember to provide inline citations of documents in the format [[D1]](), [[D2]](), [[D3]](), etc., and [[Q1]](), [[Q2]](),... if
+you want to cite the answer to a sub-question. If you have multiple citations, please cite for example
+as [[D1]]()[[Q3]](), or [[D2]]()[[D4]](), etc. Feel free to cite sub-questions in addition to documents, but make sure that you
+have docuemnt citations ([[D7]]() etc.) if possible!
+
+Again, you should be sure that the answer is supported by the information provided!
+
+Try to keep your answer concise. But also highlight uncertainties you may have should there be substantial ones,
+or assumptions you made.
+
+Here is the contextual information:
+\n-------\n
+*Answered Sub-questions (these should really matter!):
+{answered_sub_questions}
+
+And here are relevant document information that support the sub-question answers, or that are relevant for the actual question:\n
+
+{relevant_docs}
+
+\n-------\n
+\n
+And here is the question I want you to answer based on the information above:
+\n--\n
+{question}
+\n--\n\n
+Answer:"""
+)
+
+# sub_question_answer_str is empty
+INITIAL_RAG_PROMPT_NO_SUB_QUESTIONS = (
+    """{answered_sub_questions}
+{persona_specification}
+Use the information provided below
+- and only the provided information - to answer the provided question.
+The information provided below consists of a number of documents that were deemed relevant for the question.
+
+IMPORTANT RULES:
+ - If you cannot reliably answer the question solely using the provided information, say that you cannot reliably answer.
+ You may give some  additional facts you learned, but do not try to invent an answer.
+ - If the information is irrelevant, just say """
+    + f'"{UNKNOWN_ANSWER}"'
+    + """.
+ - If the information is relevant but not fully conclusive, specify that the information is not conclusive and say why.
+
+Again, you should be sure that the answer is supported by the information provided!
+
+Remember to provide inline citations of documents in the format [[D1]](), [[D2]](), [[D3]](), etc.!
+
+Try to keep your answer concise.
+
+Here are is the relevant context information:
+\n-------\n
+{relevant_docs}
+\n-------\n
+
+And here is the question I want you to answer based on the context above
+\n--\n
+{question}
+\n--\n
+
+Answer:"""
+)
+
+REVISED_RAG_PROMPT = (
+    """\n
+{persona_specification}
+Use the information provided below - and only the
+provided information - to answer the provided question.
+
+The information provided below consists of:
+    1) an initial answer that was given but found to be lacking in some way.
+    2) a number of answered sub-questions - these are very important(!) and definitely should be
+    considered to answer the question. Note that the sub-questions have a type, 'initial' and 'revised'. The 'initial'
+     ones were available for the initial answer, and the 'revised' were not. So please use the 'revised' sub-questions in
+     particular to update/extend/correct the initial answer!
+    information from the revised sub-questions
+    3) a number of documents that were also deemed relevant for the question.
+
+IMPORTANT RULES:
+ - If you cannot reliably answer the question solely using the provided information, say that you cannot reliably answer.
+ You may give some additional facts you learned, but do not try to invent an answer.
+ - If the information is empty or irrelevant, just say """
+    + f'"{UNKNOWN_ANSWER}"'
+    + """.
+ - If the information is relevant but not fully conclusive, provide and answer to the extent you can but also
+ specify that the information is not conclusive and why.
+- Ignore the exisiting citations within the answered sub-questions, like [[D1]]()... and [[Q2]]()!
+The citations you will need to use will need to refer to the documents and sub-questions that you are explicitly
+presented with below!
+
+Again, you should be sure that the answer is supported by the information provided!
+
+Remember to provide inline citations of documents in the format [[D1]](), [[D2]](), [[D3]](), etc., and [[Q1]](), [[Q2]](),... if
+you want to cite the answer to a sub-question. If you have multiple citations, please cite for example
+as [[D1]]()[[Q3]](), or [[D2]]()[[D4]](), etc. Feel free to cite sub-questions in addition to documents, but make sure that you
+have docuemnt citations ([[D7]]() etc.) if possible!
+Proper citations are important for the final answer to be verifiable! \n\n\n
+
+Try to keep your answer concise. But also highlight uncertainties you may have should there be substantial ones,
+or assumptions you made.
+
+Here is the contextual information:
+\n-------\n
+
+*Initial Answer that was found to be lacking:
+{initial_answer}
+
+*Answered Sub-questions (these should really matter! They also contain questions/answers that were not available when the original
+answer was constructed):
+{answered_sub_questions}
+
+And here are relevant document information that support the sub-question answers, or that are relevant for the actual question:\n
+
+{relevant_docs}
+
+\n-------\n
+\n
+Lastly, here is the question I want you to answer based on the information above:
+\n--\n
+{question}
+\n--\n\n
+Answer:"""
+)
+
+# sub_question_answer_str is empty
+REVISED_RAG_PROMPT_NO_SUB_QUESTIONS = (
+    """{answered_sub_questions}\n
+{persona_specification}
+Use the information provided below - and only the
+provided information - to answer the provided question.
+
+The information provided below consists of:
+    1) an initial answer that was given but found to be lacking in some way.
+    2) a number of documents that were also deemed relevant for the question.
+
+IMPORTANT RULES:
+ - If you cannot reliably answer the question solely using the provided information, say that you cannot reliably answer.
+ You may give some additional facts you learned, but do not try to invent an answer.
+ - If the information is empty or irrelevant, just say """
+    + f'"{UNKNOWN_ANSWER}"'
+    + """.
+ - If the information is relevant but not fully conclusive, provide and answer to the extent you can but also
+ specify that the information is not conclusive and why.
+
+Again, you should be sure that the answer is supported by the information provided!
+
+Remember to provide inline citations of documents in the format [[D1]](), [[D2]](), [[D3]](), etc.
+
+Try to keep your answer concise. But also highlight uncertainties you may have should there be substantial ones,
+or assumptions you made.
+
+Here is the contextual information:
+\n-------\n
+
+*Initial Answer that was found to be lacking:
+{initial_answer}
+
+And here are relevant document information that support the sub-question answers, or that are relevant for the actual question:\n
+
+{relevant_docs}
+
+\n-------\n
+\n
+Lastly, here is the question I want you to answer based on the information above:
+\n--\n
+{question}
+\n--\n\n
+Answer:"""
+)
+
+
+ENTITY_TERM_PROMPT = """ \n
+    Based on the original question and the context retieved from a dataset, please generate a list of
+    entities (e.g. companies, organizations, industries, products, locations, etc.), terms and concepts
+    (e.g. sales, revenue, etc.) that are relevant for the question, plus their relations to each other.
+
+    \n\n
+    Here is the original question:
+    \n ------- \n
+    {question}
+    \n ------- \n
+   And here is the context retrieved:
+    \n ------- \n
+    {context}
+    \n ------- \n
+
+    Please format your answer as a json object in the following format:
+
+    {{"retrieved_entities_relationships": {{
+        "entities": [{{
+            "entity_name": <assign a name for the entity>,
+            "entity_type": <specify a short type name for the entity, such as 'company', 'location',...>
+        }}],
+        "relationships": [{{
+            "relationship_name": <assign a name for the relationship>,
+            "relationship_type": <specify a short type name for the relationship, such as 'sales_to', 'is_location_of',...>,
+            "relationship_entities": [<related entity name 1>, <related entity name 2>, ...]
+        }}],
+        "terms": [{{
+            "term_name": <assign a name for the term>,
+            "term_type": <specify a short type name for the term, such as 'revenue', 'market_share',...>,
+            "term_similar_to": <list terms that are similar to this term>
+        }}]
+    }}
+    }}
+   """

backend/onyx/agent_search/shared_graph_utils/utils.py

@@ -0,0 +1,264 @@
+import ast
+import json
+import re
+from collections.abc import Callable
+from collections.abc import Iterator
+from collections.abc import Sequence
+from datetime import datetime
+from datetime import timedelta
+from typing import Any
+from typing import cast
+from uuid import UUID
+
+from langchain_core.messages import BaseMessage
+from sqlalchemy.orm import Session
+
+from onyx.agent_search.models import ProSearchConfig
+from onyx.agent_search.shared_graph_utils.models import EntityRelationshipTermExtraction
+from onyx.chat.models import AnswerStyleConfig
+from onyx.chat.models import CitationConfig
+from onyx.chat.models import DocumentPruningConfig
+from onyx.chat.models import PromptConfig
+from onyx.configs.chat_configs import CHAT_TARGET_CHUNK_PERCENTAGE
+from onyx.configs.chat_configs import MAX_CHUNKS_FED_TO_CHAT
+from onyx.configs.constants import DEFAULT_PERSONA_ID
+from onyx.context.search.enums import LLMEvaluationType
+from onyx.context.search.models import InferenceSection
+from onyx.context.search.models import RetrievalDetails
+from onyx.context.search.models import SearchRequest
+from onyx.db.persona import get_persona_by_id
+from onyx.db.persona import Persona
+from onyx.llm.interfaces import LLM
+from onyx.tools.tool_constructor import SearchToolConfig
+from onyx.tools.tool_implementations.search.search_tool import SearchTool
+
+
+def normalize_whitespace(text: str) -> str:
+    """Normalize whitespace in text to single spaces and strip leading/trailing whitespace."""
+    import re
+
+    return re.sub(r"\s+", " ", text.strip())
+
+
+# Post-processing
+def format_docs(docs: Sequence[InferenceSection]) -> str:
+    formatted_doc_list = []
+
+    for doc_nr, doc in enumerate(docs):
+        formatted_doc_list.append(f"Document D{doc_nr + 1}:\n{doc.combined_content}")
+
+    return "\n\n".join(formatted_doc_list)
+
+
+def format_docs_content_flat(docs: Sequence[InferenceSection]) -> str:
+    formatted_doc_list = []
+
+    for _, doc in enumerate(docs):
+        formatted_doc_list.append(f"\n...{doc.combined_content}\n")
+
+    return "\n\n".join(formatted_doc_list)
+
+
+def clean_and_parse_list_string(json_string: str) -> list[dict]:
+    # Remove any prefixes/labels before the actual JSON content
+    json_string = re.sub(r"^.*?(?=\[)", "", json_string, flags=re.DOTALL)
+
+    # Remove markdown code block markers and any newline prefixes
+    cleaned_string = re.sub(r"```json\n|\n```", "", json_string)
+    cleaned_string = cleaned_string.replace("\\n", " ").replace("\n", " ")
+    cleaned_string = " ".join(cleaned_string.split())
+
+    # Try parsing with json.loads first, fall back to ast.literal_eval
+    try:
+        return json.loads(cleaned_string)
+    except json.JSONDecodeError:
+        try:
+            return ast.literal_eval(cleaned_string)
+        except (ValueError, SyntaxError) as e:
+            raise ValueError(f"Failed to parse JSON string: {cleaned_string}") from e
+
+
+def clean_and_parse_json_string(json_string: str) -> dict[str, Any]:
+    # Remove markdown code block markers and any newline prefixes
+    cleaned_string = re.sub(r"```json\n|\n```", "", json_string)
+    cleaned_string = cleaned_string.replace("\\n", " ").replace("\n", " ")
+    cleaned_string = " ".join(cleaned_string.split())
+    # Parse the cleaned string into a Python dictionary
+    return json.loads(cleaned_string)
+
+
+def format_entity_term_extraction(
+    entity_term_extraction_dict: EntityRelationshipTermExtraction,
+) -> str:
+    entities = entity_term_extraction_dict.entities
+    terms = entity_term_extraction_dict.terms
+    relationships = entity_term_extraction_dict.relationships
+
+    entity_strs = ["\nEntities:\n"]
+    for entity in entities:
+        entity_str = f"{entity.entity_name} ({entity.entity_type})"
+        entity_strs.append(entity_str)
+
+    entity_str = "\n - ".join(entity_strs)
+
+    relationship_strs = ["\n\nRelationships:\n"]
+    for relationship in relationships:
+        relationship_name = relationship.relationship_name
+        relationship_type = relationship.relationship_type
+        relationship_entities = relationship.relationship_entities
+        relationship_str = (
+            f"""{relationship_name} ({relationship_type}): {relationship_entities}"""
+        )
+        relationship_strs.append(relationship_str)
+
+    relationship_str = "\n - ".join(relationship_strs)
+
+    term_strs = ["\n\nTerms:\n"]
+    for term in terms:
+        term_str = f"{term.term_name} ({term.term_type}): similar to {', '.join(term.term_similar_to)}"
+        term_strs.append(term_str)
+
+    term_str = "\n - ".join(term_strs)
+
+    return "\n".join(entity_strs + relationship_strs + term_strs)
+
+
+def _format_time_delta(time: timedelta) -> str:
+    seconds_from_start = f"{((time).seconds):03d}"
+    microseconds_from_start = f"{((time).microseconds):06d}"
+    return f"{seconds_from_start}.{microseconds_from_start}"
+
+
+def generate_log_message(
+    message: str,
+    node_start_time: datetime,
+    graph_start_time: datetime | None = None,
+) -> str:
+    current_time = datetime.now()
+
+    if graph_start_time is not None:
+        graph_time_str = _format_time_delta(current_time - graph_start_time)
+    else:
+        graph_time_str = "N/A"
+
+    node_time_str = _format_time_delta(current_time - node_start_time)
+
+    return f"{graph_time_str} ({node_time_str} s): {message}"
+
+
+def get_test_config(
+    db_session: Session, primary_llm: LLM, fast_llm: LLM, search_request: SearchRequest
+) -> tuple[ProSearchConfig, SearchTool]:
+    persona = get_persona_by_id(DEFAULT_PERSONA_ID, None, db_session)
+    document_pruning_config = DocumentPruningConfig(
+        max_chunks=int(
+            persona.num_chunks
+            if persona.num_chunks is not None
+            else MAX_CHUNKS_FED_TO_CHAT
+        ),
+        max_window_percentage=CHAT_TARGET_CHUNK_PERCENTAGE,
+    )
+
+    answer_style_config = AnswerStyleConfig(
+        citation_config=CitationConfig(
+            # The docs retrieved by this flow are already relevance-filtered
+            all_docs_useful=True
+        ),
+        document_pruning_config=document_pruning_config,
+        structured_response_format=None,
+    )
+
+    search_tool_config = SearchToolConfig(
+        answer_style_config=answer_style_config,
+        document_pruning_config=document_pruning_config,
+        retrieval_options=RetrievalDetails(),  # may want to set dedupe_docs=True
+        rerank_settings=None,  # Can use this to change reranking model
+        selected_sections=None,
+        latest_query_files=None,
+        bypass_acl=False,
+    )
+
+    prompt_config = PromptConfig.from_model(persona.prompts[0])
+
+    search_tool = SearchTool(
+        db_session=db_session,
+        user=None,
+        persona=persona,
+        retrieval_options=search_tool_config.retrieval_options,
+        prompt_config=prompt_config,
+        llm=primary_llm,
+        fast_llm=fast_llm,
+        pruning_config=search_tool_config.document_pruning_config,
+        answer_style_config=search_tool_config.answer_style_config,
+        selected_sections=search_tool_config.selected_sections,
+        chunks_above=search_tool_config.chunks_above,
+        chunks_below=search_tool_config.chunks_below,
+        full_doc=search_tool_config.full_doc,
+        evaluation_type=(
+            LLMEvaluationType.BASIC
+            if persona.llm_relevance_filter
+            else LLMEvaluationType.SKIP
+        ),
+        rerank_settings=search_tool_config.rerank_settings,
+        bypass_acl=search_tool_config.bypass_acl,
+    )
+
+    config = ProSearchConfig(
+        search_request=search_request,
+        # chat_session_id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+        chat_session_id=UUID("edda10d5-6cef-45d8-acfb-39317552a1f4"),  # Joachim
+        # chat_session_id=UUID("d1acd613-2692-4bc3-9d65-c6d3da62e58e"),  # Evan
+        message_id=1,
+        use_persistence=True,
+        primary_llm=primary_llm,
+        fast_llm=fast_llm,
+        search_tool=search_tool,
+    )
+
+    return config, search_tool
+
+
+def get_persona_prompt(persona: Persona | None) -> str:
+    if persona is None:
+        return ""
+    else:
+        return "\n".join([x.system_prompt for x in persona.prompts])
+
+
+def make_question_id(level: int, question_nr: int) -> str:
+    return f"{level}_{question_nr}"
+
+
+def parse_question_id(question_id: str) -> tuple[int, int]:
+    level, question_nr = question_id.split("_")
+    return int(level), int(question_nr)
+
+
+def _dispatch_nonempty(
+    content: str, dispatch_event: Callable[[str, int], None], num: int
+) -> None:
+    if content != "":
+        dispatch_event(content, num)
+
+
+def dispatch_separated(
+    token_itr: Iterator[BaseMessage],
+    dispatch_event: Callable[[str, int], None],
+    sep: str = "\n",
+) -> list[str | list[str | dict[str, Any]]]:
+    num = 1
+    streamed_tokens: list[str | list[str | dict[str, Any]]] = [""]
+    for message in token_itr:
+        content = cast(str, message.content)
+        if sep in content:
+            sub_question_parts = content.split(sep)
+            _dispatch_nonempty(sub_question_parts[0], dispatch_event, num)
+            num += 1
+            _dispatch_nonempty(
+                "".join(sub_question_parts[1:]).strip(), dispatch_event, num
+            )
+        else:
+            _dispatch_nonempty(content, dispatch_event, num)
+        streamed_tokens.append(content)
+
+    return streamed_tokens

backend/onyx/auth/noauth_user.py

@@ -23,6 +23,7 @@ def load_no_auth_user_preferences(store: KeyValueStore) -> UserPreferences:
         preferences_data = cast(
             Mapping[str, Any], store.load(KV_NO_AUTH_USER_PREFERENCES_KEY)
         )
+        print("preferences_data", preferences_data)
         return UserPreferences(**preferences_data)
     except KvKeyNotFoundError:
         return UserPreferences(

backend/onyx/background/celery/celery_utils.py

@@ -14,6 +14,7 @@ from onyx.connectors.interfaces import PollConnector
 from onyx.connectors.interfaces import SlimConnector
 from onyx.connectors.models import Document
 from onyx.db.connector_credential_pair import get_connector_credential_pair
+from onyx.db.enums import ConnectorCredentialPairStatus
 from onyx.db.enums import TaskStatus
 from onyx.db.models import TaskQueueState
 from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface
@@ -41,14 +42,21 @@ def _get_deletion_status(
         return None
 
     redis_connector = RedisConnector(tenant_id, cc_pair.id)
-    if not redis_connector.delete.fenced:
-        return None
+    if redis_connector.delete.fenced:
+        return TaskQueueState(
+            task_id="",
+            task_name=redis_connector.delete.fence_key,
+            status=TaskStatus.STARTED,
+        )
 
-    return TaskQueueState(
-        task_id="",
-        task_name=redis_connector.delete.fence_key,
-        status=TaskStatus.STARTED,
-    )
+    if cc_pair.status == ConnectorCredentialPairStatus.DELETING:
+        return TaskQueueState(
+            task_id="",
+            task_name=redis_connector.delete.fence_key,
+            status=TaskStatus.PENDING,
+        )
+
+    return None
 
 
 def get_deletion_attempt_snapshot(

backend/onyx/background/celery/tasks/connector_deletion/tasks.py

@@ -44,11 +44,11 @@ def check_for_connector_deletion_task(
         timeout=CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT,
     )
 
-    try:
-        # these tasks should never overlap
-        if not lock_beat.acquire(blocking=False):
-            return None
+    # these tasks should never overlap
+    if not lock_beat.acquire(blocking=False):
+        return None
 
+    try:
         # collect cc_pair_ids
         cc_pair_ids: list[int] = []
         with get_session_with_tenant(tenant_id) as db_session:

backend/onyx/background/celery/tasks/doc_permission_syncing/tasks.py

@@ -102,11 +102,11 @@ def check_for_doc_permissions_sync(self: Task, *, tenant_id: str | None) -> bool
         timeout=CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT,
     )
 
-    try:
-        # these tasks should never overlap
-        if not lock_beat.acquire(blocking=False):
-            return None
+    # these tasks should never overlap
+    if not lock_beat.acquire(blocking=False):
+        return None
 
+    try:
         # get all cc pairs that need to be synced
         cc_pair_ids_to_sync: list[int] = []
         with get_session_with_tenant(tenant_id) as db_session:

backend/onyx/background/celery/tasks/external_group_syncing/tasks.py

@@ -102,11 +102,11 @@ def check_for_external_group_sync(self: Task, *, tenant_id: str | None) -> bool
         timeout=CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT,
     )
 
-    try:
-        # these tasks should never overlap
-        if not lock_beat.acquire(blocking=False):
-            return None
+    # these tasks should never overlap
+    if not lock_beat.acquire(blocking=False):
+        return None
 
+    try:
         cc_pair_ids_to_sync: list[int] = []
         with get_session_with_tenant(tenant_id) as db_session:
             cc_pairs = get_all_auto_sync_cc_pairs(db_session)

backend/onyx/background/celery/tasks/indexing/tasks.py

@@ -63,6 +63,7 @@ from onyx.redis.redis_connector_index import RedisConnectorIndex
 from onyx.redis.redis_connector_index import RedisConnectorIndexPayload
 from onyx.redis.redis_pool import get_redis_client
 from onyx.redis.redis_pool import redis_lock_dump
+from onyx.redis.redis_pool import SCAN_ITER_COUNT_DEFAULT
 from onyx.utils.logger import setup_logger
 from onyx.utils.variable_functionality import global_version
 from shared_configs.configs import INDEXING_MODEL_SERVER_HOST
@@ -204,6 +205,10 @@ def get_unfenced_index_attempt_ids(db_session: Session, r: redis.Redis) -> list[
 def check_for_indexing(self: Task, *, tenant_id: str | None) -> int | None:
     """a lightweight task used to kick off indexing tasks.
     Occcasionally does some validation of existing state to clear up error conditions"""
+    debug_tenants = {
+        "tenant_i-043470d740845ec56",
+        "tenant_82b497ce-88aa-4fbd-841a-92cae43529c8",
+    }
     time_start = time.monotonic()
 
     tasks_created = 0
@@ -219,11 +224,11 @@ def check_for_indexing(self: Task, *, tenant_id: str | None) -> int | None:
         timeout=CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT,
     )
 
-    try:
-        # these tasks should never overlap
-        if not lock_beat.acquire(blocking=False):
-            return None
+    # these tasks should never overlap
+    if not lock_beat.acquire(blocking=False):
+        return None
 
+    try:
         locked = True
 
         # check for search settings swap
@@ -246,15 +251,25 @@ def check_for_indexing(self: Task, *, tenant_id: str | None) -> int | None:
                     )
 
         # gather cc_pair_ids
+        lock_beat.reacquire()
         cc_pair_ids: list[int] = []
         with get_session_with_tenant(tenant_id) as db_session:
-            lock_beat.reacquire()
             cc_pairs = fetch_connector_credential_pairs(db_session)
             for cc_pair_entry in cc_pairs:
                 cc_pair_ids.append(cc_pair_entry.id)
 
         # kick off index attempts
         for cc_pair_id in cc_pair_ids:
+            # debugging logic - remove after we're done
+            if tenant_id in debug_tenants:
+                ttl = redis_client.ttl(OnyxRedisLocks.CHECK_INDEXING_BEAT_LOCK)
+                task_logger.info(
+                    f"check_for_indexing cc_pair lock: "
+                    f"tenant={tenant_id} "
+                    f"cc_pair={cc_pair_id} "
+                    f"ttl={ttl}"
+                )
+
             lock_beat.reacquire()
 
             redis_connector = RedisConnector(tenant_id, cc_pair_id)
@@ -263,22 +278,58 @@ def check_for_indexing(self: Task, *, tenant_id: str | None) -> int | None:
                     db_session
                 )
                 for search_settings_instance in search_settings_list:
+                    if tenant_id in debug_tenants:
+                        ttl = redis_client.ttl(OnyxRedisLocks.CHECK_INDEXING_BEAT_LOCK)
+                        task_logger.info(
+                            f"check_for_indexing cc_pair search settings lock: "
+                            f"tenant={tenant_id} "
+                            f"cc_pair={cc_pair_id} "
+                            f"ttl={ttl}"
+                        )
+
                     redis_connector_index = redis_connector.new_index(
                         search_settings_instance.id
                     )
                     if redis_connector_index.fenced:
                         continue
 
+                    if tenant_id in debug_tenants:
+                        ttl = redis_client.ttl(OnyxRedisLocks.CHECK_INDEXING_BEAT_LOCK)
+                        task_logger.info(
+                            f"check_for_indexing get_connector_credential_pair_from_id: "
+                            f"tenant={tenant_id} "
+                            f"cc_pair={cc_pair_id} "
+                            f"ttl={ttl}"
+                        )
+
                     cc_pair = get_connector_credential_pair_from_id(
                         cc_pair_id, db_session
                     )
                     if not cc_pair:
                         continue
 
+                    if tenant_id in debug_tenants:
+                        ttl = redis_client.ttl(OnyxRedisLocks.CHECK_INDEXING_BEAT_LOCK)
+                        task_logger.info(
+                            f"check_for_indexing get_last_attempt_for_cc_pair: "
+                            f"tenant={tenant_id} "
+                            f"cc_pair={cc_pair_id} "
+                            f"ttl={ttl}"
+                        )
+
                     last_attempt = get_last_attempt_for_cc_pair(
                         cc_pair.id, search_settings_instance.id, db_session
                     )
 
+                    if tenant_id in debug_tenants:
+                        ttl = redis_client.ttl(OnyxRedisLocks.CHECK_INDEXING_BEAT_LOCK)
+                        task_logger.info(
+                            f"check_for_indexing cc_pair should index: "
+                            f"tenant={tenant_id} "
+                            f"cc_pair={cc_pair_id} "
+                            f"ttl={ttl}"
+                        )
+
                     search_settings_primary = False
                     if search_settings_instance.id == search_settings_list[0].id:
                         search_settings_primary = True
@@ -311,6 +362,15 @@ def check_for_indexing(self: Task, *, tenant_id: str | None) -> int | None:
                                 cc_pair.id, None, db_session
                             )
 
+                    if tenant_id in debug_tenants:
+                        ttl = redis_client.ttl(OnyxRedisLocks.CHECK_INDEXING_BEAT_LOCK)
+                        task_logger.info(
+                            f"check_for_indexing cc_pair try_creating_indexing_task: "
+                            f"tenant={tenant_id} "
+                            f"cc_pair={cc_pair_id} "
+                            f"ttl={ttl}"
+                        )
+
                     # using a task queue and only allowing one task per cc_pair/search_setting
                     # prevents us from starving out certain attempts
                     attempt_id = try_creating_indexing_task(
@@ -331,14 +391,42 @@ def check_for_indexing(self: Task, *, tenant_id: str | None) -> int | None:
                         )
                         tasks_created += 1
 
+                    if tenant_id in debug_tenants:
+                        ttl = redis_client.ttl(OnyxRedisLocks.CHECK_INDEXING_BEAT_LOCK)
+                        task_logger.info(
+                            f"check_for_indexing cc_pair try_creating_indexing_task finished: "
+                            f"tenant={tenant_id} "
+                            f"cc_pair={cc_pair_id} "
+                            f"ttl={ttl}"
+                        )
+
+        # debugging logic - remove after we're done
+        if tenant_id in debug_tenants:
+            ttl = redis_client.ttl(OnyxRedisLocks.CHECK_INDEXING_BEAT_LOCK)
+            task_logger.info(
+                f"check_for_indexing unfenced lock: "
+                f"tenant={tenant_id} "
+                f"ttl={ttl}"
+            )
+
+        lock_beat.reacquire()
+
         # Fail any index attempts in the DB that don't have fences
         # This shouldn't ever happen!
         with get_session_with_tenant(tenant_id) as db_session:
-            lock_beat.reacquire()
             unfenced_attempt_ids = get_unfenced_index_attempt_ids(
                 db_session, redis_client
             )
             for attempt_id in unfenced_attempt_ids:
+                # debugging logic - remove after we're done
+                if tenant_id in debug_tenants:
+                    ttl = redis_client.ttl(OnyxRedisLocks.CHECK_INDEXING_BEAT_LOCK)
+                    task_logger.info(
+                        f"check_for_indexing unfenced attempt id lock: "
+                        f"tenant={tenant_id} "
+                        f"ttl={ttl}"
+                    )
+
                 lock_beat.reacquire()
 
                 attempt = get_index_attempt(db_session, attempt_id)
@@ -356,9 +444,18 @@ def check_for_indexing(self: Task, *, tenant_id: str | None) -> int | None:
                     attempt.id, db_session, failure_reason=failure_reason
                 )
 
+        # debugging logic - remove after we're done
+        if tenant_id in debug_tenants:
+            ttl = redis_client.ttl(OnyxRedisLocks.CHECK_INDEXING_BEAT_LOCK)
+            task_logger.info(
+                f"check_for_indexing validate fences lock: "
+                f"tenant={tenant_id} "
+                f"ttl={ttl}"
+            )
+
+        lock_beat.reacquire()
         # we want to run this less frequently than the overall task
         if not redis_client.exists(OnyxRedisSignals.VALIDATE_INDEXING_FENCES):
-            lock_beat.reacquire()
             # clear any indexing fences that don't have associated celery tasks in progress
             # tasks can be in the queue in redis, in reserved tasks (prefetched by the worker),
             # or be currently executing
@@ -370,7 +467,6 @@ def check_for_indexing(self: Task, *, tenant_id: str | None) -> int | None:
                 task_logger.exception("Exception while validating indexing fences")
 
             redis_client.set(OnyxRedisSignals.VALIDATE_INDEXING_FENCES, 1, ex=60)
-
     except SoftTimeLimitExceeded:
         task_logger.info(
             "Soft time limit exceeded, task is being terminated gracefully."
@@ -405,7 +501,9 @@ def validate_indexing_fences(
     )
 
     # validate all existing indexing jobs
-    for key_bytes in r.scan_iter(RedisConnectorIndex.FENCE_PREFIX + "*"):
+    for key_bytes in r.scan_iter(
+        RedisConnectorIndex.FENCE_PREFIX + "*", count=SCAN_ITER_COUNT_DEFAULT
+    ):
         lock_beat.reacquire()
         with get_session_with_tenant(tenant_id) as db_session:
             validate_indexing_fence(

backend/onyx/background/celery/tasks/pruning/tasks.py

@@ -89,11 +89,11 @@ def check_for_pruning(self: Task, *, tenant_id: str | None) -> bool | None:
         timeout=CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT,
     )
 
-    try:
-        # these tasks should never overlap
-        if not lock_beat.acquire(blocking=False):
-            return None
+    # these tasks should never overlap
+    if not lock_beat.acquire(blocking=False):
+        return None
 
+    try:
         cc_pair_ids: list[int] = []
         with get_session_with_tenant(tenant_id) as db_session:
             cc_pairs = get_connector_credential_pairs(db_session)

backend/onyx/background/celery/tasks/shared/RetryDocumentIndex.py

@@ -28,13 +28,35 @@ class RetryDocumentIndex:
         wait=wait_random_exponential(multiplier=1, max=MAX_WAIT),
         stop=stop_after_delay(STOP_AFTER),
     )
-    def delete_single(self, doc_id: str) -> int:
-        return self.index.delete_single(doc_id)
+    def delete_single(
+        self,
+        doc_id: str,
+        *,
+        tenant_id: str | None,
+        chunk_count: int | None,
+    ) -> int:
+        return self.index.delete_single(
+            doc_id,
+            tenant_id=tenant_id,
+            chunk_count=chunk_count,
+        )
 
     @retry(
         retry=retry_if_exception_type(httpx.ReadTimeout),
         wait=wait_random_exponential(multiplier=1, max=MAX_WAIT),
         stop=stop_after_delay(STOP_AFTER),
     )
-    def update_single(self, doc_id: str, fields: VespaDocumentFields) -> int:
-        return self.index.update_single(doc_id, fields)
+    def update_single(
+        self,
+        doc_id: str,
+        *,
+        tenant_id: str | None,
+        chunk_count: int | None,
+        fields: VespaDocumentFields,
+    ) -> int:
+        return self.index.update_single(
+            doc_id,
+            tenant_id=tenant_id,
+            chunk_count=chunk_count,
+            fields=fields,
+        )