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 ###

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,15 @@ 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,
+        with_commit=True,
+    )

controller/project/manager.py

@@ -15,6 +15,8 @@
     data_slice,
     information_source,
     general,
+    attribute,
+    embedding
 )
 from submodules.model import daemon
 from fast_api.types import HuddleData, ProjectSize
@@ -25,6 +27,8 @@
 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 +57,47 @@ 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 = ["__ACCESS_GROUPS", "__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, "__ACCESS_GROUPS")
+    access_users_attribute = attribute.get_by_name(project_id, "__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, "__ACCESS_GROUPS")
+    access_users = attribute.get_by_name(project_id, "__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(