streamlit-octostar-utils 0.4.2.dev25__tar.gz → 0.5.0.dev2__tar.gz

{streamlit_octostar_utils-0.4.2.dev25 → streamlit_octostar_utils-0.5.0.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: streamlit-octostar-utils
-Version: 0.4.2.dev25
+Version: 0.5.0.dev2
 Summary: 
 License: MIT
 License-File: LICENSE

{streamlit_octostar_utils-0.4.2.dev25 → streamlit_octostar_utils-0.5.0.dev2}/pyproject.toml

@@ -5,7 +5,7 @@ include = '\.pyi?$'
 
 [tool.poetry]
 name = "streamlit-octostar-utils"
-version = "0.4.2-dev.25"
+version = "0.5.0-dev.2"
 description = ""
 license = "MIT"
 authors = ["Octostar"]

{streamlit_octostar_utils-0.4.2.dev25 → streamlit_octostar_utils-0.5.0.dev2}/streamlit_octostar_utils/api_crafter/celery.py

@@ -9,7 +9,6 @@ import subprocess
 from fastapi import Query
 import time
 import os
-import pickle
 import atexit
 import redis
 import uuid
@@ -68,7 +67,31 @@ class CeleryQueueConfig:
         self.options = options
 
 
+class TaskResult:
+    """Wrapper for task results that include binary parts alongside JSON data.
+    Tasks returning binary data (e.g. images) should return a TaskResult
+    so that serialized_io writes them as multipart parts instead of attempting
+    JSON serialization on bytes."""
+
+    def __init__(self, data, part=None):
+        self.data = data
+        self.part = part
+
+
 class CelerySerialized:
+    """Serializes task data to a boundary-delimited multipart file.
+
+    Format: metadata JSON part followed by optional binary/streamed parts,
+    separated by boundary markers (the task_id). Replaces pickle entirely.
+    """
+
+    BOUNDARY_PREFIX = b"--"
+    BOUNDARY_SUFFIX = b"\r\n"
+    BOUNDARY_END = b"--\r\n"
+    CONTENT_TYPE_JSON = b"Content-Type: application/json\r\n"
+    CONTENT_TYPE_BYTES = b"Content-Type: application/octet-stream\r\n"
+    HEADER_END = b"\r\n"
+
     def __init__(self, folder, redis_client, data=None):
         self.folder = folder
         self.data = data
@@ -77,18 +100,70 @@ class CelerySerialized:
     def set_task_id(self, task_id):
         self.task_id = task_id
 
-    def dump(self):
+    def _boundary(self):
+        return self.task_id.encode()
+
+    def _write_boundary(self, f):
+        f.write(self.BOUNDARY_PREFIX + self._boundary() + self.BOUNDARY_SUFFIX)
+
+    def _write_end_boundary(self, f):
+        f.write(self.BOUNDARY_PREFIX + self._boundary() + self.BOUNDARY_END)
+
+    def dump(self, parts=None, part_is_list=False):
+        """Write metadata + optional parts in multipart format.
+        parts: optional list of bytes objects to write as additional parts.
+        part_is_list: whether the original part was a list (preserves type on read).
+        """
         assert self.task_id
+        if isinstance(self.data, dict):
+            metadata = self.data
+        else:
+            metadata = {"data": self.data}
+        metadata["part_count"] = len(parts) if parts else 0
+        metadata["part_is_list"] = part_is_list
         with RedisFileLock(self.redis_client, os.path.join(self.folder, self.task_id)):
-            with open(os.path.join(self.folder, self.task_id), "wb") as target_file:
-                pickle.dump(self.data, file=target_file, protocol=pickle.HIGHEST_PROTOCOL)
+            with open(os.path.join(self.folder, self.task_id), "wb") as f:
+                self._write_boundary(f)
+                f.write(self.CONTENT_TYPE_JSON)
+                f.write(self.HEADER_END)
+                f.write(json.dumps(metadata).encode())
+                f.write(b"\r\n")
+                if parts:
+                    for part in parts:
+                        self._write_boundary(f)
+                        f.write(self.CONTENT_TYPE_BYTES)
+                        f.write(self.HEADER_END)
+                        f.write(part)
+                        f.write(b"\r\n")
+                self._write_end_boundary(f)
 
     def load(self):
+        """Read multipart file. Returns (metadata_dict, list_of_bytes_parts)."""
         assert self.task_id
+        boundary = self.BOUNDARY_PREFIX + self._boundary()
+        end_boundary = self.BOUNDARY_PREFIX + self._boundary() + b"--"
         with RedisFileLock(self.redis_client, os.path.join(self.folder, self.task_id)):
-            with open(os.path.join(self.folder, self.task_id), "rb") as source_file:
-                data = pickle.load(source_file)
-        return data
+            with open(os.path.join(self.folder, self.task_id), "rb") as f:
+                raw = f.read()
+        sections = raw.split(boundary)
+        metadata = None
+        parts = []
+        for section in sections:
+            section = section.strip(b"\r\n")
+            if not section or section == b"--":
+                continue
+            header_end = section.find(b"\r\n\r\n")
+            if header_end == -1:
+                continue
+            header = section[:header_end]
+            body = section[header_end + 4:]
+            if body.endswith(b"\r\n"):
+                body = body[:-2]
+            if b"application/json" in header:
+                metadata = json.loads(body)
+            else:
+                parts.append(body)
+        return metadata or {}, parts
 
 
 class CeleryExecutor(object):
@@ -523,9 +598,16 @@ class CeleryExecutor(object):
             task_id = task.request.id
             serialized_data = CelerySerialized(folder=self.in_folder, redis_client=self.redis_client)
             serialized_data.set_task_id(task_id)
-            data = serialized_data.load()
+            metadata, parts = serialized_data.load()
             del serialized_data
-            args, kwargs = data.get("args", []), data.get("kwargs", {})
+            args, kwargs = metadata.get("args", []), metadata.get("kwargs", {})
+
+            part_count = metadata.get("part_count", 0)
+            if part_count > 0:
+                if metadata.get("part_is_list", part_count > 1):
+                    args = [parts] + args
+                else:
+                    args = [parts[0]] + args
 
             if self.app.conf.task_always_eager:
                 queue = task.request.delivery_info.get("routing_key", self.app.conf.task_default_routing_key)
@@ -536,9 +618,25 @@ class CeleryExecutor(object):
             queue = task.request.delivery_info.get("routing_key", self.app.conf.task_default_routing_key)
             task.request.resources = (self.resource_registry or {}).get(queue, {})
             out_data = task_fn(task, *args, **kwargs)
-            serialized_data = CelerySerialized(folder=self.out_folder, data=out_data, redis_client=self.redis_client)
+            if isinstance(out_data, TaskResult):
+                out_part_is_list = isinstance(out_data.part, list)
+                if out_data.part is None:
+                    out_parts = None
+                elif out_part_is_list:
+                    out_parts = out_data.part
+                else:
+                    out_parts = [out_data.part]
+                serialized_data = CelerySerialized(
+                    folder=self.out_folder, data=out_data.data, redis_client=self.redis_client
+                )
+            else:
+                out_parts = None
+                out_part_is_list = False
+                serialized_data = CelerySerialized(
+                    folder=self.out_folder, data=out_data, redis_client=self.redis_client
+                )
             serialized_data.set_task_id(task_id)
-            serialized_data.dump()
+            serialized_data.dump(parts=out_parts, part_is_list=out_part_is_list)
             del serialized_data
             if os.path.isfile(os.path.join(self.in_folder, task_id)):
                 with RedisFileLock(self.redis_client, os.path.join(self.in_folder, task_id)):
@@ -579,7 +677,7 @@ class CeleryExecutor(object):
 
         return decorator
 
-    async def send_task(self, task_fn, args=None, kwargs=None, **options) -> str:
+    async def send_task(self, task_fn, args=None, kwargs=None, part=None, **options) -> str:
         args = args if args is not None else []
         kwargs = kwargs if kwargs is not None else {}
         if self.app.conf.task_always_eager and "dev_preload" not in self.app.conf:
@@ -635,14 +733,21 @@ class CeleryExecutor(object):
                 await asyncio.get_running_loop().run_in_executor(
                     self.set_thread_pool, _check_queue_llen, queue_name
                 )
-            await asyncio.get_running_loop().run_in_executor(
-                self.io_thread_pool,
-                _write_task_data,
-                self.in_folder,
-                args,
-                kwargs,
-                task_id,
-            )
+
+            if part is not None:
+                await self._write_task_data_with_part(
+                    task_id, args, kwargs, part
+                )
+            else:
+                await asyncio.get_running_loop().run_in_executor(
+                    self.io_thread_pool,
+                    _write_task_data,
+                    self.in_folder,
+                    args,
+                    kwargs,
+                    task_id,
+                )
+
             await asyncio.get_running_loop().run_in_executor(
                 self.set_thread_pool, _send_task, task_fn, task_id, options
             )
@@ -664,6 +769,50 @@ class CeleryExecutor(object):
                 sem.release()
         return task_id
 
+    async def _write_task_data_with_part(self, task_id, args, kwargs, part):
+        """Write task data with a streamed part to the multipart file.
+        The part becomes the first arg on the worker side.
+        """
+        boundary = CelerySerialized.BOUNDARY_PREFIX + task_id.encode()
+        boundary_line = boundary + CelerySerialized.BOUNDARY_SUFFIX
+        end_boundary_line = boundary + CelerySerialized.BOUNDARY_END
+
+        is_list = isinstance(part, list)
+        items = part if is_list else [part]
+        part_count = len(items)
+
+        metadata = {"args": args, "kwargs": kwargs, "part_count": part_count, "part_is_list": is_list}
+        metadata_bytes = json.dumps(metadata).encode()
+
+        file_path = os.path.join(self.in_folder, task_id)
+        with open(file_path, "wb") as f:
+            f.write(boundary_line)
+            f.write(CelerySerialized.CONTENT_TYPE_JSON)
+            f.write(CelerySerialized.HEADER_END)
+            f.write(metadata_bytes)
+            f.write(b"\r\n")
+
+            for item in items:
+                f.write(boundary_line)
+                f.write(CelerySerialized.CONTENT_TYPE_BYTES)
+                f.write(CelerySerialized.HEADER_END)
+                if hasattr(item, "__aiter__"):
+                    async for chunk in item:
+                        f.write(chunk if isinstance(chunk, bytes) else chunk.encode())
+                elif hasattr(item, "read"):
+                    while True:
+                        chunk = await item.read(65536)
+                        if not chunk:
+                            break
+                        f.write(chunk if isinstance(chunk, bytes) else chunk.encode())
+                elif isinstance(item, bytes):
+                    f.write(item)
+                else:
+                    raise TypeError(f"Unsupported part item type: {type(item)}")
+                f.write(b"\r\n")
+
+            f.write(end_boundary_line)
+
     async def terminate_task(self, task_id):
         def _terminate_task(celery_app, task_id):
             celery_app.control.revoke(task_id, terminate=True)
@@ -717,8 +866,15 @@ class CeleryExecutor(object):
         def _read_task_data(out_folder, task_id):
             serialized_data = CelerySerialized(folder=out_folder, redis_client=self.redis_client)
             serialized_data.set_task_id(task_id)
-            result = serialized_data.load()
-            return result
+            metadata, parts = serialized_data.load()
+            data = metadata.get("data", metadata)
+            part_count = metadata.get("part_count", 0)
+            if part_count > 0:
+                if metadata.get("part_is_list", part_count > 1):
+                    return TaskResult(data=data, part=parts)
+                else:
+                    return TaskResult(data=data, part=parts[0])
+            return data
 
         def _remove_task_data(celery_app, in_folder, out_folder, task_id):
             celery_app.AsyncResult(task_id).forget()
@@ -744,10 +900,10 @@ class CeleryExecutor(object):
             )
         return result
 
-    async def send_and_wait_task(self, task_fn, args=None, kwargs=None, timeout=60, **options):
+    async def send_and_wait_task(self, task_fn, args=None, kwargs=None, part=None, timeout=60, **options):
         args = args if args is not None else []
         kwargs = kwargs if kwargs is not None else {}
-        task_id = await self.send_task(task_fn, args, kwargs, **options)
+        task_id = await self.send_task(task_fn, args, kwargs, part=part, **options)
         ready = False
         state = None
         start_time = time.time()

{streamlit_octostar_utils-0.4.2.dev25 → streamlit_octostar_utils-0.5.0.dev2}/streamlit_octostar_utils/api_crafter/fastapi.py

@@ -1,5 +1,5 @@
 import asyncio
-from fastapi import Request, Body
+from fastapi import Request
 from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
 from typing import List, Optional, Literal, Any
@@ -17,8 +17,6 @@ import traceback
 from copy import copy
 import logging
     
-from octostar.client import make_client
-
 MAX_ERROR_MESSAGE_BYTES = 256
 MAX_ERROR_TRACEBACK_BYTES = 10240
 
@@ -99,100 +97,6 @@ class Route(ABC):
         return self
 
 
-class OctostarRoute(Route):
-    def __init__(self, app, tasks_routes, celery_executor=None, router=None):
-        self.app = app
-        self._router = router
-        self.routed_funcs = []
-        self.tasks_routes = tasks_routes
-        self.celery_executor = celery_executor
-        self.endpoints = {}
-        self.define_routes()
-
-    def register_route(self, op, octostar_task):
-        self.endpoints[op.strip("/")] = octostar_task
-
-    def define_routes(self):
-        if self.celery_executor:
-
-            @Route.route(self, path="/task-state/{task_id}")
-            async def get_task_status(task_id: str) -> JSONResponse:
-                task_status = await self.tasks_routes.get_task(task_id, pop=False)
-                task_status = task_status.model_dump(mode="json")["data"]["task_state"]
-                return JSONResponse(task_status)
-
-            @Route.route(self, path="/task-result/{task_id}")
-            async def get_task_result(task_id: str) -> JSONResponse:
-                return_data = await self.tasks_routes.get_task(task_id, pop=True)
-                return_data = return_data.model_dump(mode="json")["data"]["data"]
-                return JSONResponse(return_data)
-
-            @Route.route(self, path="/{op}", methods=["POST"])
-            async def send_task(
-                op: str,
-                os_context: dict = Body(...),
-                jwt: str = Body(...),
-                params: dict = Body(dict()),
-            ) -> str:
-                """
-                Any request coming from Octostar (e.g. manifest) should enter from here.
-                """
-                path_params = []
-                op = op.split("/")
-                if len(op) > 1:
-                    path_params = op[1:]
-                op = op[0]
-                query_params = params
-                client = make_client(jwt)
-                if op not in self.endpoints.keys():
-                    raise StarletteHTTPException(401, f"Route {op} is forbidden for NiFi.")
-                task_id = await self.celery_executor.send_task(
-                    self.endpoints[op], args=[os_context, client, query_params]
-                )
-                return task_id
-
-        else:
-
-            @Route.route(self, path="/{op}", methods=["POST"])
-            async def call_task(
-                op: str,
-                os_context: dict = Body(...),
-                jwt: str = Body(...),
-                params: dict = Body(dict()),
-            ) -> str:
-                """
-                Any request coming from Octostar (e.g. manifest) should enter from here.
-                """
-                path_params = []
-                op = op.split("/")
-                if len(op) > 1:
-                    path_params = op[1:]
-                op = op[0]
-                query_params = params
-                client = make_client(jwt)
-                if op not in self.endpoints.keys():
-                    raise StarletteHTTPException(401, f"Route {op} is forbidden for NiFi.")
-                result = await self.endpoints[op](os_context, client, query_params)
-                return result
-
-    @staticmethod
-    def octostar_task(celery_executor, *args, **opts):
-        def decorator(func):
-            if celery_executor:
-                serialized_func = celery_executor.serialized_io(func)
-                task_func = celery_executor.app.task(*args, **opts)(serialized_func)
-            else:
-
-                @wraps(func)
-                def octostar_func(*args, **kwargs):
-                    return func(None, *args, **kwargs)
-
-                task_func = octostar_func
-            return task_func
-
-        return decorator
-
-
 class CommonModels(object):
     class OKResponseModel(BaseModel):
         message: str = "OK"

{streamlit_octostar_utils-0.4.2.dev25 → streamlit_octostar_utils-0.5.0.dev2}/streamlit_octostar_utils/api_crafter/nifi.py

@@ -238,6 +238,156 @@ class NifiFragmenter(object):
             pointer.get("merge_params") or {}, defragmenter_config, lambda _, v2: v2
         )
 
+    _REQUIRED_FRAGMENT_FIELDS = ("index", "count", "identifier")
+
+    @staticmethod
+    def get_fragment_info(entity, fragmenter_keylist):
+        """Read fragment metadata (identifier, count, index, root_uid, merge_params)
+        for a given fragmenter level. Read-only -- does not mutate the entity.
+
+        Args:
+            entity: A NifiEntity or NifiEntityProxy.
+            fragmenter_keylist: Dot-separated key path into the fragment config
+                (e.g. "document_pages" or "audio_split").
+
+        Returns:
+            dict with keys like identifier, count, index, root_uid, merge_params.
+            Empty dict if fragmenter_keylist is empty or intermediate keys are
+            missing (entity not fragmented at this level).
+
+        Raises:
+            KeyError: If the final fragment key is missing from the config.
+            ValueError: If the fragment info exists but lacks required fields
+                (index, count, identifier).
+        """
+        if not fragmenter_keylist:
+            return {}
+        pointer = entity.request["config"]["fragment"]
+        for k in fragmenter_keylist.split(".")[:-1]:
+            if not pointer.get(k):
+                return {}
+            pointer = pointer[k]
+        info = pointer[fragmenter_keylist.split(".")[-1]]
+        missing = [f for f in NifiFragmenter._REQUIRED_FRAGMENT_FIELDS if f not in info]
+        if missing:
+            raise RuntimeError(
+                f"Fragment info for '{fragmenter_keylist}' is missing required "
+                f"field(s): {', '.join(missing)}"
+            )
+        return info
+
+    @staticmethod
+    def identify_fragment_groups(nifi_batches):
+        """Find all fragmented entities grouped by their active fragmenter level.
+
+        Args:
+            nifi_batches: List of NifiEntityBatch objects.
+
+        Returns:
+            dict mapping fragmenter_keylist to list of entities at that level.
+            Empty dict if no fragments found. Callers use get_fragment_info()
+            to fetch metadata per entity, and filter by index==0 to find roots.
+        """
+        all_entities = list(itertools.chain(*[b.entities for b in nifi_batches]))
+        groups = {}
+        for e in all_entities:
+            stack = e.request["config"].get("fragment", {}).get("fragments_stack", [])
+            if stack:
+                groups.setdefault(stack[0], []).append(e)
+        return groups
+
+    @staticmethod
+    def build_fragment_tree_from_children_entities(root_entity, fragmenter_keylist):
+        """Recursively build a tree from a root fragment entity by walking
+        its children_entities.
+
+        Args:
+            root_entity: The root entity (index 0) to start from.
+            fragmenter_keylist: The fragmenter level to build for.
+
+        Returns:
+            Nested dict with keys:
+                "entity": NifiEntity/NifiEntityProxy
+                "index": int
+                "merge_params": dict or None
+                "children": list of child trees
+        """
+        info = NifiFragmenter.get_fragment_info(root_entity, fragmenter_keylist)
+        child_fragments = []
+        for e in root_entity.children_entities:
+            try:
+                child_info = NifiFragmenter.get_fragment_info(e, fragmenter_keylist)
+                if child_info:
+                    child_fragments.append(e)
+            except (AttributeError, KeyError):
+                pass
+        return {
+            "entity": root_entity,
+            "index": info.get("index"),
+            "merge_params": info.get("merge_params"),
+            "children": [
+                NifiFragmenter.build_fragment_tree_from_children_entities(child, fragmenter_keylist)
+                for child in child_fragments
+            ],
+        }
+
+    @staticmethod
+    def extract_tree_entities(tree):
+        """Flatten a fragment tree into a list of all entities (pre-order).
+
+        Args:
+            tree: Fragment tree node (from build_fragment_tree_from_children_entities).
+
+        Returns:
+            List of entities in pre-order traversal.
+        """
+        entities = [tree["entity"]]
+        for child in tree.get("children", []):
+            entities.extend(NifiFragmenter.extract_tree_entities(child))
+        return entities
+
+    @staticmethod
+    def iterate_fragments_tree(tree, order="post"):
+        """Yield tree nodes in traversal order.
+
+        Args:
+            tree: Fragment tree node (from build_fragment_tree_from_children_entities).
+            order: "post" (children first, default) or "pre" (parent first).
+
+        Yields:
+            dict nodes with "entity", "index", "merge_params", "children" keys.
+        """
+        children = sorted(tree.get("children", []), key=lambda x: x["index"])
+        if order == "pre":
+            yield tree
+        for child in children:
+            yield from NifiFragmenter.iterate_fragments_tree(child, order)
+        if order == "post":
+            yield tree
+
+    @staticmethod
+    def reduce_fragments_tree(tree, leaf_fn, parent_fn):
+        """Bottom-up tree reduction. Processes leaves first, then folds results up.
+
+        Args:
+            tree: Fragment tree node (from build_fragment_tree_from_children_entities).
+            leaf_fn: Callable(node) -> result, called on nodes with no children.
+            parent_fn: Callable(node, child_results) -> result, called on
+                nodes with children. child_results is a list of results from
+                child nodes, sorted by index.
+
+        Returns:
+            The result from the root node.
+        """
+        children = sorted(tree.get("children", []), key=lambda x: x["index"])
+        if not children:
+            return leaf_fn(tree)
+        child_results = [
+            NifiFragmenter.reduce_fragments_tree(child, leaf_fn, parent_fn)
+            for child in children
+        ]
+        return parent_fn(tree, child_results)
+
 
 class NifiEntityBatch(object):
     def __init__(self, entities, config, config_key):
@@ -777,25 +927,68 @@ class NifiEntity(object):
             entity_type = self.record["entity_type"]
             return entity_type == type or type in self.request["ontology_info"]["parents"]
 
-    def is_fragmented(self) -> bool:
-        return bool(self.request["config"].get("fragment", {}).get("fragments_stack"))
-
-    def is_root_fragment(self, entity) -> bool:
-        def _is_sub_fragment_recursive(fragment: dict) -> bool:
-            if not isinstance(fragment, dict):
-                return False
-            if all(k in fragment for k in ["index", "count", "identifier"]):
-                return fragment.get("index", 0) != 0
-            for value in fragment.values():
-                if isinstance(value, dict):
-                    if _is_sub_fragment_recursive(value):
-                        return True
-            return False
-
+    def is_fragmented(self, fragment_name_or_idx=None) -> bool:
+        """Check whether this entity is part of a fragmentation.
+
+        Args:
+            fragment_name_or_idx: If None (default), returns True if the entity
+                belongs to any fragmentation level. If an int, checks whether the
+                fragments_stack has an entry at that index. If a string, checks
+                whether that fragmenter keylist is present in the stack.
+
+        Returns:
+            True if the entity is fragmented (at the specified level, if given).
+        """
+        stack = self.request["config"].get("fragment", {}).get("fragments_stack", [])
+        if fragment_name_or_idx is None:
+            return bool(stack)
+        if isinstance(fragment_name_or_idx, int):
+            return abs(fragment_name_or_idx) <= len(stack)
+        return fragment_name_or_idx in stack
+
+    def is_root_fragment(self, fragment_name_or_idx=-1, recurse=True) -> bool:
+        """Check whether this entity is a root fragment (index == 0).
+
+        Args:
+            fragment_name_or_idx: Which fragmentation level to check. An int
+                indexes into fragments_stack (default -1 = oldest level),
+                a string matches by fragmenter keylist name.
+            recurse: If True (default), check from the starting level towards
+                index 0 (most recent) and return True only if the entity is
+                root at every checked level. If False, only check the single
+                specified level.
+
+        Returns:
+            True if the entity is a root fragment at the specified level(s),
+            or if the entity is not fragmented at all.
+        """
         if not self.is_fragmented():
             return True
-        fragment = entity.request.get("config", {}).get("fragment", {})
-        return not _is_sub_fragment_recursive(fragment)
+        fragments_stack = self.request["config"]["fragment"]["fragments_stack"]
+
+        if isinstance(fragment_name_or_idx, int):
+            try:
+                resolved = (
+                    fragment_name_or_idx
+                    if fragment_name_or_idx >= 0
+                    else len(fragments_stack) + fragment_name_or_idx
+                )
+                key = fragments_stack[resolved]
+            except (IndexError, ValueError):
+                return True
+        else:
+            key = fragment_name_or_idx
+            if key not in fragments_stack:
+                return True
+            resolved = fragments_stack.index(key)
+
+        if recurse:
+            keys_to_check = fragments_stack[:resolved + 1]
+            return all(
+                NifiFragmenter.get_fragment_info(self, k).get("index", 0) == 0
+                for k in keys_to_check
+            )
+        return NifiFragmenter.get_fragment_info(self, key).get("index", 0) == 0
 
     def get_fragment_root_uid(self, fragment_name_or_idx) -> str:
         fragment_config = self.request.get("config", {}).get("fragment", {})
@@ -1216,18 +1409,20 @@ class NifiRoute(Route):
             op = op[0]
             query_params = request.query_params
             processor_suffix = query_params["processor_suffix"]
-            body = await request.json()
             processor_name = "processor." + self.processor_name + "." + op + "." + processor_suffix
             if op not in self.endpoints.keys():
                 raise StarletteHTTPException(403, f"Route {op} is forbidden for NiFi.")
-            task_id = await self.celery_executor.send_task(self.endpoints[op], args=[body, processor_name])
+            task_id = await self.celery_executor.send_task(
+                self.endpoints[op], args=[processor_name], part=request.stream()
+            )
             return task_id
 
     @staticmethod
     def nifi_task(celery_executor, *args, **opts):
         def decorator(func):
             @wraps(func)
-            def nifi_func(task, body, processor_name, *args, **kwargs):
+            def nifi_func(task, body_bytes, processor_name, *args, **kwargs):
+                body = json.loads(body_bytes)
                 with NifiContextManager(body) as nifi_context:
                     entity_batches = nifi_context.receive_input(body, processor_name)
                     entity_batches = func(

{streamlit_octostar_utils-0.4.2.dev25 → streamlit_octostar_utils-0.5.0.dev2}/streamlit_octostar_utils/ontology/inheritance.py

@@ -15,13 +15,13 @@ def get_label_keys(type, ontology):
     return list(label_keys.keys())
 
 def get_label(record, ontology):
-    if record.get("os_materialized_label") not in (None, ""):
-        return str(record.get("os_materialized_label"))
-    if record.get("os_custom_label") not in (None, ""):
-        return str(record.get("os_custom_label"))
+    if record.get("os_entity_label_materialized") not in (None, ""):
+        return str(record.get("os_entity_label_materialized")).strip()
+    if record.get("os_entity_label") not in (None, ""):
+        return str(record.get("os_entity_label")).strip()
     label_keys = get_label_keys(record["entity_type"], ontology)
     fields = [record.get(field) for field in label_keys]
     fields = [f for f in fields if f is not None]
     fields = [str(f) for f in fields if f]
     label = " ".join(fields)
-    return label or None
+    return label.strip() or None