Cognition integration provider

alembic/versions/6868ac66ea92_adds_cognition_group_management.py

@@ -0,0 +1,107 @@
+"""adds cognition group management
+
+Revision ID: 6868ac66ea92
+Revises: 36f087da55b1
+Create Date: 2025-06-26 12:58:16.408919
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "6868ac66ea92"
+down_revision = "36f087da55b1"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "group",
+        sa.Column("id", postgresql.UUID(as_uuid=True), nullable=False),
+        sa.Column("organization_id", postgresql.UUID(as_uuid=True), nullable=True),
+        sa.Column("name", sa.String(), nullable=True),
+        sa.Column("description", sa.String(), nullable=True),
+        sa.Column("created_at", sa.DateTime(), nullable=True),
+        sa.Column("created_by", postgresql.UUID(as_uuid=True), nullable=True),
+        sa.Column("meta_data", sa.JSON(), nullable=True),
+        sa.ForeignKeyConstraint(["created_by"], ["user.id"], ondelete="SET NULL"),
+        sa.ForeignKeyConstraint(
+            ["organization_id"], ["organization.id"], ondelete="CASCADE"
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("name"),
+        schema="cognition",
+    )
+    op.create_index(
+        op.f("ix_cognition_group_created_by"),
+        "group",
+        ["created_by"],
+        unique=False,
+        schema="cognition",
+    )
+    op.create_index(
+        op.f("ix_cognition_group_organization_id"),
+        "group",
+        ["organization_id"],
+        unique=False,
+        schema="cognition",
+    )
+    op.create_table(
+        "group_member",
+        sa.Column("id", postgresql.UUID(as_uuid=True), nullable=False),
+        sa.Column("group_id", postgresql.UUID(as_uuid=True), nullable=True),
+        sa.Column("user_id", postgresql.UUID(as_uuid=True), nullable=True),
+        sa.Column("created_at", sa.DateTime(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["group_id"], ["cognition.group.id"], ondelete="CASCADE"
+        ),
+        sa.ForeignKeyConstraint(["user_id"], ["user.id"], ondelete="CASCADE"),
+        sa.PrimaryKeyConstraint("id"),
+        schema="cognition",
+    )
+    op.create_index(
+        op.f("ix_cognition_group_member_group_id"),
+        "group_member",
+        ["group_id"],
+        unique=False,
+        schema="cognition",
+    )
+    op.create_index(
+        op.f("ix_cognition_group_member_user_id"),
+        "group_member",
+        ["user_id"],
+        unique=False,
+        schema="cognition",
+    )
+    op.add_column("user", sa.Column("oidc_identifier", sa.String(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("user", "oidc_identifier")
+    op.drop_index(
+        op.f("ix_cognition_group_member_user_id"),
+        table_name="group_member",
+        schema="cognition",
+    )
+    op.drop_index(
+        op.f("ix_cognition_group_member_group_id"),
+        table_name="group_member",
+        schema="cognition",
+    )
+    op.drop_table("group_member", schema="cognition")
+    op.drop_index(
+        op.f("ix_cognition_group_organization_id"),
+        table_name="group",
+        schema="cognition",
+    )
+    op.drop_index(
+        op.f("ix_cognition_group_created_by"), table_name="group", schema="cognition"
+    )
+    op.drop_table("group", schema="cognition")
+    # ### end Alembic commands ###

api/transfer.py

@@ -196,7 +196,7 @@ def __calculate_missing_attributes(project_id: str, user_id: str) -> None:
             i += 1
             if i >= 60:
                 i = 0
-                daemon.reset_session_token_in_thread()
+                general.remove_and_refresh_session(request_new=True)
             if tokenization.is_doc_bin_creation_running_or_queued(project_id):
                 time.sleep(2)
                 continue
@@ -211,7 +211,7 @@ def __calculate_missing_attributes(project_id: str, user_id: str) -> None:
                 break
             if i >= 60:
                 i = 0
-                daemon.reset_session_token_in_thread()
+                general.remove_and_refresh_session(request_new=True)
 
             current_att_id = attribute_ids[0]
             current_att = attribute.get(project_id, current_att_id)

app.py

@@ -16,6 +16,7 @@
 from fast_api.routes.organization import router as org_router
 from fast_api.routes.project import router as project_router
 from fast_api.routes.project_setting import router as project_setting_router
+from fast_api.routes.project_internal import router as project_internal_router
 from fast_api.routes.misc import router as misc_router
 from fast_api.routes.comment import router as comment_router
 from fast_api.routes.attribute import router as attribute_router
@@ -43,6 +44,7 @@
     PREFIX_ORGANIZATION,
     PREFIX_PROJECT,
     PREFIX_PROJECT_SETTING,
+    PREFIX_PROJECT_INTERNAL,
     PREFIX_MISC,
     PREFIX_COMMENT,
     PREFIX_ATTRIBUTE,
@@ -121,6 +123,9 @@
 fastapi_app_internal.include_router(
     record_internal_router, prefix=PREFIX_RECORD_INTERNAL, tags=["record-internal"]
 )
+fastapi_app_internal.include_router(
+    project_internal_router, prefix=PREFIX_PROJECT_INTERNAL, tags=["project-internal"]
+)
 
 
 routes = [

controller/auth/kratos.py

@@ -263,3 +263,15 @@ def check_user_exists(email: str) -> bool:
             if i["traits"]["email"].lower() == email.lower():
                 return True
     return False
+
+
+def get_user_from_search(email: str) -> bool:
+    request = requests.get(
+        f"{KRATOS_ADMIN_URL}/identities?preview_credentials_identifier_similar={quote(email)}"
+    )
+    if request.ok:
+        identities = request.json()
+        for i in identities:
+            if i["traits"]["email"].lower() == email.lower():
+                return i
+    return None

controller/monitor/manager.py

@@ -115,3 +115,14 @@ def cancel_parse_cognition_file_task(
         transformation_key,
         with_commit=True,
     )
+
+
+def cancel_integration_task(
+    task_info: Dict[str, Any],
+) -> None:
+
+    integration_id = task_info.get("integrationId")
+
+    task_monitor.set_integration_task_to_failed(
+        integration_id, error_message="Cancelled by task manager"
+    )

controller/project/manager.py

@@ -15,16 +15,24 @@
     data_slice,
     information_source,
     general,
+    attribute,
+    embedding,
 )
 from submodules.model import daemon
 from fast_api.types import HuddleData, ProjectSize
 from controller.task_master import manager as task_master_manager
 from submodules.model.enums import TaskType, RecordTokenizationScope
 from submodules.model.business_objects import util as db_util
+from submodules.model.integration_objects.helper import (
+    REFINERY_ATTRIBUTE_ACCESS_GROUPS,
+    REFINERY_ATTRIBUTE_ACCESS_USERS,
+)
 from submodules.s3 import controller as s3
 from service.search import search
 from controller.auth import kratos
 from submodules.model.util import sql_alchemy_to_dict
+from controller.embedding import connector
+
 
 ALL_PROJECTS_WHITELIST = {
     "id",
@@ -53,6 +61,77 @@ def get_all_projects(organization_id: str) -> List[Project]:
     return project.get_all(organization_id)
 
 
+def get_all_projects_with_access_management(organization_id: str) -> List[Project]:
+    return project.get_all_with_access_management(organization_id)
+
+
+def activate_access_management(project_id):
+    relative_position = attribute.get_relative_position(project_id)
+    if relative_position is None:
+        relative_position = 1
+    else:
+        relative_position += 1
+    filter_attributes = [
+        REFINERY_ATTRIBUTE_ACCESS_GROUPS,
+        REFINERY_ATTRIBUTE_ACCESS_USERS,
+    ]
+    attribute.create(
+        project_id=project_id,
+        relative_position=relative_position,
+        name=filter_attributes[0],
+        data_type=enums.DataTypes.PERMISSION.value,
+        user_created=False,
+        visibility=enums.AttributeVisibility.HIDE.value,
+        with_commit=True,
+        state=enums.AttributeState.AUTOMATICALLY_CREATED.value,
+    )
+    attribute.create(
+        project_id=project_id,
+        relative_position=relative_position + 1,
+        name=filter_attributes[1],
+        data_type=enums.DataTypes.PERMISSION.value,
+        user_created=False,
+        visibility=enums.AttributeVisibility.HIDE.value,
+        with_commit=True,
+        state=enums.AttributeState.AUTOMATICALLY_CREATED.value,
+    )
+    all_embeddings = embedding.get_all_embeddings_by_project_id(project_id)
+    for embedding_item in all_embeddings:
+        prev_filter_attributes = embedding_item.filter_attributes or []
+        new_filter_attributes = list(set(prev_filter_attributes + filter_attributes))
+        embedding_item.filter_attributes = new_filter_attributes
+        general.commit()
+        if connector.update_attribute_payloads_for_neural_search(
+            project_id, str(embedding_item.id)
+        ):
+            embedding.update_embedding_filter_attributes(
+                project_id,
+                str(embedding_item.id),
+                new_filter_attributes,
+                with_commit=True,
+            )
+
+
+def deactivate_access_management(project_id: str) -> None:
+    record.delete_access_management_attributes(project_id)
+    access_groups_attribute = attribute.get_by_name(
+        project_id, REFINERY_ATTRIBUTE_ACCESS_GROUPS
+    )
+    access_users_attribute = attribute.get_by_name(
+        project_id, REFINERY_ATTRIBUTE_ACCESS_USERS
+    )
+    if access_groups_attribute:
+        attribute.delete(project_id, access_groups_attribute.id, with_commit=True)
+    if access_users_attribute:
+        attribute.delete(project_id, access_users_attribute.id, with_commit=True)
+
+
+def is_access_management_activated(project_id: str) -> bool:
+    access_groups = attribute.get_by_name(project_id, REFINERY_ATTRIBUTE_ACCESS_GROUPS)
+    access_users = attribute.get_by_name(project_id, REFINERY_ATTRIBUTE_ACCESS_USERS)
+    return access_groups is not None and access_users is not None
+
+
 def get_all_projects_by_user(organization_id) -> List[Project]:
     projects = project.get_all_by_user_organization_id(organization_id)
     project_dicts = sql_alchemy_to_dict(