Cognition ETL Provider (#201)

* perf: add etl task table

* perf(etl_task): add llm_config

* perf: rename EtlTask class

* perf: add business objects

* fix: optional file ids in etl task create

* perf: add file_path to etl_task

* perf: file paths

* perf: add tokenizer column

* perf: add file_size col

* perf: add split config

* perf: add split config to etl task

* fix: set split types to chunk and shrink

* perf: update enum name

* perf: default ETLExtractorPDF to PDF2MD

* perf: fkey alignment

* perf: add ETLTransformer

* perf: remove deleted cols

* perf: add monitor.set_etl_task_to_failed

* perf: EnumKern

* perf: add etl_task_id to integration records

* perf: enum alignment

* perf: add project update for integration

* perf: etl and integration deletions

* fix: thread start in general

* perf: add ETLExtractorPDF from_string

* perf: integration_objects cleanup

* perf: etl_task.cache_config

* perf: error print

* fix: overwrite etl_task_id fkey

* perf: make get_or_create_etl_task a submodule function

* fix: get_or_create_etl_task args

* perf: align to integrations

* perf: add ETLCacheKeys

* perf: add file-cache to ETLCacheKey enum

* perf: minor fixes

* New table

* Adds new enums

* Tmp commot

* tmp commit

* Helper method

* Helper for queue view

* util changes

* Adds optional fields for overwriting

* Adds transf key & extr key

* model iwth /

* PArse scope proj

* perf: etl utils update

* perf: make enum EnumKern

* fix: fallback config for integrations

* fix: cleanup integration fallback

* perf: update etl_utils

* perf: etl task original file name

* perf: models update
etl utils enhancements

* fix: localhost reference

* perf: integration - refinery load

* perf: update integration etl config

* perf: add get_splitting_key method

* perf: update config generation

* perf: update cache keys

* perf: rename get hashed str fn

* perf: function defs

* perf: update cache keys

* perf: etl file cache deletion

* perf: rework delete_etl cache

* fix: delete etl cache

* perf: query task queue for active tmp doc

* perf: enhance fn args

* Config preset creation

* Change PDF to documents

* perf: add langchain enums

* Removes outdated project columns

* perf: file_type infer

* style: error message

* style: error message

* Submodule change for refinery-test

* perf: add from_mimetype

* small change

* perf: rename ETLExtractor enums

* update

* fix: enum fix

* fix: llmindicator -> llmidentifier

* perf: minor enhancements

* perf: enum file type enhancement

* fix: integration notify config

* Fix query

* Fix project creation without tokenizer

* Fix image types

* perf: enhance file type inference

* fix: hash prompt for transformation cache

* perf: integration records

* perf: track integration progress

* perf: add supported extractors enum fn

* perf: add extractor validations

* perf: update dataset model

* fix: markdown_file etl config loading step

* perf: add meta_data to etl_task

* perf: remove dataset.llm_config col

* perf: various enhancements

---------

Co-authored-by: JWittmeyer <jens.wittmeyer@kern.ai>
Co-authored-by: Lina <lina.lumburovska@kern.ai>

Author: 3 people

business_objects/general.py

@@ -77,7 +77,7 @@ def force_remove_and_refresh_session_by_id(session_id: str) -> bool:
         if session_id not in session_lookup:
             return False
     # context vars cant be closed from a different context but we can work around it by using a thread (which creates a new context) with the same id
-    daemon.run_without_db_token(__close_in_context(session_id))
+    daemon.run_without_db_token(__close_in_context, session_id)
     return True
 
 

business_objects/monitor.py

@@ -1,10 +1,11 @@
 from typing import Any, List, Optional
 import datetime
-from . import general
-from submodules.model import enums, telemetry
-from submodules.model.models import TaskQueue, Organization
+from submodules.model.business_objects import general
+from submodules.model import enums
+from submodules.model.models import TaskQueue, Organization, CognitionIntegration
 from submodules.model.util import prevent_sql_injection
 from submodules.model.session import session
+from submodules.model.global_objects import etl_task as etl_task_db_bo
 from submodules.model.cognition_objects import (
     macro as macro_db_bo,
     markdown_file as markdown_file_db_bo,
@@ -207,9 +208,9 @@ def set_integration_task_to_failed(
         enums.CognitionMarkdownFileState
     ] = enums.CognitionMarkdownFileState.FAILED,
     with_commit: bool = True,
-) -> None:
+) -> CognitionIntegration:
     # argument `state` is a workaround for cognition-gateway/api/routes/integrations.delete_many
-    integration_db_bo.update(
+    return integration_db_bo.update(
         id=integration_id,
         state=state,
         finished_at=datetime.datetime.now(datetime.timezone.utc),
@@ -220,6 +221,26 @@ def set_integration_task_to_failed(
     )
 
 
+def set_etl_task_to_failed(
+    id: str,
+    is_active: bool = False,
+    error_message: Optional[str] = None,
+    state: Optional[
+        enums.CognitionMarkdownFileState
+    ] = enums.CognitionMarkdownFileState.FAILED,
+    with_commit: bool = True,
+) -> None:
+    # argument `state` is a workaround for cognition-gateway/api/routes/integrations.delete_many
+    etl_task_db_bo.update(
+        id=id,
+        state=state,
+        finished_at=datetime.datetime.now(datetime.timezone.utc),
+        is_active=is_active,
+        error_message=error_message,
+        with_commit=with_commit,
+    )
+
+
 def __select_running_information_source_payloads(
     project_id: Optional[str] = None,
     only_running: bool = False,

business_objects/task_queue.py

@@ -33,6 +33,23 @@ def get_orphan_tasks() -> List[TaskQueue]:
     )
 
 
+def get_all_queued_etl_task_for_conversation(
+    org_id: str, project_id: str, conversation_id: str
+) -> Optional[List[TaskQueue]]:
+    return (
+        session.query(TaskQueue)
+        .filter(
+            TaskQueue.organization_id == org_id,
+            TaskQueue.task_type == enums.TaskType.EXECUTE_ETL.value,
+            text(f"task_info->'tmp_doc_metadata'->>'project_id' = '{project_id}'"),
+            text(
+                f"task_info->'tmp_doc_metadata'->>'conversation_id' = '{conversation_id}'"
+            ),
+        )
+        .all()
+    )
+
+
 def get_likely_failed_tasks(days: int = 1) -> List[TaskQueue]:
     return (
         session.query(TaskQueue)

cognition_objects/environment_variable.py

@@ -6,7 +6,6 @@
 from ..models import (
     CognitionEnvironmentVariable,
     CognitionMarkdownDataset,
-    CognitionProject,
     GraphRAGIndex,
 )
 from ..util import prevent_sql_injection
@@ -64,6 +63,22 @@ def get_by_name_and_org_id(
     )
 
 
+def get_by_id_and_org_id(
+    org_id: str,
+    id: str,
+) -> CognitionEnvironmentVariable:
+
+    return (
+        session.query(CognitionEnvironmentVariable)
+        .filter(
+            CognitionEnvironmentVariable.organization_id == org_id,
+            CognitionEnvironmentVariable.project_id == None,
+            CognitionEnvironmentVariable.id == id,
+        )
+        .first()
+    )
+
+
 def get_dataset_env_var_value(
     dataset_id: str, org_id: str, scope: Literal["extraction", "transformation"]
 ) -> Union[str, None]:
@@ -122,43 +137,6 @@ def get_all_by_project_id(project_id: str) -> List[CognitionEnvironmentVariable]
     )
 
 
-def get_cognition_project_env_var_value(cognition_project_id: str) -> str:
-
-    env_var_id = cast(
-        CognitionProject.llm_config.op("->")("transformation").op("->>")("envVarId"),
-        UUID,
-    )
-    v = (
-        session.query(CognitionEnvironmentVariable.value)
-        .join(CognitionProject, env_var_id == CognitionEnvironmentVariable.id)
-        .filter(
-            CognitionProject.id == cognition_project_id,
-        )
-        .first()
-    )
-    if v and v[0]:
-        return str(v[0])
-
-
-def get_cognition_project_extraction_env_var_value(
-    cognition_project_id: str, envVar: str
-) -> str:
-    env_var_id = cast(
-        CognitionProject.llm_config.op("->")("extraction").op("->>")(envVar),
-        UUID,
-    )
-    v = (
-        session.query(CognitionEnvironmentVariable.value)
-        .join(CognitionProject, env_var_id == CognitionEnvironmentVariable.id)
-        .filter(
-            CognitionProject.id == cognition_project_id,
-        )
-        .first()
-    )
-    if v and v[0]:
-        return str(v[0])
-
-
 def get_cognition_graphrag_env_var_value(org_id: str, graphrag_index_id: str) -> str:
     # currently not in use because of fixed env var
 

cognition_objects/etl_config_presets.py

@@ -0,0 +1,84 @@
+from typing import Any, List, Optional, Dict
+from ..business_objects import general
+from ..session import session
+from ..models import (
+    ETLConfigPresets,
+)
+
+
+def get(config_id: str) -> ETLConfigPresets:
+    return (
+        session.query(ETLConfigPresets)
+        .filter(
+            ETLConfigPresets.id == config_id,
+        )
+        .first()
+    )
+
+
+def get_all_in_org(
+    org_id: str,
+) -> List[ETLConfigPresets]:
+
+    return (
+        session.query(ETLConfigPresets)
+        .filter(
+            ETLConfigPresets.organization_id == org_id,
+        )
+        .order_by(ETLConfigPresets.created_at.asc())
+        .all()
+    )
+
+
+def create(
+    org_id: str,
+    user_id: str,
+    name: str,
+    description: str,
+    etl_config: Dict[str, Any],
+    add_config: Dict[str, Any],
+    with_commit: bool = True,
+) -> ETLConfigPresets:
+    etl_config: ETLConfigPresets = ETLConfigPresets(
+        created_by=user_id,
+        organization_id=org_id,
+        name=name,
+        description=description,
+        etl_config=etl_config,
+        add_config=add_config,
+    )
+    general.add(etl_config, with_commit)
+    return etl_config
+
+
+def update(
+    org_id: str,
+    etl_config_id: str,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    etl_config: Optional[Dict[str, Any]] = None,
+    add_config: Optional[Dict[str, Any]] = None,
+    with_commit: bool = True,
+) -> ETLConfigPresets:
+    etl_config_item: ETLConfigPresets = get(etl_config_id)
+    if not etl_config_item or str(etl_config_item.organization_id) != org_id:
+        raise Exception("ETL Config not found")
+
+    if name is not None:
+        etl_config_item.name = name
+    if description is not None:
+        etl_config_item.description = description
+    if etl_config is not None:
+        etl_config_item.etl_config = etl_config
+    if add_config is not None:
+        etl_config_item.add_config = add_config
+    general.flush_or_commit(with_commit)
+    return etl_config_item
+
+
+def delete(org_id: str, etl_config_id: str, with_commit: bool = True) -> None:
+    session.query(ETLConfigPresets).filter(
+        ETLConfigPresets.organization_id == org_id,
+        ETLConfigPresets.id == etl_config_id,
+    ).delete()
+    general.flush_or_commit(with_commit)

cognition_objects/integration.py

@@ -3,9 +3,11 @@
 from sqlalchemy import func
 from sqlalchemy.orm.attributes import flag_modified
 
+
 from ..business_objects import general
+from ..integration_objects import manager as integration_records_bo
 from ..session import session
-from ..models import CognitionIntegration, CognitionGroup
+from ..models import CognitionIntegration, CognitionGroup, EtlTask
 from ..enums import (
     CognitionMarkdownFileState,
     CognitionIntegrationType,
@@ -147,6 +149,49 @@ def get_last_synced_at(
     return result[0] if result else None
 
 
+def get_active_etl_tasks(
+    integration_id: str,
+) -> List[EtlTask]:
+    IntegrationModel = integration_records_bo.integration_model(integration_id)
+    return (
+        session.query(EtlTask)
+        .filter(EtlTask.is_active == True)
+        .join(
+            IntegrationModel,
+            (EtlTask.id == IntegrationModel.etl_task_id)
+            & (IntegrationModel.integration_id == integration_id),
+        )
+        .all()
+    )
+
+
+def get_all_etl_tasks(
+    integration_id: str,
+) -> List[EtlTask]:
+    IntegrationModel = integration_records_bo.integration_model(integration_id)
+    return (
+        session.query(EtlTask)
+        .join(
+            IntegrationModel,
+            (IntegrationModel.etl_task_id == EtlTask.id)
+            & (IntegrationModel.integration_id == integration_id),
+        )
+        .all()
+    )
+
+
+def get_integration_progress(
+    integration_id: str,
+) -> float:
+    count_all_records = integration_records_bo.count(integration_id)
+    all_tasks = get_all_etl_tasks(integration_id)
+    finished_tasks = [task for task in all_tasks if task.state in FINISHED_STATES]
+
+    if count_all_records == 0:
+        return 0.0
+    return round((len(finished_tasks) / count_all_records) * 100.0, 2)
+
+
 def count_org_integrations(org_id: str) -> Dict[str, int]:
     counts = (
         session.query(CognitionIntegration.type, func.count(CognitionIntegration.id))
@@ -201,6 +246,7 @@ def create(
 
 def update(
     id: str,
+    project_id: Optional[str] = None,
     updated_by: Optional[str] = None,
     name: Optional[str] = None,
     description: Optional[str] = None,
@@ -220,6 +266,8 @@ def update(
     if not integration:
         return None
 
+    if project_id is not None and integration.project_id is None:
+        integration.project_id = project_id
     if updated_by is not None:
         integration.updated_by = updated_by
     if name is not None:
@@ -279,6 +327,16 @@ def execution_finished(id: str) -> bool:
 def delete_many(
     ids: List[str], delete_cognition_groups: bool = True, with_commit: bool = True
 ) -> None:
+    for id in ids:
+        integration_records, IntegrationModel = (
+            integration_records_bo.get_all_by_integration_id(id)
+        )
+        integration_records_bo.delete_many(
+            IntegrationModel,
+            ids=[rec.id for rec in integration_records],
+            with_commit=True,
+        )
+
     (
         session.query(CognitionIntegration)
         .filter(CognitionIntegration.id.in_(ids))
@@ -290,6 +348,7 @@ def delete_many(
             .filter(CognitionGroup.meta_data.op("->>")("integration_id").in_(ids))
             .delete(synchronize_session=False)
         )
+
     general.flush_or_commit(with_commit)