Release updates (#205)

* fix prepare-record-export

* fix create-outlier-slice via create-data-slice

* Fixes wizard wait for tokenization

* Hotfix

* Uses new free api route

* Fix export projects with key

* Fixes pw snapshot upload

* Submodule update

* Remove double org request

* Cache kratos

* slight changes

* cache update

* Fixes gql request foruser assignment

* Remove print

* fixes snakecase key prepare_record_export

* Submodules merge

---------

Co-authored-by: JWittmeyer <[email protected]>
Co-authored-by: Lina <[email protected]>

Author: 3 people

api/transfer.py

@@ -441,7 +441,7 @@ def __calculate_missing_attributes(project_id: str, user_id: str) -> None:
             if i >= 60:
                 i = 0
                 ctx_token = general.remove_and_refresh_session(ctx_token, True)
-            if tokenization.is_doc_bin_creation_running(project_id):
+            if tokenization.is_doc_bin_creation_running_or_queued(project_id):
                 time.sleep(2)
                 continue
             else:

controller/auth/kratos.py

@@ -1,60 +1,133 @@
-from typing import Union, Any, List, Optional, Dict
-from requests import Response
+from typing import Union, Any, List, Dict
 import os
 import requests
 import logging
+from datetime import datetime, timedelta
 
 logging.basicConfig(level=logging.INFO)
 logger: logging.Logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
 KRATOS_ADMIN_URL = os.getenv("KRATOS_ADMIN_URL")
 
+# user_id -> {"identity" -> full identity, "simple" -> {"id": str, "mail": str, "firstName": str, "lastName": str}}
+# "collected" -> timestamp
+KRATOS_IDENTITY_CACHE: Dict[str, Any] = {}
+KRATOS_IDENTITY_CACHE_TIMEOUT = timedelta(minutes=30)
+
+
+def __get_cached_values() -> Dict[str, Dict[str, Any]]:
+    global KRATOS_IDENTITY_CACHE
+    if not KRATOS_IDENTITY_CACHE or len(KRATOS_IDENTITY_CACHE) == 0:
+        __refresh_identity_cache()
+    elif (
+        KRATOS_IDENTITY_CACHE["collected"] + KRATOS_IDENTITY_CACHE_TIMEOUT
+        < datetime.now()
+    ):
+        __refresh_identity_cache()
+    return KRATOS_IDENTITY_CACHE
+
+
+def __refresh_identity_cache():
+    global KRATOS_IDENTITY_CACHE
+    request = requests.get(f"{KRATOS_ADMIN_URL}/identities")
+    if request.ok:
+        collected = datetime.now()
+        KRATOS_IDENTITY_CACHE = {
+            identity["id"]: {
+                "identity": identity,
+                "simple": __parse_identity_to_simple(identity),
+            }
+            for identity in request.json()
+        }
+        KRATOS_IDENTITY_CACHE["collected"] = collected
+    else:
+        KRATOS_IDENTITY_CACHE = {}
+
+
+def __get_identity(user_id: str, only_simple: bool = True) -> Dict[str, Any]:
+    if not isinstance(user_id, str):
+        user_id = str(user_id)
+    cache = __get_cached_values()
+    if user_id in cache:
+        if only_simple:
+            return cache[user_id]["simple"]
+        return cache[user_id]
+
+    if len(user_id) == 36:
+        # check not new entry outside cache
+        request = requests.get(f"{KRATOS_ADMIN_URL}/identities/{user_id}")
+        if request.ok:
+            identity = request.json()
+            if identity["id"] == user_id:
+                KRATOS_IDENTITY_CACHE[user_id] = {
+                    "identity": identity,
+                    "simple": __parse_identity_to_simple(identity),
+                }
+                if only_simple:
+                    return KRATOS_IDENTITY_CACHE[user_id]["simple"]
+                return KRATOS_IDENTITY_CACHE[user_id]
+    # e.g. if id "GOLD_STAR" is requested => wont be in cache but expects a dummy dict
+    if only_simple:
+        return __parse_identity_to_simple({"id": user_id})
+    return {
+        "identity": {
+            "id": user_id,
+            "traits": {"email": None, "name": {"first": None, "last": None}},
+        }
+    }
+
+
+def __parse_identity_to_simple(identity: Dict[str, Any]) -> Dict[str, str]:
+    r = {
+        "id": identity["id"],
+        "mail": None,
+        "firstName": None,
+        "lastName": None,
+    }
+    if "traits" in identity:
+        r["mail"] = identity["traits"]["email"]
+        if "name" in identity["traits"]:
+            r["firstName"] = identity["traits"]["name"]["first"]
+            r["lastName"] = identity["traits"]["name"]["last"]
+    return r
+
 
 def get_userid_from_mail(user_mail: str) -> str:
-    for identity in requests.get(f"{KRATOS_ADMIN_URL}/identities").json():
-        if identity["traits"]["email"] == user_mail:
-            return identity["id"]
+    values = __get_cached_values()
+    for key in values:
+        if key == "collected":
+            continue
+        if values[key]["simple"]["mail"] == user_mail:
+            return key
     return None
 
 
 def resolve_user_mail_by_id(user_id: str) -> str:
-    res: Response = requests.get("{}/identities/{}".format(KRATOS_ADMIN_URL, user_id))
-    data: Any = res.json()
-    if res.status_code == 200 and data["traits"]:
-        return data["traits"]["email"]
+    i = __get_identity(user_id)
+    if i:
+        return i["mail"]
     return None
 
 
 def resolve_user_name_by_id(user_id: str) -> str:
-    res: Response = requests.get("{}/identities/{}".format(KRATOS_ADMIN_URL, user_id))
-    data: Any = res.json()
-    if res.status_code == 200 and data["traits"]:
-        return data["traits"]["name"]
+    i = __get_identity(user_id, False)
+    if i:
+        i = i["identity"]
+        return i["traits"]["name"]
     return None
 
 
 def resolve_all_user_ids(
     relevant_ids: List[str], as_list: bool = True
-) -> List[Dict[str, str]]:
+) -> Union[Dict[str, Dict[str, str]], List[Dict[str, str]]]:
     final = [] if as_list else {}
     for id in relevant_ids:
-        r = requests.get(f"{KRATOS_ADMIN_URL}/identities/{id}").json()
-        d = {
-            "id": id,
-            "mail": None,
-            "firstName": None,
-            "lastName": None,
-        }
-        if "traits" in r:
-            traits = r["traits"]
-            d["mail"] = traits["email"]
-            d["firstName"] = traits["name"]["first"]
-            d["lastName"] = traits["name"]["last"]
+        i = __get_identity(id)
         if as_list:
-            final.append(d)
+            final.append(i)
         else:
-            final[id] = d
+            final[id] = i
     return final
 
 
@@ -63,25 +136,16 @@ def expand_user_mail_name(
 ) -> List[Dict[str, str]]:
     final = []
     for user in users:
-        r = requests.get(f"{KRATOS_ADMIN_URL}/identities/{user[user_id_key]}").json()
-        d = {
-            "mail": None,
-            "firstName": None,
-            "lastName": None,
-        }
-        if "traits" in r:
-            traits = r["traits"]
-            d["mail"] = traits["email"]
-            d["firstName"] = traits["name"]["first"]
-            d["lastName"] = traits["name"]["last"]
-        user = {**user, **d}
+        i = __get_identity(user[user_id_key])
+        user = {**user, **i}
         final.append(user)
     return final
 
 
 def resolve_user_name_and_email_by_id(user_id: str) -> dict:
-    res: Response = requests.get("{}/identities/{}".format(KRATOS_ADMIN_URL, user_id))
-    data: Any = res.json()
-    if res.status_code == 200 and data["traits"]:
-        return data["traits"]["name"], data["traits"]["email"]
+    i = __get_identity(user_id, False)
+    if i:
+        i = i["identity"]
+    if i and "traits" in i and i["traits"]:
+        return i["traits"]["name"], i["traits"]["email"]
     return None

controller/auth/manager.py

@@ -18,10 +18,10 @@
 
 
 def get_organization_id_by_info(info) -> Organization:
-    organization: Organization = get_user_by_info(info).organization
-    if not organization:
+    user = get_user_by_info(info)
+    if not user or not user.organization_id:
         raise GraphQLError("User is not associated to an organization")
-    return organization
+    return str(user.organization_id)
 
 
 def get_user_by_info(info) -> User:
@@ -76,7 +76,7 @@ def check_project_access_dep(request: Request, project_id: str):
 
 
 def check_project_access(info, project_id: str) -> None:
-    organization_id: str = get_organization_id_by_info(info).id
+    organization_id: str = get_organization_id_by_info(info)
     project: Project = project_manager.get_project_with_orga_id(
         organization_id, project_id
     )

controller/data_slice/manager.py

@@ -1,4 +1,3 @@
-import json
 from typing import Dict, Any, List, Optional
 
 from submodules.model import DataSlice
@@ -79,7 +78,7 @@ def create_data_slice(
         project_id=project_id,
         created_by=user_id,
         name=name,
-        filter_raw=json.loads(filter_raw),
+        filter_raw=filter_raw,
         filter_data=filter_data,
         static=static,
         slice_type=slice_type,

controller/organization/manager.py

@@ -56,14 +56,18 @@ def get_user_info(user) -> User:
     return user_expanded
 
 
-def get_all_users(organization_id: str, user_role: Optional[str] = None) -> List[User]:
+def get_all_users(
+    organization_id: str, user_role: Optional[str] = None, as_dict: bool = True
+) -> List[User]:
     parsed = None
     if user_role:
         try:
             parsed = enums.UserRoles[user_role.upper()]
         except KeyError:
             raise ValueError(f"Invalid UserRoles: {user_role}")
     all_users = user.get_all(organization_id, parsed)
+    if not as_dict:
+        return all_users
     all_users_dict = sql_alchemy_to_dict(
         all_users, column_whitelist=USER_INFO_WHITELIST
     )

controller/record/manager.py

@@ -164,7 +164,7 @@ def edit_records(
         tokenization_service.request_tokenize_project(project_id, user_id)
         time.sleep(1)
         # wait for tokenization to finish, the endpoint itself handles missing docbins
-        while tokenization.is_doc_bin_creation_running(project_id):
+        while tokenization.is_doc_bin_creation_running_or_queued(project_id):
             time.sleep(0.5)
 
     except Exception as e:
@@ -255,7 +255,7 @@ def __check_and_prep_edit_records(
             useable_embeddings[embedding_item.attribute_id] = []
         useable_embeddings[embedding_item.attribute_id].append(embedding_item)
 
-    if tokenization.is_doc_bin_creation_running(project_id):
+    if tokenization.is_doc_bin_creation_running_or_queued(project_id):
         errors_found.append(
             "tokenization is currently running. Wait for it to finish before editing records."
         )

controller/task_queue/handler/tokenization.py

@@ -5,7 +5,9 @@
     attribute as attribute_db_bo,
 )
 from controller.tokenization import tokenization_service
-from submodules.model.business_objects.tokenization import is_doc_bin_creation_running
+from submodules.model.business_objects.tokenization import (
+    is_doc_bin_creation_running_or_queued,
+)
 from submodules.model.enums import RecordTokenizationScope
 from ..util import if_task_queue_send_websocket
 
@@ -51,4 +53,4 @@ def __start_task(task: Dict[str, Any]) -> bool:
 
 
 def __check_finished(task: Dict[str, Any]) -> bool:
-    return not is_doc_bin_creation_running(task["project_id"])
+    return not is_doc_bin_creation_running_or_queued(task["project_id"], True)

controller/transfer/cognition/constants.py

@@ -242,4 +242,4 @@ def get_embeddings(self):
     "en": "doc2query/msmarco-t5-base-v1",
 }
 
-FREE_API_REQUEST_URL = "https://free.api.kern.ai/inference"
+FREE_API_REQUEST_URL = "https://k8s.freeapi.kern.ai/inference"

controller/transfer/cognition/import_wizard.py

@@ -193,11 +193,17 @@ def __finalize_setup(
         if c > 120:
             token_ref.request_new()
             c = 0
-        if tokenization_db_bo.is_doc_bin_creation_running(reference_project_id):
+        if tokenization_db_bo.is_doc_bin_creation_running_or_queued(
+            reference_project_id
+        ):
             continue
-        if tokenization_db_bo.is_doc_bin_creation_running(question_project_id):
+        if tokenization_db_bo.is_doc_bin_creation_running_or_queued(
+            question_project_id
+        ):
             continue
-        if tokenization_db_bo.is_doc_bin_creation_running(relevance_project_id):
+        if tokenization_db_bo.is_doc_bin_creation_running_or_queued(
+            relevance_project_id
+        ):
             continue
         break
 
@@ -313,6 +319,7 @@ def __finalize_setup_for(
     token_ref: TokenRef,
 ) -> Token:
     target_data = {"TARGET_LANGUAGE": project_language}
+
     # attributes
     attributes = project_type.get_attributes()
     target_data["target_type"] = "ac"

fast_api/models.py

@@ -305,7 +305,12 @@ class UpdateProjectStatusBody(BaseModel):
 
 
 class PrepareProjectExportBody(BaseModel):
-    export_options: StrictStr
+    export_options: Any
+    key: Optional[StrictStr] = None
+
+
+class PrepareRecordExportBody(BaseModel):
+    export_options: Any
     key: Optional[StrictStr] = None
 
 

fast_api/routes/data_browser.py

@@ -163,7 +163,7 @@ def create_data_slice(
 ):
     name = body.name
     static = body.static
-    filter_raw = body.filterRaw
+    filter_raw = json.loads(body.filterRaw)
     filter_data = [json.loads(item) for item in body.filterData]
 
     user = auth_manager.get_user_by_info(request.state.info)

fast_api/routes/data_slices.py

@@ -1,7 +1,7 @@
 from typing import Optional
 
 from fastapi import APIRouter, Depends, Request
-from fast_api.routes.client_response import pack_json_result, wrap_content_for_frontend
+from fast_api.routes.client_response import pack_json_result
 from submodules.model.util import sql_alchemy_to_dict
 from typing import List
 from controller.data_slice import manager
@@ -28,10 +28,12 @@ def get_data_slices(
         sql_alchemy_to_dict(ds, for_frontend=False)
         for ds in manager.get_all_data_slices(project_id, slice_type)
     ]
+
     for v in values:
         v["filterData"] = json.dumps(v["filter_data"])
         v["filterRaw"] = json.dumps(v["filter_raw"])
         del v["count_sql"]
+
     return pack_json_result(
         {"data": {"dataSlices": values}},
     )