Cognition integration provider

business_objects/embedding.py

@@ -321,8 +321,11 @@ def __build_payload_selector(
             if (
                 data_type != enums.DataTypes.TEXT.value
                 and data_type != enums.DataTypes.LLM_RESPONSE.value
+                and data_type != enums.DataTypes.PERMISSION.value
             ):
                 payload_selector += f"'{attr}', (r.\"data\"->>'{attr}')::{data_type}"
+            if data_type == enums.DataTypes.PERMISSION.value:  
+                payload_selector += f"'{attr}', r.\"data\"->'{attr}'"
             else:
                 payload_selector += f"'{attr}', r.\"data\"->>'{attr}'"
         payload_selector = f"json_build_object({payload_selector}) payload"

business_objects/monitor.py

@@ -9,6 +9,7 @@
     markdown_file as markdown_file_db_bo,
     file_extraction as file_extraction_db_bo,
     file_transformation as file_transformation_db_bo,
+    integration as integration_db_bo,
 )
 
 FILE_CACHING_IN_PROGRESS_STATES = [
@@ -197,6 +198,16 @@ def set_parse_cognition_file_task_to_failed(
     general.commit()
 
 
+def set_integration_task_to_failed(
+    integration_id: str,
+    with_commit: bool = False,
+) -> None:
+    integration = integration_db_bo.get_by_id(integration_id)
+    if integration:
+        integration.state = enums.CognitionMarkdownFileState.FAILED.value
+        general.flush_or_commit(with_commit)
+
+
 def __select_running_information_source_payloads(
     project_id: Optional[str] = None,
     only_running: bool = False,

business_objects/project.py

@@ -8,10 +8,7 @@
 
 from .. import enums
 from ..session import session
-from ..models import (
-    Project,
-    Record,
-)
+from ..models import Project, Record, Attribute
 from ..util import prevent_sql_injection
 
 QUEUE_PROJECT_NAME = "@@HIDDEN_QUEUE_PROJECT@@"
@@ -155,6 +152,39 @@ def get_all(organization_id: str) -> List[Project]:
     )
 
 
+def get_all_with_access_management(organization_id: str) -> List[Project]:
+    return (
+        session.query(Project)
+        .join(Attribute, Project.id == Attribute.project_id)
+        .filter(
+            Project.organization_id == organization_id,
+            Attribute.name.in_(["__ACCESS_GROUPS", "__ACCESS_USERS"]),  #
+            Attribute.name.in_(["__ACCESS_GROUPS", "__ACCESS_USERS"]),
+            Attribute.user_created == False,
+            Attribute.data_type == enums.DataTypes.PERMISSION.value,
+            Attribute.state == enums.AttributeState.AUTOMATICALLY_CREATED.value,
+        )
+        .distinct()
+        .all()
+    )
+
+
+def check_access_management_active(project_id: str) -> bool:
+    return (
+        session.query(Project)
+        .join(Attribute, Project.id == Attribute.project_id)
+        .filter(
+            Project.id == project_id,
+            Attribute.name.in_(["__ACCESS_GROUPS", "__ACCESS_USERS"]),
+            Attribute.user_created == False,
+            Attribute.data_type == enums.DataTypes.PERMISSION.value,
+            Attribute.state == enums.AttributeState.AUTOMATICALLY_CREATED.value,
+        )
+        .count()
+        > 0
+    )
+
+
 def get_all_by_user_organization_id(organization_id: str) -> List[Project]:
     projects = (
         session.query(Project).filter(Project.organization_id == organization_id).all()

business_objects/record.py

@@ -1,6 +1,6 @@
 from __future__ import with_statement
 from typing import List, Dict, Any, Optional, Tuple, Iterable
-from sqlalchemy import cast, Text
+from sqlalchemy import cast, Text, String
 from sqlalchemy.orm.attributes import flag_modified
 from sqlalchemy.sql.expression import bindparam
 from sqlalchemy import update
@@ -609,7 +609,7 @@ def count_missing_tokenized_records(project_id: str) -> int:
     query = f"""
     SELECT COUNT(*)
     FROM (
-        {get_records_without_tokenization(project_id, None, query_only = True)}
+        {get_records_without_tokenization(project_id, None, query_only=True)}
     ) record_query
     """
     return general.execute_first(query)[0]
@@ -807,6 +807,25 @@ def delete_user_created_attribute(
     general.flush_or_commit(with_commit)
 
 
+def delete_access_management_attributes(
+    project_id: str, with_commit: bool = True
+) -> None:
+    access_groups_attribute_item = attribute.get_by_name(project_id, "__ACCESS_GROUPS")
+    access_users_attribute_item = attribute.get_by_name(project_id, "__ACCESS_USERS")
+
+    if access_users_attribute_item and access_groups_attribute_item:
+        record_items = get_all(project_id=project_id)
+        for i, record_item in enumerate(record_items):
+            if record_item.data.get(access_groups_attribute_item.name):
+                del record_item.data[access_groups_attribute_item.name]
+            if record_item.data.get(access_users_attribute_item.name):
+                del record_item.data[access_users_attribute_item.name]
+            flag_modified(record_item, "data")
+            if (i + 1) % 1000 == 0:
+                general.flush_or_commit(with_commit)
+        general.flush_or_commit(with_commit)
+
+
 def delete_duplicated_rats(with_commit: bool = False) -> None:
     # no project so run for all to prevent expensive join with record table
     query = """
@@ -925,3 +944,19 @@ def get_first_no_text_column(project_id: str, record_id: str) -> str:
     WHERE r.project_id = '{project_id}' AND r.id = '{record_id}'
     """
     return general.execute_first(query)[0]
+
+
+def get_record_ids_by_running_ids(project_id: str, running_ids: List[int]) -> List[str]:
+    return [
+        row[0]
+        for row in (
+            session.query(cast(Record.id, String))
+            .filter(
+                Record.project_id == project_id,
+                Record.data[attribute.get_running_id_name(project_id)]
+                .as_integer()
+                .in_(running_ids),
+            )
+            .all()
+        )
+    ]

cognition_objects/group.py

@@ -0,0 +1,123 @@
+from datetime import datetime
+from typing import List, Optional
+from ..business_objects import general
+from ..session import session
+from ..models import CognitionGroup
+
+
+def get(group_id: str) -> CognitionGroup:
+    return session.query(CognitionGroup).filter(CognitionGroup.id == group_id).first()
+
+
+def get_with_organization_id(organization_id: str, group_id: str) -> CognitionGroup:
+    return (
+        session.query(CognitionGroup)
+        .filter(
+            CognitionGroup.organization_id == organization_id,
+            CognitionGroup.id == group_id,
+        )
+        .first()
+    )
+
+
+def get_all(organization_id: str) -> List[CognitionGroup]:
+    return (
+        session.query(CognitionGroup)
+        .filter(CognitionGroup.organization_id == organization_id)
+        .order_by(CognitionGroup.name.asc())
+        .all()
+    )
+
+
+def get_all_by_integration_id(
+    organization_id: str, integration_id: str
+) -> List[CognitionGroup]:
+    integration_id_json = CognitionGroup.meta_data.op("->>")("integration_id")
+
+    return (
+        session.query(CognitionGroup)
+        .filter(
+            CognitionGroup.organization_id == organization_id,
+            integration_id_json == integration_id,
+        )
+        .order_by(CognitionGroup.name.asc())
+        .all()
+    )
+
+
+def get_all_by_integration_id_permission_grouped(
+    organization_id: str, integration_id: str
+) -> List[CognitionGroup]:
+    integration_id_json = CognitionGroup.meta_data.op("->>")("integration_id")
+
+    integration_groups = (
+        session.query(CognitionGroup)
+        .filter(
+            CognitionGroup.organization_id == organization_id,
+            integration_id_json == integration_id,
+        )
+        .all()
+    )
+    integration_groups_by_permission = {}
+    for group in integration_groups:
+        permission_id = group.meta_data.get("permission_id")
+        integration_groups_by_permission[permission_id] = group
+    return integration_groups_by_permission
+
+
+def get_by_name(organization_id: str, name: str):
+    return (
+        session.query(CognitionGroup)
+        .filter(
+            CognitionGroup.organization_id == organization_id,
+            CognitionGroup.name == name,
+        )
+        .first()
+    )
+
+
+def create_group(
+    organization_id: str,
+    name: str,
+    description: str,
+    created_by: str,
+    created_at: Optional[datetime] = None,
+    with_commit: bool = True,
+    meta_data: Optional[dict] = None,
+) -> CognitionGroup:
+    group = CognitionGroup(
+        organization_id=organization_id,
+        name=name,
+        description=description,
+        created_by=created_by,
+        created_at=created_at,
+        meta_data=meta_data,
+    )
+    general.add(group, with_commit)
+    return group
+
+
+def update_group(
+    group_id: str,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    with_commit: bool = True,
+    meta_data: Optional[dict] = None,
+) -> CognitionGroup:
+    group = get(group_id)
+
+    if name is not None:
+        group.name = name
+    if description is not None:
+        group.description = description
+    if meta_data is not None:
+        group.meta_data = meta_data
+    general.flush_or_commit(with_commit)
+
+    return group
+
+
+def delete(organization_id: str, group_id: str, with_commit: bool = True) -> None:
+    group = get_with_organization_id(organization_id, group_id)
+    if group:
+        general.delete(group, with_commit)

cognition_objects/group_member.py

@@ -0,0 +1,103 @@
+from datetime import datetime
+from typing import Optional
+from ..business_objects import general, user
+from . import group
+from ..session import session
+from ..models import CognitionGroupMember
+
+
+def get(group_id: str, id: str):
+    return (
+        session.query(CognitionGroupMember)
+        .filter(
+            CognitionGroupMember.group_id == group_id, CognitionGroupMember.id == id
+        )
+        .first()
+    )
+
+
+def get_by_group_and_user(group_id: str, user_id: str) -> CognitionGroupMember:
+    return (
+        session.query(CognitionGroupMember)
+        .filter(
+            CognitionGroupMember.group_id == group_id,
+            CognitionGroupMember.user_id == user_id,
+        )
+        .first()
+    )
+
+
+def get_by_user_id(user_id: str) -> list:
+    return (
+        session.query(CognitionGroupMember)
+        .filter(CognitionGroupMember.user_id == user_id)
+        .all()
+    )
+
+
+def get_all_by_group(group_id: str) -> list:
+    return (
+        session.query(CognitionGroupMember)
+        .filter(CognitionGroupMember.group_id == group_id)
+        .all()
+    )
+
+
+def get_all_by_group_count(group_id: str) -> int:
+    return (
+        session.query(CognitionGroupMember)
+        .filter(CognitionGroupMember.group_id == group_id)
+        .count()
+    )
+
+
+def create(
+    group_id: str,
+    user_id: str,
+    created_at: Optional[datetime] = None,
+    with_commit: bool = True,
+) -> CognitionGroupMember:
+    already_exist = get_by_group_and_user(group_id=group_id, user_id=user_id)
+    if already_exist:
+        return already_exist
+
+    group_item = group.get(group_id)
+    user_item = user.get(user_id)
+    if not group_item or not user_item:
+        raise Exception("Group or user not found")
+    if group_item.organization_id != user_item.organization_id:
+        raise Exception("User not in the same organization as the group")
+
+    group_member = CognitionGroupMember(
+        group_id=group_id,
+        user_id=user_id,
+        created_at=created_at,
+    )
+    general.add(group_member, with_commit)
+    return group_member
+
+
+def delete_by_group_and_user_id(
+    group_id: str, user_id: str, with_commit: bool = True
+) -> None:
+    group_member = get_by_group_and_user(group_id, user_id)
+    if group_member:
+        general.delete(group_member, with_commit)
+
+
+def delete_by_user_id(user_id: str, with_commit: bool = True) -> None:
+    group_memberships = (
+        session.query(CognitionGroupMember)
+        .filter(CognitionGroupMember.user_id == user_id)
+        .all()
+    )
+    for membership in group_memberships:
+        general.delete(membership, with_commit=False)
+    general.flush_or_commit(with_commit)
+
+
+def clear_by_group_id(group_id: str, with_commit: bool = True) -> None:
+    group_memberships = get_all_by_group(group_id)
+    for membership in group_memberships:
+        general.delete(membership, with_commit=False)
+    general.flush_or_commit(with_commit)