bug fix

* update

* fix

* finalize`

* remove unnecessary prints

* fix

* k

.github/workflows/pr-python-checks.yml

@@ -44,7 +44,7 @@ jobs:
     - name: Check import order with reorder-python-imports
       run: |
         cd backend
-        find ./danswer -name "*.py" | xargs reorder-python-imports --py311-plus
+        find ./onyx -name "*.py" | xargs reorder-python-imports --py311-plus
 
     - name: Check code formatting with Black
       run: |

.github/workflows/pr-python-connector-tests.yml

@@ -23,6 +23,10 @@ env:
   # Jira
   JIRA_USER_EMAIL: ${{ secrets.JIRA_USER_EMAIL }}
   JIRA_API_TOKEN: ${{ secrets.JIRA_API_TOKEN }}
+
+  GONG_ACCESS_KEY: ${{ secrets.GONG_ACCESS_KEY }}
+  GONG_ACCESS_KEY_SECRET: ${{ secrets.GONG_ACCESS_KEY_SECRET }}
+
   # Google
   GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON_STR: ${{ secrets.GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON_STR }}
   GOOGLE_DRIVE_OAUTH_CREDENTIALS_JSON_STR_TEST_USER_1: ${{ secrets.GOOGLE_DRIVE_OAUTH_CREDENTIALS_JSON_STR_TEST_USER_1 }}

.gitignore

@@ -1,12 +1,25 @@
-.env
+# editors
+.vscode
+.zed
+
+# macos
 .DS_store
+
+# python
 .venv
 .mypy_cache
 .idea
-/deployment/data/nginx/app.conf
-.vscode/
-*.sw?
-/backend/tests/regression/answer_quality/search_test_config.yaml
+
+# testing
 /web/test-results/
 backend/onyx/agent_search/main/test_data.json
 backend/tests/regression/answer_quality/test_data.json
+
+# secret files
+.env
+jira_test_env
+
+# others
+/deployment/data/nginx/app.conf
+*.sw?
+/backend/tests/regression/answer_quality/search_test_config.yaml

.pre-commit-config.yaml

@@ -1,12 +1,13 @@
 repos:
   - repo: https://github.com/psf/black
-    rev: 23.3.0
+    rev: 25.1.0
     hooks:
     - id: black
       language_version: python3.11
 
-  - repo: https://github.com/asottile/reorder_python_imports
-    rev: v3.9.0
+  # this is a fork which keeps compatibility with black
+  - repo: https://github.com/wimglenn/reorder-python-imports-black
+    rev: v3.14.0
     hooks:
     - id: reorder-python-imports
       args: ['--py311-plus', '--application-directories=backend/']
@@ -18,14 +19,14 @@ repos:
   # These settings will remove unused imports with side effects
   # Note: The repo currently does not and should not have imports with side effects
   - repo: https://github.com/PyCQA/autoflake
-    rev: v2.2.0
+    rev: v2.3.1
     hooks:
       - id: autoflake
         args: [ '--remove-all-unused-imports', '--remove-unused-variables', '--in-place' , '--recursive']
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.0.286
+    rev: v0.11.4
     hooks:
       - id: ruff
   - repo: https://github.com/pre-commit/mirrors-prettier

backend/Dockerfile.model_server

@@ -46,6 +46,7 @@ WORKDIR /app
 
 # Utils used by model server
 COPY ./onyx/utils/logger.py /app/onyx/utils/logger.py
+COPY ./onyx/utils/middleware.py /app/onyx/utils/middleware.py
 
 # Place to fetch version information
 COPY ./onyx/__init__.py /app/onyx/__init__.py

backend/alembic/versions/027381bce97c_add_shortcut_option_for_users.py

@@ -5,6 +5,7 @@ Revises: 6fc7886d665d
 Create Date: 2025-01-14 12:14:00.814390
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/08a1eda20fe1_add_earliest_indexing_to_connector.py

@@ -5,6 +5,7 @@ Revises: 8a87bd6ec550
 Create Date: 2024-07-23 11:12:39.462397
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/0a2b51deb0b8_add_starter_prompts.py

@@ -5,6 +5,7 @@ Revises: 5f4b8568a221
 Create Date: 2024-03-02 23:23:49.960309
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.dialects import postgresql

backend/alembic/versions/0a98909f2757_enable_encrypted_fields.py

@@ -5,6 +5,7 @@ Revises: 570282d33c49
 Create Date: 2024-05-05 19:30:34.317972
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.sql import table

backend/alembic/versions/0ebb1d516877_add_ccpair_deletion_failure_message.py

@@ -5,6 +5,7 @@ Revises: 52a219fb5233
 Create Date: 2024-09-10 15:03:48.233926
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/0f7ff6d75b57_add_index_to_index_attempt_time_created.py

@@ -5,6 +5,7 @@ Revises: 369644546676
 Create Date: 2025-01-10 14:01:14.067144
 
 """
+
 from alembic import op
 
 # revision identifiers, used by Alembic.

backend/alembic/versions/15326fcec57e_introduce_onyx_apis.py

@@ -5,6 +5,7 @@ Revises: 77d07dffae64
 Create Date: 2023-11-11 20:51:24.228999
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/173cae5bba26_port_config_store.py

@@ -5,6 +5,7 @@ Revises: e50154680a5c
 Create Date: 2024-03-19 15:30:44.425436
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.dialects import postgresql

backend/alembic/versions/177de57c21c9_display_custom_llm_models.py

@@ -5,6 +5,7 @@ Revises: 4ee1287bd26a
 Create Date: 2024-11-21 11:49:04.488677
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.dialects import postgresql

backend/alembic/versions/1a03d2c2856b_add_indexes_to_document__tag.py

@@ -5,6 +5,7 @@ Revises: 9c00a2bccb83
 Create Date: 2025-02-18 10:45:13.957807
 
 """
+
 from alembic import op
 
 # revision identifiers, used by Alembic.

backend/alembic/versions/1b10e1fda030_add_additional_data_to_notifications.py

@@ -5,6 +5,7 @@ Revises: 6756efa39ada
 Create Date: 2024-10-15 19:26:44.071259
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.dialects import postgresql

backend/alembic/versions/1b8206b29c5d_add_user_delete_cascades.py

@@ -5,6 +5,7 @@ Revises: 35e6853a51d5
 Create Date: 2024-09-18 11:48:59.418726
 
 """
+
 from alembic import op
 
 

backend/alembic/versions/213fd978c6d8_notifications.py

@@ -5,6 +5,7 @@ Revises: 5fc1f54cc252
 Create Date: 2024-08-10 11:13:36.070790
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/23957775e5f5_remove_feedback_foreignkey_constraint.py

@@ -5,6 +5,7 @@ Revises: bc9771dccadf
 Create Date: 2024-06-27 16:04:51.480437
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/2666d766cb9b_google_oauth2.py

@@ -5,6 +5,7 @@ Revises: 6d387b3196c2
 Create Date: 2023-05-05 15:49:35.716016
 
 """
+
 import fastapi_users_db_sqlalchemy
 import sqlalchemy as sa
 from alembic import op

backend/alembic/versions/26b931506ecb_default_chosen_assistants_to_none.py

@@ -5,6 +5,7 @@ Revises: 2daa494a0851
 Create Date: 2024-11-12 13:23:29.858995
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.dialects import postgresql

backend/alembic/versions/27c6ecc08586_permission_framework.py

@@ -5,6 +5,7 @@ Revises: 2666d766cb9b
 Create Date: 2023-05-24 18:45:17.244495
 
 """
+
 import fastapi_users_db_sqlalchemy
 import sqlalchemy as sa
 from alembic import op

backend/alembic/versions/2955778aa44c_add_chunk_count_to_document.py

@@ -5,6 +5,7 @@ Revises: c0aab6edb6dd
 Create Date: 2025-01-04 11:39:43.268612
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/2cdeff6d8c93_set_built_in_to_default.py

@@ -5,6 +5,7 @@ Revises: f5437cc136c5
 Create Date: 2025-02-11 14:57:51.308775
 
 """
+
 from alembic import op
 
 

backend/alembic/versions/2d2304e27d8c_add_above_below_to_persona.py

@@ -5,6 +5,7 @@ Revises: 4b08d97e175a
 Create Date: 2024-08-21 19:15:15.762948
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/2daa494a0851_add_group_sync_time.py

@@ -5,6 +5,7 @@ Revises: c0fd6e4da83a
 Create Date: 2024-11-11 10:57:22.991157
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/2f80c6a2550f_add_chat_session_specific_temperature_.py

@@ -5,6 +5,7 @@ Revises: 33ea50e88f24
 Create Date: 2025-01-31 10:30:27.289646
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/30c1d5744104_persona_datetime_aware.py

@@ -5,6 +5,7 @@ Revises: 7f99be1cb9f5
 Create Date: 2023-10-16 23:21:01.283424
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/325975216eb3_add_icon_color_and_icon_shape_to_persona.py

@@ -5,6 +5,7 @@ Revises: 91ffac7e65b3
 Create Date: 2024-07-24 21:29:31.784562
 
 """
+
 import random
 from alembic import op
 import sqlalchemy as sa

backend/alembic/versions/33cb72ea4d80_single_tool_call_per_message.py

@@ -5,6 +5,7 @@ Revises: 5b29123cd710
 Create Date: 2024-11-01 12:51:01.535003
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/33ea50e88f24_foreign_key_input_prompts.py

@@ -5,6 +5,7 @@ Revises: a6df6b88ef81
 Create Date: 2025-01-29 10:54:22.141765
 
 """
+
 from alembic import op
 
 

backend/alembic/versions/351faebd379d_add_curator_fields.py

@@ -5,6 +5,7 @@ Revises: ee3f4b47fad5
 Create Date: 2024-08-15 22:37:08.397052
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/35e518e0ddf4_properly_cascade.py

@@ -5,6 +5,7 @@ Revises: 91a0a4d62b14
 Create Date: 2024-09-20 21:24:04.891018
 
 """
+
 from alembic import op
 
 

backend/alembic/versions/35e6853a51d5_server_default_chosen_assistants.py

@@ -5,6 +5,7 @@ Revises: c99d76fcd298
 Create Date: 2024-09-13 13:20:32.885317
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.dialects import postgresql

backend/alembic/versions/369644546676_add_composite_index_for_index_attempt_.py

@@ -5,6 +5,7 @@ Revises: 2955778aa44c
 Create Date: 2025-01-08 15:38:17.224380
 
 """
+
 from alembic import op
 from sqlalchemy import text
 

backend/alembic/versions/3781a5eb12cb_add_chunk_stats_table.py

@@ -5,6 +5,7 @@ Revises: df46c75b714e
 Create Date: 2025-03-10 10:02:30.586666
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/3879338f8ba1_add_tool_table.py

@@ -5,6 +5,7 @@ Revises: f1c6478c3fd8
 Create Date: 2024-05-11 16:11:23.718084
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/38eda64af7fe_add_chat_session_sharing.py

@@ -5,6 +5,7 @@ Revises: 776b3bbe9092
 Create Date: 2024-03-27 19:41:29.073594
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/3934b1bc7b62_update_github_connector_repo_name_to_.py

@@ -5,6 +5,7 @@ Revises: b7c2b63c4a03
 Create Date: 2025-03-05 10:50:30.516962
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 import json

backend/alembic/versions/3b25685ff73c_move_is_public_to_cc_pair.py

@@ -5,6 +5,7 @@ Revises: e0a68a81d434
 Create Date: 2023-10-05 18:47:09.582849
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/3bd4c84fe72f_improved_index.py

@@ -5,6 +5,7 @@ Revises: 8f43500ee275
 Create Date: 2025-02-26 13:07:56.217791
 
 """
+
 from alembic import op
 
 

backend/alembic/versions/3c5e35aa9af0_polling_document_count.py

@@ -5,6 +5,7 @@ Revises: 27c6ecc08586
 Create Date: 2023-06-14 23:45:51.760440
 
 """
+
 import sqlalchemy as sa
 from alembic import op
 

backend/alembic/versions/3c6531f32351_add_back_input_prompts.py

@@ -5,6 +5,7 @@ Revises: aeda5f2df4f6
 Create Date: 2025-01-13 12:49:51.705235
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 import fastapi_users_db_sqlalchemy

backend/alembic/versions/401c1ac29467_add_tables_for_ui_based_llm_.py

@@ -5,6 +5,7 @@ Revises: 703313b75876
 Create Date: 2024-04-13 18:07:29.153817
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.dialects import postgresql

backend/alembic/versions/43cbbb3f5e6a_rename_index_origin_to_index_recursively.py

@@ -5,6 +5,7 @@ Revises: e1392f05e840
 Create Date: 2024-08-01 12:38:54.466081
 
 """
+
 from alembic import op
 
 # revision identifiers, used by Alembic.

backend/alembic/versions/44f856ae2a4a_add_cloud_embedding_model.py

@@ -5,6 +5,7 @@ Revises: d716b0791ddd
 Create Date: 2024-06-28 20:01:05.927647
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/4505fd7302e1_added_is_internet_to_dbdoc.py

@@ -5,6 +5,7 @@ Revises: c18cdf4b497e
 Create Date: 2024-06-18 20:46:09.095034
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/465f78d9b7f9_larger_access_tokens_for_oauth.py

@@ -5,6 +5,7 @@ Revises: 3c5e35aa9af0
 Create Date: 2023-07-18 17:33:40.365034
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/46625e4745d4_remove_native_enum.py

@@ -5,6 +5,7 @@ Revises: 9d97fecfab7f
 Create Date: 2023-10-27 11:38:33.803145
 
 """
+
 from alembic import op
 from sqlalchemy import String
 

backend/alembic/versions/46b7a812670f_fix_user__external_user_group_id_fk.py

@@ -5,6 +5,7 @@ Revises: f32615f71aeb
 Create Date: 2024-09-23 12:58:03.894038
 
 """
+
 from alembic import op
 
 # revision identifiers, used by Alembic.

backend/alembic/versions/4738e4b3bae1_pg_file_store.py

@@ -5,6 +5,7 @@ Revises: e91df4e935ef
 Create Date: 2024-03-20 18:53:32.461518
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/47433d30de82_create_indexattempt_table.py

@@ -5,6 +5,7 @@ Revises:
 Create Date: 2023-05-04 00:55:32.971991
 
 """
+
 import sqlalchemy as sa
 from alembic import op
 from sqlalchemy.dialects import postgresql

backend/alembic/versions/475fcefe8826_add_name_to_api_key.py

@@ -5,6 +5,7 @@ Revises: ecab2b3f1a3b
 Create Date: 2024-04-11 11:05:18.414438
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/4794bc13e484_update_prompt_length.py

@@ -0,0 +1,51 @@
+"""update prompt length
+
+Revision ID: 4794bc13e484
+Revises: f7505c5b0284
+Create Date: 2025-04-02 11:26:36.180328
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "4794bc13e484"
+down_revision = "f7505c5b0284"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.alter_column(
+        "prompt",
+        "system_prompt",
+        existing_type=sa.TEXT(),
+        type_=sa.String(length=5000000),
+        existing_nullable=False,
+    )
+    op.alter_column(
+        "prompt",
+        "task_prompt",
+        existing_type=sa.TEXT(),
+        type_=sa.String(length=5000000),
+        existing_nullable=False,
+    )
+
+
+def downgrade() -> None:
+    op.alter_column(
+        "prompt",
+        "system_prompt",
+        existing_type=sa.String(length=5000000),
+        type_=sa.TEXT(),
+        existing_nullable=False,
+    )
+    op.alter_column(
+        "prompt",
+        "task_prompt",
+        existing_type=sa.String(length=5000000),
+        type_=sa.TEXT(),
+        existing_nullable=False,
+    )

backend/alembic/versions/47e5bef3a1d7_add_persona_categories.py

@@ -5,6 +5,7 @@ Revises: dfbe9e93d3c7
 Create Date: 2024-11-05 18:55:02.221064
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/48d14957fe80_add_support_for_custom_tools.py

@@ -5,6 +5,7 @@ Revises: b85f02ec1308
 Create Date: 2024-06-09 14:58:19.946509
 
 """
+
 from alembic import op
 import fastapi_users_db_sqlalchemy
 import sqlalchemy as sa

backend/alembic/versions/495cb26ce93e_create_knowlege_graph_tables.py

@@ -0,0 +1,219 @@
+"""create knowlege graph tables
+
+Revision ID: 495cb26ce93e
+Revises: 6a804aeb4830
+Create Date: 2025-03-19 08:51:14.341989
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+
+# revision identifiers, used by Alembic.
+revision = "495cb26ce93e"
+down_revision = "6a804aeb4830"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "kg_entity_type",
+        sa.Column("id_name", sa.String(), primary_key=True, nullable=False, index=True),
+        sa.Column("description", sa.String(), nullable=True),
+        sa.Column("grounding", sa.String(), nullable=False),
+        sa.Column("clustering", postgresql.JSONB, nullable=False, server_default="{}"),
+        sa.Column(
+            "classification_requirements",
+            postgresql.JSONB,
+            nullable=False,
+            server_default="{}",
+        ),
+        sa.Column("cluster_count", sa.Integer(), nullable=True),
+        sa.Column(
+            "extraction_sources", postgresql.JSONB, nullable=False, server_default="{}"
+        ),
+        sa.Column("active", sa.Boolean(), nullable=False, default=False),
+        sa.Column(
+            "time_updated",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("now()"),
+            onupdate=sa.text("now()"),
+        ),
+        sa.Column(
+            "time_created", sa.DateTime(timezone=True), server_default=sa.text("now()")
+        ),
+        sa.Column("grounded_source_name", sa.String(), nullable=True, unique=True),
+        sa.Column(
+            "ge_determine_instructions", postgresql.ARRAY(sa.String()), nullable=True
+        ),
+        sa.Column("ge_grounding_signature", sa.String(), nullable=True),
+    )
+
+    # Create KGRelationshipType table
+    op.create_table(
+        "kg_relationship_type",
+        sa.Column("id_name", sa.String(), primary_key=True, nullable=False, index=True),
+        sa.Column("name", sa.String(), nullable=False, index=True),
+        sa.Column(
+            "source_entity_type_id_name", sa.String(), nullable=False, index=True
+        ),
+        sa.Column(
+            "target_entity_type_id_name", sa.String(), nullable=False, index=True
+        ),
+        sa.Column("definition", sa.Boolean(), nullable=False, default=False),
+        sa.Column("clustering", postgresql.JSONB, nullable=False, server_default="{}"),
+        sa.Column("cluster_count", sa.Integer(), nullable=True),
+        sa.Column("type", sa.String(), nullable=False, index=True),
+        sa.Column("active", sa.Boolean(), nullable=False, default=True),
+        sa.Column(
+            "time_updated",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("now()"),
+            onupdate=sa.text("now()"),
+        ),
+        sa.Column(
+            "time_created", sa.DateTime(timezone=True), server_default=sa.text("now()")
+        ),
+        sa.ForeignKeyConstraint(
+            ["source_entity_type_id_name"], ["kg_entity_type.id_name"]
+        ),
+        sa.ForeignKeyConstraint(
+            ["target_entity_type_id_name"], ["kg_entity_type.id_name"]
+        ),
+    )
+
+    # Create KGEntity table
+    op.create_table(
+        "kg_entity",
+        sa.Column("id_name", sa.String(), primary_key=True, nullable=False, index=True),
+        sa.Column("name", sa.String(), nullable=False, index=True),
+        sa.Column("document_id", sa.String(), nullable=True, index=True),
+        sa.Column(
+            "alternative_names",
+            postgresql.ARRAY(sa.String()),
+            nullable=False,
+            server_default="{}",
+        ),
+        sa.Column("entity_type_id_name", sa.String(), nullable=False, index=True),
+        sa.Column("description", sa.String(), nullable=True),
+        sa.Column(
+            "keywords",
+            postgresql.ARRAY(sa.String()),
+            nullable=False,
+            server_default="{}",
+        ),
+        sa.Column("cluster_count", sa.Integer(), nullable=True),
+        sa.Column(
+            "acl", postgresql.ARRAY(sa.String()), nullable=False, server_default="{}"
+        ),
+        sa.Column("boosts", postgresql.JSONB, nullable=False, server_default="{}"),
+        sa.Column("event_time", sa.DateTime(timezone=True), nullable=True),
+        sa.Column(
+            "time_updated",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("now()"),
+            onupdate=sa.text("now()"),
+        ),
+        sa.Column(
+            "time_created", sa.DateTime(timezone=True), server_default=sa.text("now()")
+        ),
+        sa.ForeignKeyConstraint(["entity_type_id_name"], ["kg_entity_type.id_name"]),
+    )
+    op.create_index("ix_entity_type_acl", "kg_entity", ["entity_type_id_name", "acl"])
+    op.create_index(
+        "ix_entity_name_search", "kg_entity", ["name", "entity_type_id_name"]
+    )
+
+    # Create KGRelationship table
+    op.create_table(
+        "kg_relationship",
+        sa.Column("id_name", sa.String(), primary_key=True, nullable=False, index=True),
+        sa.Column("source_node", sa.String(), nullable=False, index=True),
+        sa.Column("target_node", sa.String(), nullable=False, index=True),
+        sa.Column("type", sa.String(), nullable=False, index=True),
+        sa.Column("relationship_type_id_name", sa.String(), nullable=False, index=True),
+        sa.Column("cluster_count", sa.Integer(), nullable=True),
+        sa.Column(
+            "time_updated",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("now()"),
+            onupdate=sa.text("now()"),
+        ),
+        sa.Column(
+            "time_created", sa.DateTime(timezone=True), server_default=sa.text("now()")
+        ),
+        sa.ForeignKeyConstraint(["source_node"], ["kg_entity.id_name"]),
+        sa.ForeignKeyConstraint(["target_node"], ["kg_entity.id_name"]),
+        sa.ForeignKeyConstraint(
+            ["relationship_type_id_name"], ["kg_relationship_type.id_name"]
+        ),
+        sa.UniqueConstraint(
+            "source_node",
+            "target_node",
+            "type",
+            name="uq_kg_relationship_source_target_type",
+        ),
+    )
+    op.create_index(
+        "ix_kg_relationship_nodes", "kg_relationship", ["source_node", "target_node"]
+    )
+
+    # Create KGTerm table
+    op.create_table(
+        "kg_term",
+        sa.Column("id_term", sa.String(), primary_key=True, nullable=False, index=True),
+        sa.Column(
+            "entity_types",
+            postgresql.ARRAY(sa.String()),
+            nullable=False,
+            server_default="{}",
+        ),
+        sa.Column(
+            "time_updated",
+            sa.DateTime(timezone=True),
+            server_default=sa.text("now()"),
+            onupdate=sa.text("now()"),
+        ),
+        sa.Column(
+            "time_created", sa.DateTime(timezone=True), server_default=sa.text("now()")
+        ),
+    )
+    op.create_index("ix_search_term_entities", "kg_term", ["entity_types"])
+    op.create_index("ix_search_term_term", "kg_term", ["id_term"])
+    op.add_column(
+        "document",
+        sa.Column("kg_processed", sa.Boolean(), nullable=False, server_default="false"),
+    )
+    op.add_column(
+        "document",
+        sa.Column("kg_data", postgresql.JSONB(), nullable=False, server_default="{}"),
+    )
+    op.add_column(
+        "connector",
+        sa.Column(
+            "kg_extraction_enabled",
+            sa.Boolean(),
+            nullable=False,
+            server_default="false",
+        ),
+    )
+
+    op.add_column(
+        "document_by_connector_credential_pair",
+        sa.Column("has_been_kg_processed", sa.Boolean(), nullable=True),
+    )
+
+
+def downgrade() -> None:
+    # Drop tables in reverse order of creation to handle dependencies
+    op.drop_table("kg_term")
+    op.drop_table("kg_relationship")
+    op.drop_table("kg_entity")
+    op.drop_table("kg_relationship_type")
+    op.drop_table("kg_entity_type")
+    op.drop_column("connector", "kg_extraction_enabled")
+    op.drop_column("document_by_connector_credential_pair", "has_been_kg_processed")
+    op.drop_column("document", "kg_data")
+    op.drop_column("document", "kg_processed")

backend/alembic/versions/4a951134c801_moved_status_to_connector_credential_.py

@@ -5,6 +5,7 @@ Revises: 7477a5f5d728
 Create Date: 2024-08-10 19:20:34.527559
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/4b08d97e175a_change_default_prune_freq.py

@@ -5,6 +5,7 @@ Revises: d9ec13955951
 Create Date: 2024-08-20 15:28:52.993827
 
 """
+
 from alembic import op
 
 # revision identifiers, used by Alembic.

backend/alembic/versions/4d58345da04a_lowercase_user_emails.py

@@ -5,6 +5,7 @@ Revises: f1ca58b2f2ec
 Create Date: 2025-01-29 07:48:46.784041
 
 """
+
 import logging
 from typing import cast
 from alembic import op

backend/alembic/versions/4ee1287bd26a_add_multiple_slack_bot_support.py

@@ -5,6 +5,7 @@ Revises: 47e5bef3a1d7
 Create Date: 2024-11-06 13:15:53.302644
 
 """
+
 from typing import cast
 from alembic import op
 import sqlalchemy as sa

backend/alembic/versions/50b683a8295c_add_additional_retrieval_controls_to_.py

@@ -5,6 +5,7 @@ Revises: 7da0ae5ad583
 Create Date: 2023-11-27 17:23:29.668422
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/52a219fb5233_add_last_synced_and_last_modified_to_document_table.py

@@ -5,6 +5,7 @@ Revises: f7e58d357687
 Create Date: 2024-08-28 17:40:46.077470
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.sql import func

backend/alembic/versions/54a74a0417fc_danswerbot_onyxbot.py

@@ -5,6 +5,7 @@ Revises: 94dc3d0236f8
 Create Date: 2024-12-11 18:05:05.490737
 
 """
+
 from alembic import op
 
 

backend/alembic/versions/55546a7967ee_assistant_rework.py

@@ -5,6 +5,7 @@ Revises: 61ff3651add4
 Create Date: 2024-09-18 17:00:23.755399
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.dialects import postgresql

backend/alembic/versions/570282d33c49_track_onyxbot_explicitly.py

@@ -5,6 +5,7 @@ Revises: 7547d982db8f
 Create Date: 2024-05-04 17:49:28.568109
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/57b53544726e_add_document_set_tables.py

@@ -5,6 +5,7 @@ Revises: 800f48024ae9
 Create Date: 2023-09-20 16:59:39.097177
 
 """
+
 from alembic import op
 import fastapi_users_db_sqlalchemy
 import sqlalchemy as sa

backend/alembic/versions/5809c0787398_add_chat_sessions.py

@@ -5,6 +5,7 @@ Revises: d929f0c1c6af
 Create Date: 2023-09-04 15:29:44.002164
 
 """
+
 import fastapi_users_db_sqlalchemy
 from alembic import op
 import sqlalchemy as sa

backend/alembic/versions/5b29123cd710_nullable_search_settings_for_historic_.py

@@ -5,6 +5,7 @@ Revises: 949b4a92a401
 Create Date: 2024-10-30 19:37:59.630704
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/5c7fdadae813_match_any_keywords_flag_for_standard_.py

@@ -5,6 +5,7 @@ Revises: efb35676026c
 Create Date: 2024-09-13 18:52:59.256478
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/5d12a446f5c0_add_api_version_and_deployment_name_to_.py

@@ -5,6 +5,7 @@ Revises: e4334d5b33ba
 Create Date: 2024-10-08 15:56:07.975636
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/5e84129c8be3_add_docs_indexed_column_to_index_.py

@@ -5,6 +5,7 @@ Revises: e6a4bbc13fe4
 Create Date: 2023-08-10 21:43:09.069523
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/5f4b8568a221_add_removed_documents_to_index_attempt.py

@@ -5,6 +5,7 @@ Revises: dbaa756c2ccf
 Create Date: 2024-02-16 15:02:03.319907
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/5fc1f54cc252_hybrid_enum.py

@@ -5,6 +5,7 @@ Revises: 1d6ad76d1f37
 Create Date: 2024-08-06 15:35:40.278485
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/61ff3651add4_add_permission_syncing.py

@@ -5,6 +5,7 @@ Revises: 1b8206b29c5d
 Create Date: 2024-09-05 13:57:11.770413
 
 """
+
 import fastapi_users_db_sqlalchemy
 
 from alembic import op

backend/alembic/versions/643a84a42a33_add_user_configured_names_to_llmprovider.py

@@ -5,6 +5,7 @@ Revises: 0a98909f2757
 Create Date: 2024-05-07 14:54:55.493100
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/6756efa39ada_id_uuid_for_chat_session.py

@@ -5,6 +5,7 @@ Revises: 5d12a446f5c0
 Create Date: 2024-10-15 17:47:44.108537
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/6a804aeb4830_duplicated_no_harm_user_file_migration.py

@@ -5,6 +5,7 @@ Revises: 8e1ac4f39a9f
 Create Date: 2025-04-01 07:26:10.539362
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy import inspect

backend/alembic/versions/6d387b3196c2_basic_auth.py

@@ -5,6 +5,7 @@ Revises: 47433d30de82
 Create Date: 2023-05-05 14:40:10.242502
 
 """
+
 import fastapi_users_db_sqlalchemy
 import sqlalchemy as sa
 from alembic import op

backend/alembic/versions/6d562f86c78b_remove_default_bot.py

@@ -5,6 +5,7 @@ Revises: 177de57c21c9
 Create Date: 2024-11-22 11:51:29.331336
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/6fc7886d665d_make_categories_labels_and_many_to_many.py

@@ -5,6 +5,7 @@ Revises: 3c6531f32351
 Create Date: 2025-01-13 18:12:18.029112
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/703313b75876_add_tokenratelimit_tables.py

@@ -5,6 +5,7 @@ Revises: fad14119fb92
 Create Date: 2024-04-15 01:36:02.952809
 
 """
+
 import json
 from typing import cast
 from alembic import op

backend/alembic/versions/70f00c45c0f2_more_descriptive_filestore.py

@@ -5,6 +5,7 @@ Revises: 3879338f8ba1
 Create Date: 2024-05-17 17:51:41.926893
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.dialects import postgresql

backend/alembic/versions/72bdc9929a46_permission_auto_sync_framework.py

@@ -5,6 +5,7 @@ Revises: 475fcefe8826
 Create Date: 2024-04-14 21:15:28.659634
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/7547d982db8f_chat_folders.py

@@ -5,6 +5,7 @@ Revises: ef7da92f7213
 Create Date: 2024-05-02 15:18:56.573347
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 import fastapi_users_db_sqlalchemy

backend/alembic/versions/767f1c2a00eb_count_chat_tokens.py

@@ -5,6 +5,7 @@ Revises: dba7f71618f5
 Create Date: 2023-09-21 10:03:21.509899
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/76b60d407dfb_cc_pair_name_not_unique.py

@@ -5,6 +5,7 @@ Revises: b156fa702355
 Create Date: 2023-12-22 21:42:10.018804
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/776b3bbe9092_remove_remaining_enums.py

@@ -5,6 +5,7 @@ Revises: 4738e4b3bae1
 Create Date: 2024-03-22 21:34:27.629444
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/77d07dffae64_forcibly_remove_more_enum_types_from_.py

@@ -5,6 +5,7 @@ Revises: d61e513bef0a
 Create Date: 2023-11-01 12:33:01.999617
 
 """
+
 from alembic import op
 from sqlalchemy import String
 

backend/alembic/versions/78dbe7e38469_task_tracking.py

@@ -5,6 +5,7 @@ Revises: 7ccea01261f6
 Create Date: 2023-10-15 23:40:50.593262
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/795b20b85b4b_add_llm_group_permissions_control.py

@@ -5,6 +5,7 @@ Revises: 05c07bf07c00
 Create Date: 2024-07-19 11:54:35.701558
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/797089dfb4d2_persona_start_date.py

@@ -5,6 +5,7 @@ Revises: 55546a7967ee
 Create Date: 2024-09-11 14:51:49.785835
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/79acd316403a_add_api_key_table.py

@@ -5,6 +5,7 @@ Revises: 904e5138fffb
 Create Date: 2024-01-11 17:56:37.934381
 
 """
+
 from alembic import op
 import fastapi_users_db_sqlalchemy
 import sqlalchemy as sa

backend/alembic/versions/7aea705850d5_added_slack_auto_filter.py

@@ -5,6 +5,7 @@ Revises: 4505fd7302e1
 Create Date: 2024-07-10 11:01:23.581015
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/7ccea01261f6_store_chat_retrieval_docs.py

@@ -5,6 +5,7 @@ Revises: a570b80a5f20
 Create Date: 2023-10-15 10:39:23.317453
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.dialects import postgresql

backend/alembic/versions/7da0ae5ad583_add_description_to_persona.py

@@ -5,6 +5,7 @@ Revises: e86866a9c78a
 Create Date: 2023-11-27 00:16:19.959414
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/7da543f5672f_add_slackbotconfig_table.py

@@ -5,6 +5,7 @@ Revises: febe9eaa0644
 Create Date: 2023-09-24 16:34:17.526128
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.dialects import postgresql

backend/alembic/versions/7f726bad5367_slack_followup.py

@@ -5,6 +5,7 @@ Revises: 79acd316403a
 Create Date: 2024-01-15 00:19:55.991224
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 

backend/alembic/versions/7f99be1cb9f5_add_index_for_getting_documents_just_by_.py

@@ -5,6 +5,7 @@ Revises: 78dbe7e38469
 Create Date: 2023-10-15 22:48:15.487762
 
 """
+
 from alembic import op
 
 

backend/alembic/versions/800f48024ae9_add_id_to_connectorcredentialpair.py

@@ -5,6 +5,7 @@ Revises: 767f1c2a00eb
 Create Date: 2023-09-19 16:13:42.299715
 
 """
+
 from alembic import op
 import sqlalchemy as sa
 from sqlalchemy.schema import Sequence, CreateSequence