diff --git a/libs/kotaemon/kotaemon/flow/__init__.py b/libs/kotaemon/kotaemon/flow/__init__.py deleted file mode 100644 index d9b9088a..00000000 --- a/libs/kotaemon/kotaemon/flow/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -from .base import Function, Node, Param, SessionFunction, unset -from .safe import load -from .utils.modules import lazy - -__all__ = [ - "SessionFunction", - "Function", - "unset", - "load", - "Param", - "Node", - "lazy", -] diff --git a/libs/kotaemon/kotaemon/flow/backends/__init__.py b/libs/kotaemon/kotaemon/flow/backends/__init__.py deleted file mode 100644 index 522aef86..00000000 --- a/libs/kotaemon/kotaemon/flow/backends/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .base import Backend -from .http_sync import HttpSyncBackend - -__all__ = ["Backend", "HttpSyncBackend"] diff --git a/libs/kotaemon/kotaemon/flow/backends/base.py b/libs/kotaemon/kotaemon/flow/backends/base.py deleted file mode 100644 index 66d11057..00000000 --- a/libs/kotaemon/kotaemon/flow/backends/base.py +++ /dev/null @@ -1,154 +0,0 @@ -import threading -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from ..base import Function - - -class Backend: - """Track the running state of a Function in a thread-safe manner""" - - def __init__(self): - self._ff_in_run: dict[ - int, bool - ] = {} # whether the pipeline is in the run process - self._ff_prefix: dict[int, str] = {} # only root node has prefix as empty "" - self._ff_name: dict[int, str] = {} # only root node has name as empty "" - self._ff_run_id: dict[int, str] = {} # the current run id - self._ff_flow_name: dict[int, str] = {} # the run name - self._func: "Function" - - @property - def in_run(self) -> bool: - """Whether the node is in run process""" - return self._ff_in_run.get(threading.get_ident(), False) - - @in_run.setter - def in_run(self, value: bool): - self._ff_in_run[threading.get_ident()] = value - - @in_run.deleter - def in_run(self): - del self._ff_in_run[threading.get_ident()] - - @property - def prefix(self) -> str: - """Prefix of the execution flow""" - return self._ff_prefix.get(threading.get_ident(), "") - - @prefix.setter - def prefix(self, value: str): - self._ff_prefix[threading.get_ident()] = value - - @prefix.deleter - def prefix(self): - del self._ff_prefix[threading.get_ident()] - - @property - def name(self) -> str: - """Name of the function in the function flow""" - return self._ff_name.get(threading.get_ident(), "") - - @name.setter - def name(self, value: str): - self._ff_name[threading.get_ident()] = value - - @name.deleter - def name(self): - del self._ff_name[threading.get_ident()] - - @property - def run_id(self) -> str: - """Return execution id""" - return self._ff_run_id.get(threading.get_ident(), "") - - @run_id.setter - def run_id(self, value: str): - self._ff_run_id[threading.get_ident()] = value - - @run_id.deleter - def run_id(self): - self._ff_run_id.pop(threading.get_ident(), None) - - @property - def flow_name(self) -> str: - """Name of the execution flow""" - return self._ff_flow_name.get(threading.get_ident(), "") - - @flow_name.setter - def flow_name(self, value: str): - self._ff_flow_name[threading.get_ident()] = value - - @flow_name.deleter - def flow_name(self): - self._ff_flow_name.pop(threading.get_ident(), None) - - @property - def qualidx(self) -> str: - """Return the qualified execution ids for this node""" - - return f"{self.flow_name}|{self.run_id}|{self.abs_path}" - - @property - def parent_qualidx(self) -> str: - """Return the qualified execution ids for the parent node""" - ident = threading.get_ident() - return f"{self.flow_name}|{self.run_id}|{self._ff_prefix.get(ident, '')}" - - @property - def flow_qualidx(self): - """Return the qualified execution flow id""" - return f"{self.flow_name}|{self.run_id}" - - @property - def abs_path(self) -> str: - """Get the node absolute path in execution flow. - - Note: only available during execution - - Path to node is similar to path to folder: - .: root node - .a: to node a - .a.a1.a2: travel from root node to node a2 - - Returns: - str: absolute path of the node - """ - ident = threading.get_ident() - - if self._ff_prefix.get(ident, "") == ".": - return f".{self._ff_name.get(ident, '')}" - - return f"{self._ff_prefix.get(ident, '')}.{self._ff_name.get(ident, '')}" - - def track(self, **kwargs): - """Track node info - - TODO: this operation is heavily depended on _prepare_child.exec, should make - that piece of code relate to this Backend. Otherwise, tough job to maintain 2 - pieces of code that relate to each other in 2 different places that do not - look relate to each other. - """ - ident = threading.get_ident() - - self._ff_in_run[ident] = True - self._ff_prefix[ident] = kwargs.get("prefix", "") - self._ff_name[ident] = kwargs.get("name", "") - self._ff_run_id[ident] = kwargs.get("run_id", "") - self._ff_flow_name[ident] = kwargs.get("flow_name", "") - - def clear(self): - """Clear the tracking info""" - ident = threading.get_ident() - self._ff_in_run.pop(ident, None) - self._ff_prefix.pop(ident, None) - self._ff_name.pop(ident, None) - self._ff_run_id.pop(ident, None) - self._ff_flow_name.pop(ident, None) - - def exec(self, run, args, kwargs): - """Execute the pipeline's run""" - return run(*args, **kwargs) - - def attach(self, func: "Function"): - self._func = func diff --git a/libs/kotaemon/kotaemon/flow/backends/http_sync.py b/libs/kotaemon/kotaemon/flow/backends/http_sync.py deleted file mode 100644 index bb82ed64..00000000 --- a/libs/kotaemon/kotaemon/flow/backends/http_sync.py +++ /dev/null @@ -1,173 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING - -from ..utils.modules import import_modules -from .base import Backend - -if TYPE_CHECKING: - from litestar.connection import Request - from litestar.response import Response - - -def local_only_func_def(func_def: dict) -> dict: - """Trim the function definition to only contain nodes residing on the same machine - - Nodes on the same machine are those that aren't child of any non-Backend node. - """ - - def is_local_node(node: dict) -> bool: - return ( - node["configs"]["default_backend"]["__type__"] == "theflow.backends.Backend" - ) - - def handle_child_nodes(node: dict) -> dict: - if not node["nodes"]: - return node - - if not is_local_node(node): - return { - "function": node["function"], - "params": node["params"], - "nodes": {}, - "configs": node["configs"], - } - - child_nodes = {} - for name, value in node["nodes"].items(): - child_nodes[name] = handle_child_nodes(value) - - return { - "function": node["function"], - "params": node["params"], - "nodes": child_nodes, - "configs": node["configs"], - } - - return handle_child_nodes(func_def) - - -def general_exception_handler(request: Request, exc: Exception) -> Response: - import logging - - from litestar.middleware.exceptions.middleware import create_exception_response - - logging.error("Application error:", exc_info=exc) - return create_exception_response(request, exc) - - -class HttpSyncBackend(Backend): - """Execute node through HTTP synchronous request-response model - - Once triggered, the caller node of this backend will store run's args, kwargs and - node information in the cache with a unique id. It then calls the http endpoint - with the id, wait for the response. The receiver node will fetch the args, kwargs - and node information from the cache with the unique id, execute the run like - normally, and store the result in the cache with the same unique id, and returns - the unique id to the caller node. The caller node will fetch the result from the - cache with the unique id and return the result to the parent pipeline. - """ - - def __init__(self, endpoint: str): - super().__init__() - self._endpoint = endpoint - self._reqm, self._uuidm = import_modules("requests", "uuid") - - def exec(self, run, args, kwargs): - """Execute the pipeline's run remotely""" - if self._func is None: - raise RuntimeError( - "The backend is not attached to a function. If you modify the backend, " - "please make sure to call `self.fl.attach(self)` in the Function where " - "you use this backend." - ) - - # store the information in a cache with specific id - uuid = self._uuidm.uuid4().hex - self._func.context.set( - name=uuid, - value={ - "args": args, - "kwargs": kwargs, - "__fl_runstates__": { - "name": self.name, - "prefix": self.prefix, - "run_id": self.run_id, - "flow_name": self.flow_name, - }, - }, - ) - - # call the http endpoint with the id, wait for the response - resp = self._reqm.get(self._endpoint, params={"id": uuid}) - resp.raise_for_status() - - # fetch the result from the cache - result = self._func.context.get(name=uuid, default=None) - if result is None: - raise RuntimeError( - f"Cannot find the result for {self.name} with id {uuid} in global cache" - ) - result = result["result"] - - return result - - @classmethod - def make(cls, func_def: dict | str, minimal: bool = True): - """Serve the function from a function definition - - Args: - func_def: The Function definition - minimal: Whether to remove any unnecessary function from the pipeline - (those that will be executed remotely). Defaults to True. - - Returns: - A Litestar app that serves the function - """ - from theflow import load - - if isinstance(func_def, str): - import yaml - - with open(func_def) as f: - func_defd: dict = yaml.safe_load(f) - else: - func_defd = func_def - - if minimal: - func_defd = local_only_func_def(func_defd) - - # load the pipeline from the function definition - func_defd["configs"]["default_backend"] = { - "__type__": "theflow.backends.Backend" - } - func = load(func_defd, safe=False) - - # wrap the function call into a litestar app - (ls,) = import_modules("litestar") - - async def predict(id: str) -> str: - # retrieve info - context = func.context.get(name=id, default=None) - if context is None: - raise RuntimeError( - f"Cannot find the context with id {id} in global cache" - ) - result = func( - *context["args"], - **context["kwargs"], - __fl_runstates__=context["__fl_runstates__"], - ) - context["result"] = result - func.context.set(name=id, value=context) - - return id - - app = ls.Litestar( - route_handlers=[ls.get("/")(predict)], - exception_handlers={ - Exception: general_exception_handler, - }, - ) - - return app diff --git a/libs/kotaemon/kotaemon/flow/base.py b/libs/kotaemon/kotaemon/flow/base.py deleted file mode 100644 index 95be3e04..00000000 --- a/libs/kotaemon/kotaemon/flow/base.py +++ /dev/null @@ -1,1757 +0,0 @@ -from __future__ import annotations - -import inspect -import logging -from abc import ABCMeta, abstractmethod -from collections import defaultdict -from copy import deepcopy -from functools import lru_cache -from typing import _GenericAlias # type: ignore -from typing import ( - Any, - Callable, - ForwardRef, - Generic, - TypeVar, - cast, - get_type_hints, - overload, -) - -from typing_extensions import dataclass_transform - -try: - from types import GenericAlias - - _generic_alias_types: tuple = (_GenericAlias, GenericAlias) -except ImportError: - _generic_alias_types = (_GenericAlias,) - -from .config import Config, ConfigProperty, DefaultConfig -from .context import Context -from .debug import likely_cyclic_pipeline -from .exceptions import ( - CyclicDependencyError, - CyclicPipelineError, - InvalidAttrDefinition, -) -from .runs.base import RunTracker -from .settings import settings -from .utils.modules import deserialize, import_dotted_string, lazy, serialize -from .utils.pretties import unflatten_dict -from .utils.typings import ( - input_signature, - is_compatible_with, - is_union_type, - output_signature, -) -from .visualization import trace_pipelne_run - -logger = logging.getLogger(__name__) - - -def is_node_type(annotation) -> bool: - """Return True if the annotation contains Function""" - if is_union_type(annotation): - return any(is_node_type(a) for a in annotation.__args__) - if isinstance(annotation, ForwardRef): - annotation = annotation._evaluate(globals(), locals(), frozenset()) - return is_node_type(annotation) - if isinstance(annotation, _generic_alias_types): - return issubclass(annotation.__origin__, NodeAttr) - if isinstance(annotation, type): - return issubclass(annotation, Function) or issubclass(annotation, NodeAttr) - return False - - -class unset_: - def __bool__(self): - return False - - def __persist_flow__(self): - type_ = f"{self.__module__}.{self.__class__.__qualname__}" - return {"__type__": type_} - - -unset = unset_() - -_Attr = TypeVar("_Attr") -_PAttr = TypeVar("_PAttr") -_NAttr = TypeVar("_NAttr", bound="Function") -_F = TypeVar("_F", bound="Function") - - -class Attr(Generic[_Attr]): - """Decriptor to store function attributes - - Args: - default: default value of the parameter - default_callback: callback function to generate default attribute value. This - callback takes in the Function object and output the default value. - auto_callback: callback function to generate attribute value. This - callback takes in the Function object and output the default value. - cache: if True, the value of the parameter will not be cached and will be - recalculated everytime it is accessed (requires `auto_callback`) - depends_on: if set, the value of the parameter will be calculated from the - values of the depends_on parameters (requires `cache=True`) - help: help message for the attribute - """ - - def __init__( - self, - default: _Attr | lazy[_Attr] | unset_ | type = unset, - *, - default_callback: Callable[[Any], _Attr] | unset_ = unset, - auto_callback: Callable[[Any], _Attr] | unset_ = unset, - cache: bool = False, - depends_on: str | list[str] | None = None, - help: str = "", - **extras, - ): - self._default: _Attr = cast(_Attr, default) - self._default_callback = default_callback - self._auto_callback = auto_callback - self._help = help - self._extras = extras - - self._cache = cache - - if isinstance(depends_on, str): - depends_on = [depends_on] - self._depends_on: list[str] | None = depends_on - self._to_check = depends_on or [] - - self._attrx = self.__class__.__name__ - - def __str__(self): - text = ", ".join( - [ - f"{key}={value}" - for key, value in self.to_dict().items() - if not key.startswith("_") - ] - ) - return f"{self.__class__.__name__}({text})" - - def __repr__(self): - return str(self) - - @overload - def __get__(self, obj: None, _: type[Function] | None) -> Attr: - ... - - @overload - def __get__(self, obj: Function, _: type[Function] | None) -> _Attr: - ... - - def __get__( - self, obj: Function | None, _: type[Function] | None = None - ) -> _Attr | Attr: - """Get the value of the parameter""" - if obj is None: - return self - - if not isinstance(obj, Function): - raise ValueError( - f"`{self.__class__.__name__}` can only be used with Function: " - f"{self._qual_name}" - ) - - if not isinstance(self._auto_callback, unset_): - if self._name in obj.__ff_cyclic_depends__: - raise CyclicDependencyError( - f"Cyclic dependency detected: {self._qual_name}: " - f"{obj.__ff_cyclic_depends__}" - ) - obj.__ff_cyclic_depends__.add(self._name) - value = self._auto_calculate_param(obj) - obj.__ff_cyclic_depends__.remove(self._name) - elif self._name in obj._attrx[self._attrx]: - value = obj._attrx[self._attrx][self._name] - elif self._default != unset: - if isinstance(self._default, lazy): - value = self._default() - else: - value = deepcopy(self._default) - value = cast(_Attr, value) - elif not isinstance(self._default_callback, unset_): - if self._name in obj.__ff_cyclic_depends__: - raise CyclicDependencyError( - f"Cyclic dependency detected: {self._qual_name}: " - f"{obj.__ff_cyclic_depends__}" - ) - obj.__ff_cyclic_depends__.add(self._name) - value = self._default_callback(obj) - obj.__ff_cyclic_depends__.remove(self._name) - else: - return unset # type: ignore - - obj._attrx[self._attrx][self._name] = value - return value - - def __set__(self, obj: Function, value: Any): - if self._auto_callback != unset: - raise ValueError( - f"Cannot set value for auto-calculated {self._attrx}: {self._qual_name}" - ) - obj._attrx[self._attrx][self._name] = value - - def __delete__(self, obj: Function): - if self._auto_callback != unset: - raise ValueError( - f"Cannot delete value for auto-calculated parameter: {self._qual_name}" - ) - - if self._name in obj._attrx[self._attrx]: - del obj._attrx[self._attrx][self._name] - - def __set_name__(self, owner: type, name: str): - self._name = name - self._owner = owner - self._qual_name = ( - f"{self._owner.__module__}.{self._owner.__name__}.{self._name}" - ) - - # validate after receiving the name and type for actionable error message - self._validate_args() - - def __persist_flow__(self): - """Return the state in a way that can be initiated""" - export = { - "__type__": f"{self.__module__}.{self.__class__.__qualname__}", - } - for key, value in self.to_dict().items(): - try: - serialized = serialize(value) - except Exception as e: - type_ = f"{self._owner.__module__}.{self._owner.__qualname__}" - logger.debug(f"{type_}.{self._name}.{key}: {e}... skip") - serialized = serialize(unset) - export[key] = serialized - - return export - - def _auto_calculate_param(self, obj: Function) -> _Attr: - """Calculate the value of the auto-parameter - - Args: - obj: the Function object - - Returns: - the value of the parameter - """ - if isinstance(self._auto_callback, unset_): - raise ValueError( - f"Cannot calculate auto-parameter without auto_callback: " - f"{self._qual_name}" - ) - - if not self._cache: - return self._auto_callback(obj) - - if not self._to_check: - for attr in dir(obj.__class__): - if attr == self._name: - continue - - # TODO: leaky abstraction of child class here. Attr - # shouldn't have knowledge about ParamAttr or NodeAttr - if isinstance(getattr(obj.__class__, attr), ParamAttr): - self._to_check.append(attr) - if isinstance(getattr(obj.__class__, attr), NodeAttr): - self._to_check.append(attr) - - ids = {} - must_recalculate = False - for target in self._to_check: - old_id = obj.__ff_depends__[self._name].get(target, -1) - new_id = id(getattr(obj, target)) - ids[target] = new_id - - if old_id != new_id: - must_recalculate = True - break - - if must_recalculate: - value = self._auto_callback(obj) - # calculate new hash - for target in self._to_check: - id_ = ids[target] if target in ids else id(getattr(obj, target)) - obj.__ff_depends__[self._name][target] = id_ - else: - value = obj._attrx[self._attrx][self._name] - - return value - - def _validate_args(self): - """Validate the __init__ args""" - if ( - sum( - [ - self._default != unset, - self._default_callback != unset, - self._auto_callback != unset, - ] - ) - > 1 - ): - raise InvalidAttrDefinition( - "Only one of `default`, `default_callback`, `auto_callback` can be set:" - f" {self._qual_name}" - ) - - if self._cache and self._auto_callback == unset: - raise InvalidAttrDefinition( - f"`cache=True` only applies for `auto_callback`: {self._qual_name}" - ) - - if self._depends_on and not self._cache: - raise InvalidAttrDefinition( - f"`depends_on` only applies for `cache=True`: {self._qual_name}" - ) - - @classmethod - def default( - cls, - *, - refresh_on_set: bool = False, - strict_type: bool = False, - **kwargs, - ): - """Method decorator to create default callback""" - # TODO: can consider remove .default, since this is not commonly used and can - # cause visual confusion with .auto, which is much more frequently used. - if kwargs.get("help"): - raise ValueError( - "Please set `help` as function docstring when use `default` decorator" - ) - - def inner(func): - help_: str = inspect.getdoc(func) or "" - return cls( - default_callback=func, - help=help_, - refresh_on_set=refresh_on_set, - strict_type=strict_type, - **kwargs, - ) - - return inner - - @classmethod - def auto( - cls, - *, - cache: bool = True, - depends_on: str | list[str] | None = None, - **kwargs, - ): - """Method decorator to create auto callback""" - if kwargs.get("help"): - raise ValueError( - "Please set `help` as function docstring when use `auto` decorator" - ) - - def inner(func): - help_: str = inspect.getdoc(func) or "" - return cls( - auto_callback=func, - help=help_, - depends_on=depends_on, - cache=cache, - **kwargs, - ) - - return inner - - def to_dict(self) -> dict: - """Return the internal state of the Param as a dict""" - return { - "__type__": f"{self.__module__}.{self.__class__.__qualname__}", - "default": self._default, - "default_callback": self._default_callback, - "auto_callback": self._auto_callback, - "help": self._help, - "depends_on": self._depends_on, - "cache": self._cache, - **self._extras, - } - - -class ParamAttr(Attr[_PAttr]): - """Control the behavior of a parameter in a Function - - Args: - default: default value of the parameter - default_callback: callback function to generate default attribute value. This - callback takes in the Function object and output the default value. - auto_callback: callback function to generate attribute value. This - callback takes in the Function object and output the default value. - cache: if True, the value of the parameter will not be cached and will be - recalculated everytime it is accessed (requires `auto_callback`) - depends_on: if set, the value of the parameter will be calculated from the - values of the depends_on parameters (requires `cache=True`) - help: help message for the attribute - refresh_on_set: if True, the original object will be refreshed when it is set - strict_type: if True, the type of the value will be checked when it is set - """ - - def __init__( - self, - default: _PAttr | lazy[_PAttr] | unset_ = unset, - *, - default_callback: Callable[[Any], _PAttr] | unset_ = unset, - auto_callback: Callable[[Any], _PAttr] | unset_ = unset, - cache: bool = False, - depends_on: str | list[str] | None = None, - help: str = "", - refresh_on_set: bool = False, - strict_type: bool = False, - **extras, - ): - super().__init__( - default=default, - default_callback=default_callback, - auto_callback=auto_callback, - cache=cache, - depends_on=depends_on, - help=help, - **extras, - ) - - # param-specific attributes - self._refresh_on_set = refresh_on_set - self._strict_type = strict_type - self._type = None - self._attrx = "ParamAttr" - - def __set__(self, obj: Function, value: Any): - if self._strict_type: - if not isinstance(value, obj.__dict__["__annotations__"][self._name]): - # TODO: more sophisicated type checking (e.g. handle Union, Optional...) - raise ValueError( - f"Value {value} is not of type {type(self._default)} " - f"for parameter {self._name}" - ) - - super().__set__(obj, value) - if self._refresh_on_set: - obj._initialize() - - def __delete__(self, obj: Function): - super().__delete__(obj) - if self._refresh_on_set: - obj._initialize() - - @overload - def __get__(self, obj: None, _: type[Function] | None) -> ParamAttr: - ... - - @overload - def __get__(self, obj: Function, _: type[Function] | None) -> _PAttr: - ... - - def __get__( - self, obj: Function | None, _: type[Function] | None = None - ) -> _PAttr | ParamAttr: - if obj is None: - return self - - value = super().__get__(obj, _) - if value == unset: - if obj.config.params_subscribe and obj.fl.prefix: - context = f"{obj.fl.flow_qualidx}|published_params" - if obj.context.has_context(context): - value = obj.context.get( - name=self._name, - default=unset, - context=context, - ) - - if value != unset: - obj._attrx[self._attrx][self._name] = value - - return value - - @classmethod - def auto( - cls, - *, - refresh_on_set: bool = False, - strict_type: bool = False, - cache: bool = True, - depends_on: str | list[str] | None = None, - **kwargs, - ) -> Callable[[Callable[[Any], _PAttr]], ParamAttr[_PAttr]]: - """Automatically set this method as param's `auto_callback` - - If cached is not enabled, the value will be recalculated everytime. Otherwise, - the value will be recalculated in case all other nodes and parameters that - this node depends on have changed. - - As comparing all values for cache can be expensive, you can list the name - of relevant nodes and params in `depends_on` to limit the comparison to - only these nodes and params. - - Args: - refresh_on_set: whether to refresh the graph when the parameter is set - strict_type: whether to check the type of the value when the param is set - cache: whether to cache the value of the parameter - depends_on: the name of nodes and params that this node depends on - """ - return super().auto( - cache=cache, - depends_on=depends_on, - refresh_on_set=refresh_on_set, - strict_type=strict_type, - **kwargs, - ) - - @classmethod - def default( - cls, - *, - refresh_on_set: bool = False, - strict_type: bool = False, - **kwargs, - ) -> Callable[[Callable[[Any], _PAttr]], ParamAttr[_PAttr]]: - """Automatically set this method as param's `default_callback` - - When this param is accessed while unset, this method will return the default - value - - Args: - refresh_on_set: whether to refresh the graph when the parameter is set - strict_type: whether to check the type of the value when the param is set - """ - return super().default( - refresh_on_set=refresh_on_set, - strict_type=strict_type, - **kwargs, - ) - - def to_dict(self) -> dict: - """Return the internal state of the Param as a dict""" - d = super().to_dict() - d["refresh_on_set"] = self._refresh_on_set - d["strict_type"] = self._strict_type - return d - - -class NodeAttr(Attr[_NAttr]): - """Control the behavior of a node in a Function - - Args: - default: default value of the parameter - default_callback: callback function to generate default attribute value. This - callback takes in the Function object and output the default value. - auto_callback: callback function to generate attribute value. This - callback takes in the Function object and output the default value. - cache: if True, the value of the parameter will not be cached and will be - recalculated everytime it is accessed (requires `auto_callback`) - depends_on: if set, the value of the parameter will be calculated from the - values of the depends_on parameters (requires `cache=True`) - help: help message for the attribute - input: the input signature of the node - output: the output signature of the node - """ - - def __init__( - self, - default: type[_NAttr] | lazy[_NAttr] | unset_ = unset, - *, - default_callback: Callable[[Any], _NAttr] | unset_ = unset, - auto_callback: Callable[[Any], _NAttr] | unset_ = unset, - cache: bool = False, - depends_on: str | list[str] | None = None, - help: str = "", - input: unset_ | dict[str, Any] = unset, - output: Any = unset, - **extras, - ): - if inspect.isclass(default) and issubclass(default, Function): - default = cast(lazy, lazy(default)) - super().__init__( - default=default, - default_callback=default_callback, - auto_callback=auto_callback, - cache=cache, - depends_on=depends_on, - help=help, - **extras, - ) - - self._input: unset_ | dict[str, Any] = input - self._output: Any = output - - has_run_method = callable(getattr(default, "run", None)) - if self._input == unset and has_run_method: - self._input, _, _ = input_signature(default.run) # type: ignore - if self._output == unset and has_run_method: - self._output = output_signature(default.run) # type: ignore - - self._attrx = "NodeAttr" - - @overload - def __get__(self, obj: None, _: type[Function] | None) -> NodeAttr: - ... - - @overload - def __get__(self, obj: Function, _: type[Function] | None) -> _NAttr: - ... - - def __get__(self, obj: Function | None, _: type[Function] | None = None): - if obj is None: - return self - - value = super().__get__(obj, _) - if obj and value: - value = cast(_NAttr, value) - value = obj._prepare_child(value, self._name) - - return value - - @classmethod - def auto( - cls, - *, - input: unset_ | dict[str, Any] = unset, - output: Any = unset, - cache: bool = True, - depends_on: str | list[str] | None = None, - **kwargs, - ) -> Callable[[Callable[[Any], _NAttr]], NodeAttr[_NAttr]]: - """Automatically set this method as param's `auto_callback` - - If cached is not enabled, the value will be recalculated everytime. Otherwise, - the value will be recalculated in case all other nodes and parameters that - this node depends on have changed. - - As comparing all values for cache can be expensive, you can list the name - of relevant nodes and params in `depends_on` to limit the comparison to - only these nodes and params. - - Args: - refresh_on_set: whether to refresh the graph when the parameter is set - strict_type: whether to check the type of the value when the param is set - cache: whether to cache the value of the parameter - depends_on: the name of nodes and params that this node depends on - """ - return super().auto( - cache=cache, - depends_on=depends_on, - input=input, - output=output, - **kwargs, - ) - - @classmethod - def default( - cls, - *, - input: unset_ | dict[str, Any] = unset, - output: Any = unset, - **kwargs, - ) -> Callable[[Callable[[Any], _NAttr]], NodeAttr[_NAttr]]: - """Automatically set this method as param's `default_callback` - - When this param is accessed while unset, this method will return the default - value - - Args: - refresh_on_set: whether to refresh the graph when the parameter is set - strict_type: whether to check the type of the value when the param is set - """ - return super().default( - input=input, - output=output, - **kwargs, - ) - - def to_dict(self) -> dict: - """Return the internal state of a node as dict""" - d = super().to_dict() - d["input"] = self._input - d["output"] = self._output - return d - - -_node_cls: type[NodeAttr] = ( - deserialize(settings.NODE_CLASS, safe=False) - if getattr(settings, "NODE_CLASS", "") - else NodeAttr -) - - -_param_cls: type[ParamAttr] = ( - deserialize(settings.PARAM_CLASS, safe=False) - if getattr(settings, "PARAM_CLASS", "") - else ParamAttr -) - - -class _ParamWrapper: - """Wrapper class to create a Param object - - This way, users can declare Param with both Param(...) and Param.auto(...)""" - - def __call__( - self, - default: _PAttr | lazy[_PAttr] | unset_ = unset, - *, - default_callback: Callable[[Any], _PAttr] | unset_ = unset, - auto_callback: Callable[[Any], _PAttr] | unset_ = unset, - cache: bool = False, - depends_on: str | list[str] | None = None, - help: str = "", - refresh_on_set: bool = False, - strict_type: bool = False, - **extras, - ) -> Any: - return _param_cls( - default=default, - default_callback=default_callback, - auto_callback=auto_callback, - cache=cache, - depends_on=depends_on, - help=help, - refresh_on_set=refresh_on_set, - strict_type=strict_type, - **extras, - ) - - def auto( - self, - *, - refresh_on_set: bool = False, - strict_type: bool = False, - cache: bool = True, - depends_on: str | list[str] | None = None, - **kwargs, - ) -> Callable[[Callable[[Any], _PAttr]], ParamAttr[_PAttr]]: - """Automatically set this method as param's `auto_callback` - - If cached is not enabled, the value will be recalculated everytime. Otherwise, - the value will be recalculated in case all other nodes and parameters that - this node depends on have changed. - - As comparing all values for cache can be expensive, you can list the name - of relevant nodes and params in `depends_on` to limit the comparison to - only these nodes and params. - - Args: - refresh_on_set: whether to refresh the graph when the parameter is set - strict_type: whether to check the type of the value when the param is set - cache: whether to cache the value of the parameter - depends_on: the name of nodes and params that this node depends on - """ - return _param_cls.auto( - cache=cache, - depends_on=depends_on, - refresh_on_set=refresh_on_set, - strict_type=strict_type, - **kwargs, - ) - - -class _NodeWrapper: - """Wrapper class to create a Node object - - This way, users can declare Node with both Node(...) and Node.auto(...)""" - - def __call__( - self, - default: type[_NAttr] | lazy[_NAttr] | unset_ = unset, - *, - default_callback: Callable[[Any], _NAttr] | unset_ = unset, - auto_callback: Callable[[Any], _NAttr] | unset_ = unset, - cache: bool = False, - depends_on: str | list[str] | None = None, - help: str = "", - input: unset_ | dict[str, Any] = unset, - output: Any = unset, - **extras, - ) -> Any: - return _node_cls( - default=default, - default_callback=default_callback, - auto_callback=auto_callback, - cache=cache, - depends_on=depends_on, - help=help, - input=input, - output=output, - **extras, - ) - - def auto( - self, - *, - input: unset_ | dict[str, Any] = unset, - output: Any = unset, - cache: bool = True, - depends_on: str | list[str] | None = None, - **kwargs, - ) -> Callable[[Callable[[Any], _NAttr]], NodeAttr[_NAttr]]: - """Automatically set this method as param's `auto_callback` - - If cached is not enabled, the value will be recalculated everytime. Otherwise, - the value will be recalculated in case all other nodes and parameters that - this node depends on have changed. - - As comparing all values for cache can be expensive, you can list the name - of relevant nodes and params in `depends_on` to limit the comparison to - only these nodes and params. - - Args: - refresh_on_set: whether to refresh the graph when the parameter is set - strict_type: whether to check the type of the value when the param is set - cache: whether to cache the value of the parameter - depends_on: the name of nodes and params that this node depends on - """ - return _node_cls.auto( - cache=cache, - depends_on=depends_on, - input=input, - output=output, - **kwargs, - ) - - -Node = _NodeWrapper() -Param = _ParamWrapper() - - -class MetaFunction(ABCMeta): - def __new__(cls, clsname, bases, attrs): - # Will deprecate in Python 3.13. Now we needs to create the obj to reliably - # obtain type annotations. - _obj: type[Function] = super().__new__( - cls, clsname, bases, attrs # type: ignore - ) - - # Make sure all nodes and params have the Node and Param descriptor - for name, value in get_type_hints(_obj).items(): - if name not in attrs.get("__annotations__", {}): - continue - if name.startswith("_"): - continue - if name in attrs and isinstance(attrs[name], (NodeAttr, ParamAttr)): - continue - desc: NodeAttr | ParamAttr - if is_node_type(value): - desc = _node_cls(default=attrs[name]) if name in attrs else _node_cls() - else: - desc = ( - _param_cls(default=attrs[name]) if name in attrs else _param_cls() - ) - attrs[name] = desc - - try: - obj: type[Function] = super().__new__( - cls, clsname, bases, attrs # type: ignore - ) - except Exception as e: - cause = getattr(e, "__cause__", None) - if isinstance(cause, InvalidAttrDefinition): - raise cause from None - raise e from None - - # Raise invalid nodes and params - for name, value in attrs.items(): - if not isinstance(value, (NodeAttr, ParamAttr)): - continue - if name.startswith("_"): - raise ValueError(f"Node and param name cannot start with _: {name}") - - if name in obj._protected_keywords(): - raise ValueError( - f'"{name}" is a protected keyword, defined by ' - f'"{obj._protected_keywords()[name]}"' - ) - return obj - - -@dataclass_transform( - eq_default=False, - kw_only_default=True, - field_specifiers=(Param, Node), # type: ignore -) -class Function(metaclass=MetaFunction): - """Base class that handle basic logic of a composable component - - Everything is a composable component. Subclass `Function` to define and run your - own flow or component. - - This class: - - Manages input parameters - - Set up _ff_config and _ff_context - - Defines base - - Initialization order: - - initiate the parameters (by default inside __init__) - - initiate the config and context - - initiate the nodes - - Private attributes: - _params: parameters that are passed by the user (e.g. during __init__ or by - setattr) - - Parameters and nodes have to be declared so that the information about the - Function is explicit. - """ - - Config = DefaultConfig - config = ConfigProperty() - - _keywords = [ - "Config", - "apply", - "config", - "context", - "describe", - "dump", - "get_from_path", - "getx", - "is_compatible", - "last_run", - "log_progress", - "missing", - "nodes", - "params", - "run", - "set", - "set_run", - "specs", - "visualize", - "withx", - "fl", - ] - - def __init__(self, _params: dict | None = None, /, **params): - self.last_run: RunTracker - self._track_child: bool = True # flag to track child nodes - self._attrx: dict[str, dict[str, Any]] = { - "NodeAttr": {}, - "ParamAttr": {}, - "AllowExtraParam": {}, - } - self.__ff_cyclic_depends__: set = set() - self.__ff_depends__: dict[str, dict[str, int]] = defaultdict(dict) - self.__ff_run_kwargs__: dict[str, Any] = {} - self._ff_params: list[str] = [] - self._ff_nodes: list[str] = [] - self._ff_config: Config = Config(cls=self.__class__) - self._ff_context: Context | None = None - - # Initialize temporary execution variables - self._variablex() - - # collect - self._ff_params, self._ff_nodes = self._collect_registered_params_and_nodes() - - self._ff_init_called = False - if _params: - self.set(_params, strict=True) - if params: - self.set(params, strict=True) - self._ff_init_called = True - - # collect middleware - middleware_section: str = self.config.middleware_section - middleware_setting = settings.MIDDLEWARE - if middleware_section not in middleware_setting: - raise ValueError( - f'Middleware section "{middleware_section}" not found in settings' - ) - middleware_switches = self.config.middleware_switches - - self._middleware = None - if middlware_cfg := middleware_setting[middleware_section]: - next_call = self._runx - for cls_name in reversed(middlware_cfg): - if not middleware_switches.get(cls_name, True): - continue - cls = import_dotted_string(cls_name, safe=False) - next_call = cls(obj=self, next_call=next_call) - self._middleware = next_call - - if not hasattr(self, "_ff_initializing"): - # TODO: this work better if we formulate config and context as independent - self._initialize() - - def _variablex(self): - """Set temporary variables, only available during execution. Refresh when - execution finishes - """ - self.__ff_run_temp_kwargs__: dict[str, Any] = {} # temp run kwargs - self._ff_childs_called: dict = {} # only available for root - - def __rshift__(self, other: Function) -> Any: - """Return a sequential function""" - if isinstance(other, SequentialFunction): - return SequentialFunction(funcs=[self, *other.funcs]) - if isinstance(self, SequentialFunction): - return SequentialFunction(funcs=[*self.funcs, other]) - if not isinstance(other, Function): - raise ValueError( - f"Can only chain Function, but receive type: {other.__class__.__name__}" - ) - return SequentialFunction(funcs=[self, other]) - - def __floordiv__(self, other: Function) -> Any: - """Return a sequential function""" - if isinstance(other, ConcurrentFunction): - return ConcurrentFunction(funcs=[self, *other.funcs]) - if isinstance(self, ConcurrentFunction): - return ConcurrentFunction(funcs=[*self.funcs, other]) - if not isinstance(other, Function): - raise ValueError( - f"Can only chain Function, but receive type: {other.__class__.__name__}" - ) - return ConcurrentFunction(funcs=[self, other]) - - def _runx(self, *args, **kwargs): - """Subclass to handle pre- and post- run""" - self.fl.in_run = True - return self.run(*args, **kwargs) - - def _post_initialize(self): - pass - - @abstractmethod - def run(self, *args, **kwargs): # type: ignore - raise NotImplementedError(f"Please implement {self.__class__.__name__}.run") - - def __call__(self, *args, **kwargs): - """Run the flow, accepting extra parameters for routing purpose""" - if not hasattr(self, "_ff_initializing"): - self._initialize() - - # might not need to pop __fl_runstates__, because it can be used by other - # operations of the Backend. - _tfrs = kwargs.pop("__fl_runstates__", {}) - if _tfrs: - self.fl.track(**_tfrs) - - if _ff_run_kwargs := kwargs.pop("_ff_run_kwargs", {}): - # TODO: another option is to communicate through context, - # because this is run-time parameter, it will not be persisted in the - # child nodes. - self.set_run(_ff_run_kwargs, temp=True) - - if not self.fl.prefix: # only root node has prefix as empty - # check validity - has_cycle, evidence = likely_cyclic_pipeline(self) - if has_cycle: - raise CyclicPipelineError( - f"Potential cyclic pipeline, please check: {evidence[:5]}" - ) - # administrative setup - self.fl.run_id = self.config.run_id - self.fl.flow_name = self.config.function_name - self.context.create_context(context=self.fl.flow_qualidx) - self.context.set("run_id", self.fl.run_id, context=self.fl.flow_qualidx) - - # publish parameters to the shared cache - if self.config.params_publish: - self.context.create_context( - context=f"{self.fl.flow_qualidx}|published_params", - ) - for k, v in self.params.items(): - self.context.set( - name=k, - value=v, - context=f"{self.fl.flow_qualidx}|published_params", - ) - for k, v in self._attrx["AllowExtraParam"].items(): - self.context.set( - name=k, - value=v, - context=f"{self.fl.flow_qualidx}|published_params", - ) - - self.context.create_context(context=self.fl.qualidx, exist_ok=True) - - # TODO: this will override kwargs passed in __call__. Should follow the - # context-based parameters sharing method - # TODO: this will raise errors in case the users pass in a lot of parameters - # and some of them don't appear in the .run method. - if self.__ff_run_kwargs__: - kwargs.update(self.__ff_run_kwargs__) - - if self.__ff_run_temp_kwargs__: - kwargs.update(self.__ff_run_temp_kwargs__) - - try: - func = self._middleware if self._middleware else self._runx - output = self.fl.exec(func, args, kwargs) - - if not self.fl.prefix: # only root node has prefix as empty - if self.config.params_publish: - self.context.clear( - None, - context=f"{self.fl.flow_qualidx}|published_params", - ) - except Exception as e: - raise e from None - finally: - self._variablex() - self.fl.clear() - - return output - - def __repr__(self): - kwargs = ", ".join( - [f"{key}={repr(getattr(self, key, None))}" for key in self._ff_params] - ) - return f"{self.__class__.__name__}({kwargs})" - - def __str__(self): - """Represent hierarchical structure of the Function""" - if self._ff_nodes: - kwargs = [] - for key in reversed(self._ff_nodes): - value = str(getattr(self, key, None)) - value = value.replace("\n", "\n ") - kwargs.append(f" ({key}): {value}") - kwargs_repr = "\n".join(kwargs) - return f"{self.__class__.__name__}(\n{kwargs_repr}\n)" - - kwargs = [] - for key in self._ff_params: - value = str(getattr(self, key, None)) - if len(value) > 20: - value = f"{value[:15]}..." - kwargs.append(f"{key}={value}") - kwargs_repr = ", ".join(kwargs) - return f"{self.__class__.__name__}({kwargs_repr})" - - def _get_context(self) -> Context | None: - return self._ff_context - - def _set_context(self, context: Context) -> None: - self._ff_context = context - - def _del_context(self) -> None: - del self._ff_context - - context = property(_get_context, _set_context, _del_context) - - @property - def nodes(self) -> list[str]: - return self._ff_nodes - - @property - def params(self) -> dict[str, Any]: - params = {} - for key in self._ff_params: - try: - params[key] = getattr(self, key) - except Exception: - params[key] = None - return params - - def __setattr__(self, name: str, value: Any) -> None: - if name.startswith("_"): - return super().__setattr__(name, value) - - if name in self._ff_nodes: - if not isinstance(value, Function): - value = self._convert_to_function(value) - elif name not in self._ff_params and name not in self._protected_keywords(): - if self.config.allow_extra: - self._attrx["AllowExtraParam"][name] = value - else: - raise AttributeError( - f"Attribute {name} is not defined in {self.__class__.__name__}" - ) - - return super().__setattr__(name, value) - - def _initialize(self): - if self._ff_context is None: - self._ff_context = deserialize(settings.CONTEXT, safe=False) - - # Initialize the backend - self.fl = deserialize(self.config.default_backend, safe=False) - self.fl.attach(self) - - if not hasattr(self, "_ff_init_called"): - raise RuntimeError( - "Please call super().__init__(**params) in your __init__ method" - ) - - self._ff_initializing = True - self._post_initialize() - self._ff_initializing = False - - @classmethod - def _collect_registered_params_and_nodes(cls) -> tuple[list[str], list[str]]: - """Return the list of all params and nodes registered in the Function - - Returns: - tuple[list[str], list[str]]: params, nodes - """ - params, nodes = [], [] - - for attr in dir(cls): - if isinstance(getattr(cls, attr), NodeAttr): - nodes.append(attr) - elif isinstance(getattr(cls, attr), ParamAttr): - params.append(attr) - - return list(sorted(set(params))), list(sorted(set(nodes))) - - @classmethod - @lru_cache - def _protected_keywords(cls) -> dict[str, type]: - """Return the protected keywords and the class that defines each of them - - This method will concatenate the `_keywords` of all classes in the mro. - """ - keywords = {} - for each_cls in cls.mro(): - for keyword in getattr(each_cls, "__dict__", {}).get("_keywords", []): - if keyword in keywords: - continue - keywords[keyword] = each_cls - return keywords - - def _convert_to_function(self, value) -> Function: - """Convert a vanilla object into a function. - - If the value is None, return as it, as likely the user wants to disable the - node. - - Args: - value: the object to be converted - - Returns: - Function: the converted object (or as it) - """ - if value is None: - return value # type: ignore - - return ProxyFunction(ff_original_obj=value) - - def _prepare_child(self, child: _F, name: str) -> _F: - """Prepare child node to enable tracking and routing""" - if not hasattr(self, "fl"): - return child - - if not self.fl.in_run: - return child - - if not self._track_child: - return child - - def exec(*args, **kwargs): - __fl_runstates__ = { - "prefix": self.fl.abs_path, - "name": ( - name - if name not in self._ff_childs_called - else f"{name}[{self._ff_childs_called[name]}]" - ), - "run_id": self.fl.run_id, - "flow_name": self.fl.flow_name, - } - self._ff_childs_called[name] = self._ff_childs_called.get(name, 0) + 1 - return child(*args, **kwargs, __fl_runstates__=__fl_runstates__) - - return exec # type: ignore - - @classmethod - def visualize(cls): - # 1 re-initialize the flow with different mode - # 2 check the argument defintion passed into `run` - # 3 run the flow with the fake argument - # 4 track the graph - return trace_pipelne_run(cls) - - @classmethod - def withx(cls, **kwargs) -> Any: # hacky way to make mypy happy - """Return lazy init object that has the supplied params as default - - Args: - kwargs: the keywords and params to be set as default - - Returns: - A new Function with the supplied keywords and params set as default - """ - return lazy(cls, **kwargs) - - def apply(self, fn: Callable): - """Apply a function recursively to all nodes in a pipeline""" - for node in self._ff_nodes: - getattr(self, node).apply(fn) - fn(self) - return self - - def set(self, kwargs: dict, strict: bool = False): - """Set the keyword arguments in the function""" - kwargs = unflatten_dict(kwargs) - for name, value in kwargs.items(): - name = name.strip(".") - if name in self._ff_nodes and isinstance(value, dict): - getattr(self, name).set(value, strict=strict) - else: - try: - setattr(self, name, value) - except Exception as e: - if strict: - raise e from None - - def set_run(self, kwargs: dict, temp=False): - """Set run keyword arguments - - # TODO: should utilize context or queue to store these parameters, since the - # same node object can be used in multiple pipelines, or also can be used - # multiple times in the same pipeline. Hence, setting and clearing the - # internal attributes will override the internal attributes of the same node - # used in other pipelines / other parts of the pipeline. - - # It's tolerable for `set` though, because `set` deals with initialization - # parameters, which will need to be the same. - - # Another approach is to force clone of node in a pipeline, so that changing - # internal attribute of the node will never affect "that node" in other - # pipelines. - - # Nevertheless, a good abstraction of the context will provide much - # versatility to the users. - """ - kwargs = unflatten_dict(kwargs) - for name, value in kwargs.items(): - name = name.strip(".") - if name in self._ff_nodes and isinstance(value, dict): - getattr(self, name).set_run(value, temp=temp) - else: - if temp: - self.__ff_run_temp_kwargs__[name] = value - else: - self.__ff_run_kwargs__[name] = value - - @classmethod - def describe(cls) -> dict: - """Describe the flow - - TODO: export the route of the flow as well - """ - params, nodes = {}, {} - - for attr in dir(cls): - attr_value = getattr(cls, attr) - if isinstance(attr_value, NodeAttr): - value = attr_value.__persist_flow__() - if isinstance(attr_value._default, lazy) and issubclass( - attr_value._default._cls, Function - ): - value[ - "default" - ] = attr_value._default._cls.describe() # type:ignore - value["default_kwargs"] = { - key: value - for key, value in attr_value._default._params.items() - if not isinstance(value, lazy) - } - nodes[attr] = value - elif isinstance(attr_value, ParamAttr): - attr_val = attr_value.__persist_flow__() - attr_val["type"] = repr( - attr_value._owner.__annotations__.get(attr_value._name, Any) - ) - params[attr] = attr_val - - return { - "type": f"{cls.__module__}.{cls.__qualname__}", - "params": params, - "nodes": nodes, - } - - def dump(self, ignore_auto: bool = True, strict: bool = True) -> dict: - """Export the flow to a dictionary - - This method largely follows `theflow.utils.modules.serialize`, with the added - options to modify the behavior of serialization. - - Args: - ignore_auto: whether to ignore params and nodes that depend on others - strict: whether to raise error if any param or node cannot be serialized - """ - nodes: dict = {} - for node in self._ff_nodes: - try: - obj: Function = self.get_from_path(node) - if self.specs(node).get("auto_callback", unset) and ignore_auto: - continue - nodes[node] = obj.dump(ignore_auto=ignore_auto, strict=strict) - except Exception as e: - if strict: - raise e from None - logger.warn(e) - nodes[node] = None - - params = {} - for name, value in self.params.items(): - if self.specs(name).get("auto_callback", []) and ignore_auto: - continue - try: - params[name] = serialize(value) - except ValueError as e: - if strict: - raise e from None - logger.warn(e) - - return { - "function": f"{self.__module__}.{self.__class__.__qualname__}", - "nodes": nodes, - "params": params, - "configs": self.config.dump(), - } - - def specs(self, path: str) -> dict: - """Get specification about a param or a node - - Args: - path: the path to the node or param (.) delimited - - Returns: - the specification of the param or node - """ - path = path.strip(".") - - if "." in path: - module, subpath = path.split(".", 1) - return getattr(self, module).specs(subpath) - - definition = getattr(self.__class__, path) - if not isinstance(definition, (ParamAttr, NodeAttr)): - raise ValueError(f"{path} is not a param or a node") - - return definition.to_dict() - - def getx(self, path: str) -> Any: - """Get the Function node or param based on path""" - path = path.strip(".") - if "." in path: - module, subpath = path.split(".", 1) - return getattr(self, module).getx(subpath) - - return getattr(self, path) - - def missing(self) -> dict[str, list[str]]: - """Return the list of missing params and nodes""" - params, nodes = [], [] - for attr in self._ff_params: - if getattr(self.__class__, attr)._depends_on: - continue - try: - getattr(self, attr) - except Exception: - params.append(attr) - - for attr in self._ff_nodes: - if getattr(self.__class__, attr)._depends_on: - continue - try: - child = getattr(self, attr) - missings = child.missing() - for each in missings["params"]: - params.append(f"{attr}.{each}") - for each in missings["nodes"]: - nodes.append(f"{attr}.{each}") - except Exception: - nodes.append(attr) - - return {"params": params, "nodes": nodes} - - def get_from_path(self, path) -> Any: - """Get a node or param by path, with tracking disabled - - Args: - path: the path to the node or param (.) delimited - - Returns: - Node or param, depending on the path - """ - self._track_child = False - path = path.strip(".") - - if "." in path: - module, subpath = path.split(".", 1) - obj = getattr(self, module) - self._track_child = True - return obj.get_from_path(subpath) - - obj = getattr(self, path) - self._track_child = True - return obj - - def is_compatible(self, path, obj) -> bool: - """Check if the interface of a sample is compatible with the declared interface - - Args: - path: the path to the node or param (.) delimited - obj: the class or object to be checked - - Returns: - True if compatible, False otherwise - """ - specs = self.specs(path) - func = obj - if isinstance(obj, Function): - func = obj.run - elif isinstance(obj, type) and issubclass(obj, Function): - func = obj.run - - if specs["__type__"] == "theflow.base.ParamAttr": - return isinstance(obj, specs["type"]) - elif specs["__type__"] == "theflow.base.NodeAttr": - reference_input = specs["input"] - reference_output = specs["output"] - target_input, _, _ = input_signature(func) - target_output = output_signature(func) - - ok_input, ok_output = False, False - if reference_input == unset: - ok_input = True - else: - for name, annot in reference_input.items(): - if name not in target_input: - ok_input = False - break - ok_input = is_compatible_with(target_input[name], annot) - if not ok_input: - break - - if reference_output == unset: - ok_output = True - else: - ok_output = is_compatible_with(target_output, reference_output) - return ok_input and ok_output - - raise ValueError(f"{path} is not a param or a node") - - def log_progress(self, name: str | None = None, **kwargs): - """Log the progress to the name""" - if name is None: - name = self.fl.abs_path - - run_tracker = RunTracker(self) - run_tracker.log_progress(name, **kwargs) - - def __persist_flow__(self) -> dict: - """Persist function into a re-constructable JSON-serializable dictionary""" - export: dict = { - "__type__": f"{self.__module__}.{self.__class__.__qualname__}", - } - - for name, value in self.params.items(): - # ignore auto parameter - if self.specs(name).get("auto_callback", []): - continue - try: - export[name] = serialize(value) - except ValueError as e: - logger.warn(e) - continue - - for name in self._ff_nodes: - if self.specs(name).get("auto_callback", []): - continue - node = self.get_from_path(name).__persist_flow__() - export[name] = node - - return export - - -class SessionFunction(Function): - """Handle sesssion""" - - def start_session(self): - if not hasattr(self, "_ff_initializing"): - self._initialize() - - if not self.fl.prefix: # only root node has prefix as empty - # administrative setup - self.fl.run_id = self.config.run_id - self.fl.flow_name = self.config.function_name - self.context.create_context(context=self.fl.flow_qualidx) - self.context.set("run_id", self.fl.run_id, context=self.fl.flow_qualidx) - - self.context.create_context(context=self.fl.qualidx, exist_ok=True) - - def __call__(self, *args, **kwargs): - if _ff_run_kwargs := kwargs.pop("_ff_run_kwargs", {}): - # TODO: another option is to communicate through context, - # because this is run-time parameter, it will not be persisted in the - # child nodes. - self.set_run(_ff_run_kwargs, temp=True) - - if self.__ff_run_kwargs__: - kwargs.update(self.__ff_run_kwargs__) - - if self.__ff_run_temp_kwargs__: - kwargs.update(self.__ff_run_temp_kwargs__) - - output = ( - self._middleware(*args, **kwargs) - if self._middleware - else self._runx(*args, **kwargs) - ) - - return output - - def end_session(self): - self._variablex() - self.fl.clear() - - -class ProxyFunction(Function): - """Wrap an object to be a step. - - `ProxyFunction` demonstrates the same behavior as `Step`. The only difference is - that `ProxyFunction` doesn't know how the object will be called (e.g. `__call__` - or any methods) so it lazily exposes the methods when called. - - Cannot use this class directly because Function reserves some common _keywords - that can conflict with original object. - - Raise ValueError in case of conflict. - """ - - ff_original_obj: Callable - - def __init__(self, **params): - super().__init__(**params) - if isinstance(self.ff_original_obj, ProxyFunction): - raise ValueError( - "Unnecessary to wrap a ProxyFunction object with ProxyFunction" - ) - - def _create_callable(self, callable_obj): - middleware_section: str = self.config.middleware_section - middleware_setting = settings.MIDDLEWARE - if middleware_section not in middleware_setting: - raise ValueError( - f'Middleware section "{middleware_section}" not found in settings' - ) - middleware_switches = self.config.middleware_switches - - if middlware_cfg := middleware_setting[middleware_section]: - next_call = callable_obj - for cls_name in reversed(middlware_cfg): - if not middleware_switches.get(cls_name, True): - continue - cls = import_dotted_string(cls_name, safe=False) - next_call = cls(obj=self, next_call=next_call) - callable_obj = next_call - - def wrapper(*args, **kwargs): - if not hasattr(self, "_ff_initializing"): - self._initialize() - - _tfrs = kwargs.pop("__fl_runstates__", {}) - if _tfrs: - self.fl.track(**_tfrs) - - try: - output = callable_obj(*args, **kwargs) - except Exception as e: - raise e from None - finally: - self.fl.clear() - - return output - - return wrapper - - def __call__(self, *args, **kwargs): - if self._ff_context is None: - self._ff_context = deserialize(settings.CONTEXT, safe=False) - - return self._create_callable(getattr(self.ff_original_obj, "__call__"))( - *args, **kwargs - ) - - def __getattr__(self, name): - if "ff_original_obj" not in self._ff_params: - raise AttributeError( - f"{self.__class__.__qualname__} object has no attribute {name}" - ) - - if self._ff_context is None: - self._ff_context = deserialize(settings.CONTEXT, safe=False) - - attr = getattr(self.ff_original_obj, name) - if callable(attr): - attr = self._create_callable(attr) - return attr - - def run(self, *args, **kwargs) -> Any: - if hasattr(self.ff_original_obj, "run") and callable(self._ff_original_obj.run): - return self.ff_original_obj.run(*args, **kwargs) - raise NotImplementedError(f"{self.ff_original_obj}.run doesn't exist") - - -class SequentialFunction(Function): - """Sequential functions""" - - funcs: list[Function] = [] - - def __len__(self): - return len(self.funcs) - - def __getitem__(self, idx): - return self.funcs[idx] - - def __str__(self): - """Represent hierarchical structure of the Function""" - kwargs = [] - for idx, func in enumerate(self.funcs): - value = str(func() if isinstance(func, lazy) else func) - value = value.replace("\n", "\n ") - kwargs.append(f" ({idx}): {value}") - kwargs_repr = "\n".join(kwargs) - return f"{self.__class__.__name__}(\n{kwargs_repr}\n)" - - def run(self, *arg, **kwargs): - out = arg - for idx, func in enumerate(self.funcs): - func_: Function = func() if isinstance(func, lazy) else func - func_ = self._prepare_child(func_, f"func{idx}_{func_.__class__.__name__}") - out = func_(*arg, **kwargs) - arg = out if len(arg) > 1 else (out,) - return out - - -class ConcurrentFunction(Function): - """Run functions concurrently""" - - funcs: list[Function] = [] - - def __len__(self): - return len(self.funcs) - - def __getitem__(self, idx): - return self.funcs[idx] - - def __str__(self): - """Represent hierarchical structure of the Function""" - kwargs = [] - for idx, func in enumerate(self.funcs): - value = str(func) - value = value.replace("\n", "\n ") - kwargs.append(f" ({idx}): {value}") - kwargs_repr = "\n".join(kwargs) - return f"{self.__class__.__name__}(\n{kwargs_repr}\n)" - - def run(self, arg): - output = [] - for idx, func in enumerate(self.funcs): - func_: Function = func() if isinstance(func, lazy) else func - func_ = self._prepare_child(func_, f"func{idx}_{func_.__class__.__name__}") - output.append(func_(arg)) - return output diff --git a/libs/kotaemon/kotaemon/flow/cache/__init__.py b/libs/kotaemon/kotaemon/flow/cache/__init__.py deleted file mode 100644 index f4d874b1..00000000 --- a/libs/kotaemon/kotaemon/flow/cache/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -from .base import BaseCache -from .filebased import FileCache -from .memcached import PyMemcacheCache -from .memory import MemoryCache - -__all__ = ["BaseCache", "MemoryCache", "FileCache", "PyMemcacheCache"] diff --git a/libs/kotaemon/kotaemon/flow/cache/base.py b/libs/kotaemon/kotaemon/flow/cache/base.py deleted file mode 100644 index 7250e7ed..00000000 --- a/libs/kotaemon/kotaemon/flow/cache/base.py +++ /dev/null @@ -1,159 +0,0 @@ -import abc -from typing import Any, Callable, Optional - - -class BaseCache(abc.ABC): - @abc.abstractmethod - def __init__(self, *args, **kwargs): - """Initialize the cache""" - ... - - @abc.abstractmethod - def add(self, key: str, value: Any, timeout: Optional[int] = None) -> None: - """Add a key/value pair to the cache if it doesn't already exist - - If the key already exists, this method will do nothing. - - Args: - key: the name of the key to add - value: the value to add - timeout: the number of seconds to keep the key in the cache - """ - ... - - @abc.abstractmethod - def get(self, key: str, default: Any = None) -> Any: - """Get a value from the cache - - Args: - key: the name of the key to get - default: the value to return if the key doesn't exist - - Returns: - The value of the key, or the default value if the key doesn't exist - """ - ... - - @abc.abstractmethod - def delete(self, key: str) -> None: - """Delete a key from the cache - - Args: - key: the name of the key to delete - """ - ... - - @abc.abstractmethod - def set(self, key: str, value: Any, timeout: Optional[int] = None) -> None: - """Set a key/value pair in the cache - - If the key already exists, its value will be overwritten. - - Args: - key: the name of the key to set - value: the value to set - timeout: the number of seconds to keep the key in the cache - """ - ... - - @abc.abstractmethod - def touch(self, key: str, timeout: Optional[int] = None) -> None: - """Update the timeout for a key - - Args: - key: the name of the key to update - timeout: the number of seconds to keep the key in the cache - """ - ... - - @abc.abstractmethod - def clear(self) -> None: - """Clear all keys from the cache""" - ... - - @abc.abstractmethod - def close(self) -> None: - """Close the cache""" - ... - - @abc.abstractmethod - def incr(self, key: str, delta: int = 1) -> int: - """Increment a key's value - - Args: - key: the name of the key to increment - delta: the amount to increment the key's value by - """ - ... - - @abc.abstractmethod - def decr(self, key: str, delta: int = 1) -> int: - """Decrement a key's value - - Args: - key: the name of the key to decrement - delta: the amount to decrement the key's value by - - Returns: - The new value of the key - """ - ... - - @abc.abstractmethod - def __contains__(self, key: str) -> bool: - """Check if a key exists in the cache - - Args: - key: the name of the key to check - - Returns: - True if the key exists, False otherwise - """ - ... - - @abc.abstractmethod - def __getitem__(self, key: str) -> Any: - """Get a value from the cache - - Args: - key: the name of the key to get - - Returns: - The value of the key - """ - ... - - @abc.abstractmethod - def __setitem__(self, key: str, value: Any) -> None: - """Set a key/value pair in the cache - - If the key already exists, its value will be overwritten. - - Args: - key: the name of the key to set - value: the value to set - """ - ... - - @abc.abstractmethod - def __delitem__(self, key: str) -> None: - """Delete a key from the cache - - Args: - key: the name of the key to delete - """ - ... - - @abc.abstractmethod - def get_then_set(self, key: str, func: Callable[[Any], Any], default: Any = None): - """Get a value from the cache, and then set the value, avoiding race conditions - - Args: - key: the name of the key to get - func: a function to call to get the updated value - default: the value to return if the key doesn't exist - - Returns: - The value of the key, updated to cache - """ - ... diff --git a/libs/kotaemon/kotaemon/flow/cache/filebased.py b/libs/kotaemon/kotaemon/flow/cache/filebased.py deleted file mode 100644 index 2e137598..00000000 --- a/libs/kotaemon/kotaemon/flow/cache/filebased.py +++ /dev/null @@ -1,92 +0,0 @@ -import logging - -from .base import BaseCache - -logger = logging.getLogger(__name__) - -_local_caches = {} - - -class FileCache(BaseCache): - """A file-based cache - - A file-based case that persist the cache in to a directory on disk. Suitable for - different runs in different times to share the same cache. - """ - - def __init__(self, path): - try: - import diskcache # type: ignore - except ImportError: - raise ImportError( - "The diskcache package is required to use the FileBasedCache. " - "Please run: pip install diskcache" - ) - - path = str(path) - if path not in _local_caches: - _local_caches[path] = diskcache.Cache(path) - - self._cache = _local_caches[path] - self._lock = diskcache.RLock(self._cache, "__lock__") - - def add(self, key, value, timeout=None): - return self._cache.add(key, value, expire=timeout) - - def get(self, key, default=None): - return self._cache.get(key, default=default) - - def delete(self, key): - self._cache.delete(key) - - def set(self, key, value, timeout=None): - self._cache.set(key, value, expire=timeout) - - def touch(self, key, timeout=None): - self._cache.touch(key, expire=timeout) - - def clear(self): - self._cache.clear() - - def close(self): - self._cache.close() - - def incr(self, key, delta=1): - return self._cache.incr(key, delta) - - def decr(self, key, delta=1): - return self._cache.decr(key, delta) - - def __contains__(self, key): - return key in self._cache - - def __getitem__(self, key): - return self._cache[key] - - def __setitem__(self, key, value): - self._cache[key] = value - - def __delitem__(self, key): - del self._cache[key] - - @property - def lock(self): - return self._lock - - def __getstate__(self): - state = self.__dict__.copy() - state.pop("_lock") - return state - - def __setstate__(self, state): - import diskcache # type: ignore - - self.__dict__.update(state) - self._lock = diskcache.RLock(self._cache, "__lock__") - - def get_then_set(self, key, func, default=None): - with self._lock: - value = self._cache.get(key, default) - value = func(value) - self._cache.set(key, value) - return value diff --git a/libs/kotaemon/kotaemon/flow/cache/memcached.py b/libs/kotaemon/kotaemon/flow/cache/memcached.py deleted file mode 100644 index 1628f178..00000000 --- a/libs/kotaemon/kotaemon/flow/cache/memcached.py +++ /dev/null @@ -1,90 +0,0 @@ -import threading -from typing import Any, Callable, Optional - -from .base import BaseCache - - -class PyMemcacheCache(BaseCache): - def __init__(self, servers, **kwargs): - try: - import pymemcache # type: ignore - import pymemcache.serde # type: ignore # noqa: F401 - except ImportError: - raise ImportError( - "The pymemcache package is required to use the PyMemcacheCache. " - "Please run: pip install pymemcache" - ) - - self._servers = servers - self._kwargs = kwargs - self._caches = {} - - @property - def _cache(self): - import pymemcache - import pymemcache.serde - - ident = threading.get_ident() - if ident not in self._caches: - self._caches[ident] = pymemcache.Client( - self._servers, serde=pymemcache.serde.pickle_serde, **self._kwargs - ) - - return self._caches[ident] - - def add(self, key: str, value: Any, timeout: Optional[int] = None) -> None: - self._cache.add(key, value, expire=timeout or 0) - - def get(self, key: str, default: Any = None): - return self._cache.get(key, default) - - def delete(self, key: str) -> None: - self._cache.delete(key) - - def set(self, key: str, value: Any, timeout: Optional[int] = None) -> None: - self._cache.set(key, value, expire=timeout or 0) - - def touch(self, key, timeout: Optional[int] = None) -> None: - self._cache.touch(key, expire=timeout or 0) - - def clear(self) -> None: - self._cache.flush_all() - - def close(self) -> None: - self._cache.close() - - def incr(self, key, delta: int = 1) -> int: - return self._cache.incr(key, delta) # type: ignore - - def decr(self, key, delta: int = 1) -> int: - return self._cache.decr(key, delta) # type: ignore - - def __contains__(self, key: str) -> bool: - return self._cache.get(key) is not None - - def __getitem__(self, key: str): - return self._cache.get(key) - - def __setitem__(self, key: str, value: Any): - self._cache.set(key, value) - - def __delitem__(self, key: str): - self._cache.delete(key) - - def get_then_set(self, key: str, func: Callable[[Any], Any], default: Any = None): - for _ in range(20): - value, cas = self._cache.gets(key, default) - value = func(value) - if self._cache.cas(key, value, cas): - return value - - raise RuntimeError("Key changes very frequently, please try again later") - - def __getstate__(self): - state = self.__dict__.copy() - del state["_caches"] - return state - - def __setstate__(self, state): - state["_caches"] = {} - self.__dict__.update(state) diff --git a/libs/kotaemon/kotaemon/flow/cache/memory.py b/libs/kotaemon/kotaemon/flow/cache/memory.py deleted file mode 100644 index 9d6e2990..00000000 --- a/libs/kotaemon/kotaemon/flow/cache/memory.py +++ /dev/null @@ -1,102 +0,0 @@ -import logging -import multiprocessing -import multiprocessing.managers -import multiprocessing.synchronize -from typing import Any, Dict, Optional - -from .base import BaseCache - -logger = logging.getLogger(__name__) -MANAGERS: Dict[str, multiprocessing.managers.SyncManager] = {} -LOCKS: Dict[str, multiprocessing.synchronize.RLock] = {} -MSG_STORE: Dict[str, multiprocessing.managers.DictProxy] = {} - - -class MemoryCache(BaseCache): - """A memory-based cache - - This cache is quick to spin up and will terminate at the end of the process. It is - suitable for testing and do small runs. It makes use of multiprocessing module to - allow multiple processes to share the same cache. - - Args: - uid: a unique identifier for the cache. If not provided, a fixed value "" will - be used. - """ - - def __init__(self, uid: str = ""): - """Initialize the object""" - self.uid = uid - if uid not in MANAGERS: - MANAGERS[uid] = multiprocessing.Manager() - - if uid not in LOCKS: - LOCKS[uid] = multiprocessing.RLock() - - if uid not in MSG_STORE: - MSG_STORE[uid] = MANAGERS[uid].dict() - - def add(self, key: str, value: Any, timeout: Optional[int] = None) -> None: - if timeout is not None: - logger.info(f"Add: Timeout value ({timeout}) is ignored for memory cache") - with LOCKS[self.uid]: - if key not in MSG_STORE[self.uid]: - MSG_STORE[self.uid][key] = value - - def get(self, key: str, default: Any = None) -> Any: - with LOCKS[self.uid]: - return MSG_STORE[self.uid].get(key, default) - - def delete(self, key: str) -> None: - with LOCKS[self.uid]: - if key in MSG_STORE[self.uid]: - del MSG_STORE[self.uid][key] - - def set(self, key: str, value: Any, timeout: Optional[int] = None) -> None: - if timeout is not None: - logger.info(f"Set: Timeout value ({timeout}) is ignored for memory cache") - with LOCKS[self.uid]: - MSG_STORE[self.uid][key] = value - - def touch(self, key: str, timeout: Optional[int] = None) -> None: - if timeout is not None: - logger.info( - f"Touch ({key}): Timeout value ({timeout}) is ignored for memory cache" - ) - - def clear(self) -> None: - with LOCKS[self.uid]: - MSG_STORE[self.uid].clear() - - def close(self) -> None: - ... - - def incr(self, key: str, delta: int = 1) -> int: - with LOCKS[self.uid]: - if key not in MSG_STORE[self.uid]: - MSG_STORE[self.uid][key] = 0 - MSG_STORE[self.uid][key] += delta - return MSG_STORE[self.uid][key] - - def decr(self, key: str, delta: int = 1) -> int: - return self.incr(key, -delta) - - def __contains__(self, key: str) -> bool: - with LOCKS[self.uid]: - return key in MSG_STORE[self.uid] - - def __getitem__(self, key: str) -> Any: - with LOCKS[self.uid]: - return MSG_STORE[self.uid][key] - - def __setitem__(self, key: str, value: Any) -> None: - with LOCKS[self.uid]: - MSG_STORE[self.uid][key] = value - - def __delitem__(self, key: str) -> None: - with LOCKS[self.uid]: - del MSG_STORE[self.uid][key] - - @property - def lock(self): - return LOCKS[self.uid] diff --git a/libs/kotaemon/kotaemon/flow/callbacks.py b/libs/kotaemon/kotaemon/flow/callbacks.py deleted file mode 100644 index 163bf04d..00000000 --- a/libs/kotaemon/kotaemon/flow/callbacks.py +++ /dev/null @@ -1,15 +0,0 @@ -import time - -from theflow.base import Function - - -def run_id__timestamp(obj: Function) -> str: - return str(time.time()).replace(".", "") - - -def store_result__pipeline_name(obj: Function) -> str: - return f"{obj.__module__}.{obj.__class__.__qualname__}" - - -def function_name__class_name(obj: Function) -> str: - return f"{obj.__class__.__module__}.{obj.__class__.__qualname__}" diff --git a/libs/kotaemon/kotaemon/flow/cli.py b/libs/kotaemon/kotaemon/flow/cli.py deleted file mode 100644 index e69de29b..00000000 diff --git a/libs/kotaemon/kotaemon/flow/config.py b/libs/kotaemon/kotaemon/flow/config.py deleted file mode 100644 index b7b107b7..00000000 --- a/libs/kotaemon/kotaemon/flow/config.py +++ /dev/null @@ -1,190 +0,0 @@ -from typing import TYPE_CHECKING, Any, Optional, Type, Union - -import yaml - -if TYPE_CHECKING: - from .base import Function - -from .settings import settings -from .utils.modules import import_dotted_string - -# configs that are dictionary that will be aggregated from parent to child -# classes, rather than being overwritten -_aggregated_dict = {"middleware_switches"} - - -class DefaultConfig: - # skip storing the result if set to None - store_result = "{{ theflow.callbacks.store_result__pipeline_name }}" - run_id = "{{ theflow.callbacks.run_id__timestamp }}" - function_name = "{{ theflow.callbacks.function_name__class_name }}" - - # middleware - middleware_section = "default" - middleware_switches = { - "theflow.middleware.TrackProgressMiddleware": True, - "theflow.middleware.SkipComponentMiddleware": True, - "theflow.middleware.CachingMiddleware": False, - } - - # params - params_publish = False - params_subscribe = True - allow_extra: bool = False - - # declare default backend for deployment - default_backend = settings.BASE_BACKEND - - -class ConfigGet: - """A wrapper class for config retrieval""" - - def __init__(self, config: "Config", pipeline: "Function"): - self._config = config - self._pipeline = pipeline - - def __getattr__(self, name: str) -> Any: - attr = getattr(self._config, name) - if callable(attr): - return attr(self._pipeline) - return attr - - def dump(self) -> dict: - """Pass-through the config export""" - return self._config.dump() - - -class ConfigProperty: - """Serve as property to access the config from the pipeline instance""" - - def __get__(self, obj: Optional["Function"], obj_type: Type["Function"]) -> Any: - if obj is None: - return self - - if obj._ff_config is None: - raise ValueError("ConfigProperty can only be accessed after initialization") - return ConfigGet(obj._ff_config, obj) - - def __set__(self, obj: "Function", value: Union[dict, "Config", None]) -> None: - if not isinstance(value, Config): - raise ValueError("ConfigProperty can only be set with Config object") - - if isinstance(value, Config): - obj.__dict__["_ff_config"] = value - elif isinstance(value, dict) or value is None: - obj.__dict__["_ff_config"] = Config(value, cls=obj.__class__) - else: - raise ValueError( - f"Unknown config type: {type(value)}. Must be dict or Config" - ) - - -class Config: - """Config for the pipeline - - Config is a dict-like object that stores the configs for the pipeline. The config - resolution order is: - 1. default config - 2. pipeline.Config from parent classes to child classes in reverse MRO order - 3. config passed to the constructor - - Each value for a config can either be a scalar value, or a string to a callback - function that takes the pipeline instance as the only argument. The callback - function will be called when the config is accessed. The string to the callback - function should be in the format of `{{ module.to.function }}`. - - Args: - config: config dict or path to a yaml file (default: None) - cls: the pipeline class (default: None) - """ - - if TYPE_CHECKING: - from pathlib import Path - - store_result: "Path" - run_id: str - function_name: str - middleware_section: str - middleware_switches: dict[str, bool] - params_publish: bool - params_subscribe: bool - allow_extra: bool - - def __init__( - self, - config: Optional[Union[dict, str]] = None, - cls: Optional[Type["Function"]] = None, - ): - self._available_configs = { - key for key in DefaultConfig.__dict__.keys() if not key.startswith("_") - } - - if cls is not None: - self.update(cls) - if config: - if isinstance(config, str): - with open(config) as f: - config = yaml.safe_load(f) - self.update(config) - - def update_from_dict(self, config: dict): - """Parse the config dict""" - for key, value in config.items(): - if key.startswith("__"): - continue - - if key not in self._available_configs: - raise ValueError(f"Unknown config: {key}") - - if ( - isinstance(value, str) - and value.startswith("{{") - and value.endswith("}}") - ): - # parse to the callback function - dotted_string = value[2:-2].strip() - value = import_dotted_string(dotted_string, safe=False) - - if key in _aggregated_dict: - if not isinstance(value, dict): - raise ValueError(f"Config {key} must be a dict, got {type(value)}") - original_value = getattr(self, key, {}) - original_value.update(value) - value = original_value - - setattr(self, key, value) - - def update_from_pipeline(self, cls: Type["Function"]) -> None: - """Parse the pipeline configs from pipeline.Config""" - classes = cls.mro() - for each_cls in reversed(classes): - if hasattr(each_cls, "Config"): - self.update_from_dict(each_cls.Config.__dict__) - - def update_from_config(self, config: "Config") -> None: - """Parse the pipeline configs from another Config instance""" - self.update_from_dict(config.dump()) - - def update(self, val: Any) -> None: - from .base import Function - - if isinstance(val, dict): - self.update_from_dict(val) - elif isinstance(val, type) and issubclass(val, Function): - self.update_from_pipeline(val) - elif isinstance(val, Config): - self.update_from_config(val) - else: - raise ValueError(f"Unknown config type: {type(val)}") - - def dump(self) -> dict: - """Export the config dict""" - output = {} - for key in self._available_configs: - if callable(getattr(self, key)): - obj = getattr(self, key) - output[key] = f"{{{{ {obj.__module__}.{obj.__name__} }}}}" - else: - output[key] = getattr(self, key) - - return output diff --git a/libs/kotaemon/kotaemon/flow/context.py b/libs/kotaemon/kotaemon/flow/context.py deleted file mode 100644 index 4b2a2f79..00000000 --- a/libs/kotaemon/kotaemon/flow/context.py +++ /dev/null @@ -1,166 +0,0 @@ -from typing import Any, Optional - -from .settings import settings -from .utils.modules import deserialize - - -class Context: - """Context to handle communication - - Context is a key-value store backed by a cache that can be used to share - information between different components in a pipeline. It is also used to store - information about the pipeline itself. - - Any pipeline synchronization problem is essentially a communication problem. Here - the context serves as a communication channel and faciliates the communication - flow easily. - - The context should allow: - - Global context: shared by all pipelines in all steps in all pipelines - - Local context: shared by all steps in each pipeline - - The context should be process-safe and can be used in multi-processing environment. - """ - - def __init__(self): - """Initialize the context""" - - self._cache = deserialize(settings.CACHE, safe=False) - self._global_key = "__global_key__" - if self._global_key not in self._cache: - self._cache.set(self._global_key, {}) - if "__all_contexts__" not in self._cache: - self._cache.set("__all_contexts__", []) - - def _is_context_valid(self, context: Optional[str]) -> str: - """Check if the context name is valid - - Args: - context: name of the context - - Returns: - name of the context - """ - if context is None: - context = self._global_key - - if not isinstance(context, str): - raise ValueError( - f"Context name must be a string or None, got {type(context)}" - ) - - if context not in self._cache: - raise ValueError(f"Context {context} does not exist") - - return context - - def set(self, name: str, value: Any, context: Optional[str] = None) -> None: - """Set a value to the context - - Args: - name: name of the value - value: value to be set - context: name of the context, if None (default), use the global context - """ - - def func(x): - x[name] = value - return x - - context = self._is_context_valid(context) - self._cache.get_then_set(context, func=func, default={}) - - def get( - self, name: Optional[str], default=None, context: Optional[str] = None - ) -> Any: - """Get a value from the context - - Args: - name: name of the value. If None, get all values from the context in a dict - default: default value to return if the value does not exist - context: name of the context, if None (default), use the global context - """ - context = self._is_context_valid(context) - if name is None: - return self._cache[context] - - return self._cache[context].get(name, default) - - def clear(self, name: Optional[str], context: Optional[str]): - """Clear a value from the context - - Args: - name: name of the value. If None, clear all values from the context - context: name of the context, if None, clear global context - """ - context = self._is_context_valid(context) - if name is not None: - - def func(x): - if name in x: - del x[name] - return x - - self._cache.get_then_set(context, func=func, default={}) - else: - self._cache.set(context, {}) - - def has_context(self, context: str) -> bool: - """Check if a context exists - - Args: - context: name of the context - - Returns: - True if the context exists, False otherwise - """ - return context in self._cache - - def create_context(self, context: str, exist_ok=False) -> str: - """Create a context - - Args: - context: name of the context - exist_ok: if True, do not raise error if the context already exists - - Returns: - the string that can be used to access the local context - - Raises: - ValueError: if the context already exists - """ - if not isinstance(context, str): - raise ValueError(f"Context name must be a string, got {type(context)}") - - if context in self._cache: - if exist_ok: - return context - raise ValueError(f"Context {context} already exists") - - def func(x): - x.append(context) - return x - - self._cache.set(context, {}) - self._cache.get_then_set("__all_contexts__", func=func, default=[]) - - return context - - def get_all_contexts_keys(self) -> list: - """Get a list of all contexts - - Returns: - a list of all contexts keys - """ - return self._cache.get("__all_contexts", []) - - def get_all_contexts(self) -> dict: - """Get all contexts stored in the cache - - Returns: - a dict of all contexts - """ - result = {} - for key in self.get_all_contexts_keys(): - result[key] = self.get(None, context=key) - return result diff --git a/libs/kotaemon/kotaemon/flow/debug.py b/libs/kotaemon/kotaemon/flow/debug.py deleted file mode 100644 index 40298473..00000000 --- a/libs/kotaemon/kotaemon/flow/debug.py +++ /dev/null @@ -1,97 +0,0 @@ -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from .base import Function - - -def has_cycle(graph: dict[str, list[str]]) -> bool: - """Check if a graph has cycle - - Args: - graph: A graph represented by an adjacency list - - Returns: - True if the graph has cycle, False otherwise - """ - visited: set[str] = set() - path: set[str] = set() - - def visit(vertex): - if vertex in visited: - return False - visited.add(vertex) - path.add(vertex) - for neighbour in graph.get(vertex, []): - if neighbour in path or visit(neighbour): - return True - path.remove(vertex) - return False - - return any(visit(v) for v in graph) - - -def has_cyclic_dependency(cls: type["Function"]): - """Check if a component's nodes and params has cyclic dependency - - Args: - cls: A function - - Returns: - True if the function has cyclic dependency, False otherwise - """ - params, nodes = cls._collect_registered_params_and_nodes() - specs: dict[str, dict] = {} - graph: dict[str, list[str]] = {} - - # construct dependency graph - for attr in nodes + params: - graph[attr] = [] - spec: dict = specs.get(attr, {}) or getattr(cls, attr).to_dict() - if spec["auto_callback"] or spec["default_callback"]: - if not spec["depends_on"]: - continue - for src in spec["depends_on"]: - src_spec: dict = specs.get(src, {}) or getattr(cls, src).to_dict() - if src_spec["auto_callback"] or src_spec["default_callback"]: - graph[attr].append(src) - - return has_cycle(graph) - - -def likely_cyclic_pipeline(a: "Function", max_node_connections: int = 100): - """Check if a pipeline is likely to have circular loop - - Note, this heuristic assumes that if a pipeline has a lot of node connections, then - it is likely to have circular loop. - - Args: - a: A function - max_node_connections: Maximum number of nodes to check - """ - from collections import defaultdict - - counter: defaultdict[tuple[str, str, str], int] = defaultdict(int) - - threshold_idx = 0 - to_do_idx = 0 - to_dos = [a] - while threshold_idx < max_node_connections and to_do_idx < len(to_dos): - to_do = to_dos[to_do_idx] - for node in to_do._ff_nodes: - target = to_do.get_from_path(node) - if not target: - continue - to_dos.append(target) - triples = ( - f"{to_do.__module__}.{to_do.__class__.__name__}", - node, - f"{target.__module__}.{target.__class__.__name__}", - ) - threshold_idx += 1 - counter[triples] += 1 - to_do_idx += 1 - - result = list(counter.items()) - result = sorted(result, key=lambda x: x[1], reverse=True) - - return threshold_idx == max_node_connections, result diff --git a/libs/kotaemon/kotaemon/flow/exceptions.py b/libs/kotaemon/kotaemon/flow/exceptions.py deleted file mode 100644 index 8cdfd88b..00000000 --- a/libs/kotaemon/kotaemon/flow/exceptions.py +++ /dev/null @@ -1,10 +0,0 @@ -class InvalidAttrDefinition(AttributeError): - pass - - -class CyclicDependencyError(Exception): - pass - - -class CyclicPipelineError(Exception): - pass diff --git a/libs/kotaemon/kotaemon/flow/middleware.py b/libs/kotaemon/kotaemon/flow/middleware.py deleted file mode 100644 index 90916b5f..00000000 --- a/libs/kotaemon/kotaemon/flow/middleware.py +++ /dev/null @@ -1,215 +0,0 @@ -import logging -from abc import abstractmethod -from typing import TYPE_CHECKING, Callable - -if TYPE_CHECKING: - from .base import Function - -logger = logging.getLogger(__name__) - - -class Middleware: - """Middleware template to work on the input and output of a node""" - - def __init__(self, obj: "Function", next_call: Callable): - if obj is None: - raise ValueError("obj must be specified") - self.obj = obj - self.next_call = next_call - - @abstractmethod - def __call__(self, *args, **kwargs): - """Execute the middleware""" - ... - - -class SkipComponentMiddleware(Middleware): - """Skip executing the component if the user specifies so - - This middleware utilizes the following context key: - - good_to_run: if True, the component can be run - - from: skip components in a pipeline earlier than the specified component - - to: skip components in a pipeline later than the specified component (only - valid for parallel) - - This middleware utilizes the following context name: - - __from_run__: this context store the output of the previous run, specified - by the user - - This middleware reserves these kwargs when calling run: - - _ff_from_run: the path to the past run tracker, which will be used to - substitute output of skipped steps - - _ff_from: skip the steps earlier than the specified step - - _ff_to: skip the steps later than the specified step (only valid for parallel) - - """ - - def __call__(self, *args, **kwargs): - """Run the middleware in the context of a wrapping step - - If the step is marked as good_to_run. Run the step as usual. - If the step is not marked as good_to_run: - - Check if the step name matches the name from `from`. If so, it means that - we would want to run this step (as probably all the steps before this step - in the pipeline have been skipped). So we will run this step, and mark - good_to_run as True so that later step will be run as normal. - - Check if the step name matchs the name from `to`. If so, we will run this - step and mark good_to_run as False so that later step will be skipped. - """ - from .runs.base import RunTracker - - # Gather the from, to and from_run from the root pipeline - if _ff_from := kwargs.pop("_ff_from", None): - self.obj.context.set("from", _ff_from, context=self.obj.fl.flow_qualidx) - if _ff_to := kwargs.pop("_ff_to", None): - self.obj.context.set("to", _ff_to, context=self.obj.fl.flow_qualidx) - if _ff_from_run := kwargs.pop("_ff_from_run", None): - from_run = RunTracker(self.obj, "__from_run__") - from_run.load(run_path=_ff_from_run) - - self.obj.context.get("from", context=self.obj.fl.flow_qualidx) - if _from := self.obj.context.get("from", context=self.obj.fl.flow_qualidx): - from .utils.paths import is_parent_of_child - - if is_parent_of_child(self.obj.fl.name, _from): - self.obj.context.set("good_to_run", False, context=self.obj.fl.qualidx) - - # Decide whether to run or fetch from the cache - _ff_name = self.obj.fl.abs_path - - good_to_run: bool = True - if self.obj.context.has_context(context=self.obj.fl.parent_qualidx): - good_to_run = self.obj.context.get( - "good_to_run", default=True, context=self.obj.fl.parent_qualidx - ) - - if good_to_run is False: - from .utils.paths import is_name_matched - - if is_name_matched( - _ff_name, self.obj.context.get("from", context=self.obj.fl.flow_qualidx) - ): - self.obj.context.set( - "good_to_run", True, context=self.obj.fl.parent_qualidx - ) - logger.info(f"Run {_ff_name}. Turn good_to_run from False to True") - self.obj.log_progress(_ff_name, status="run") - return self.next_call(*args, **kwargs) - - try: - from_run = RunTracker(self.obj, which_progress="__from_run__") - output = from_run.output(name=_ff_name) - logger.info(f"Cached {_ff_name}") - self.obj.log_progress(_ff_name, status="cached") - return output - except Exception as e: - logger.warning(f"Failed to get output from run: {e}") - self.obj.log_progress(_ff_name, status="run") - return self.next_call(*args, **kwargs) - - if ( - self.obj.context.get("to", None, context=self.obj.fl.flow_qualidx) - == _ff_name - ): - self.obj.context.set("good_to_run", False, context=self.obj.fl.flow_qualidx) - - self.obj.log_progress(_ff_name, status="run") - return self.next_call(*args, **kwargs) - - -class TrackProgressMiddleware(Middleware): - """Store all information of the current run to the context - - A node can have 1 of the 3 states: - - run: the node is run normally - - cached: the node is not run, and the output is retrieved from the last run - """ - - def __call__(self, *args, **kwargs): - import inspect - - abs_pathx = self.obj.fl.abs_path - if abs_pathx == ".": - from .runs.base import RunTracker - - last_run = RunTracker(self.obj) - last_run.config = self.obj.config.dump() - self.obj.last_run = last_run - - _input = {"args": args, "kwargs": kwargs} - try: - _output = self.next_call(*args, **kwargs) - except Exception as e: - self.obj.log_progress(abs_pathx, input=_input, output=None, error=str(e)) - raise e from None - - if not ( - inspect.isgenerator(_output) - or inspect.isasyncgen(_output) - or inspect.iscoroutine(_output) - or inspect.isawaitable(_output) - ): - try: - self.obj.log_progress(abs_pathx, input=_input, output=_output) - except Exception as e: - import traceback - - logger.warning(f"Failed to log progress: {e}: {traceback.format_exc()}") - - if abs_pathx == ".": - # will be set by the previous code - last_run.persist() # type: ignore - - return _output - - -class CachingMiddleware(Middleware): - """Cache the output of a function and reuse that output if the input and - function definition is the same - """ - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - from .settings import settings - from .utils.modules import deserialize - - self._cache = deserialize(settings.CACHE, safe=False) - - def __call__(self, *args, **kwargs): - try: - hash_key = self.create_key(*args, **kwargs) - if hash_key in self._cache: - return self._cache[hash_key] - except Exception as e: - logger.exception(f"Failed to create key: {e}") - return self.next_call(*args, **kwargs) - - output = self.next_call(*args, **kwargs) - self._cache[hash_key] = output - return output - - def create_key(self, *args, **kwargs) -> str: - """Create a key based on the input and Function's definition - - Specifically, it depends on: - - the `run`'s input - - the Function's class name - - the Function's dump - - Args: - *args: positional arguments of the run - **kwargs: keyword arguments of the run - - Returns: - str: the key - """ - from .utils.hashes import naivehash - - hasher = naivehash() - content = { - "input": {"args": args, "kwargs": kwargs}, - "definition": self.obj.dump(), - "name": self.obj.__class__.__name__, - } - return hasher(content) diff --git a/libs/kotaemon/kotaemon/flow/runs/__init__.py b/libs/kotaemon/kotaemon/flow/runs/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/libs/kotaemon/kotaemon/flow/runs/base.py b/libs/kotaemon/kotaemon/flow/runs/base.py deleted file mode 100644 index 3efdaae2..00000000 --- a/libs/kotaemon/kotaemon/flow/runs/base.py +++ /dev/null @@ -1,193 +0,0 @@ -from __future__ import annotations - -import pickle -import shutil -from pathlib import Path -from typing import TYPE_CHECKING, Any - -import yaml - -if TYPE_CHECKING: - from ..base import Function - from ..context import Context - -from ..storage import storage - - -class RunStructure: - """The structure of a run directory""" - - progress = "progress.pkl" - input = "input.pkl" - output = "output.pkl" - config = "config.yaml" - # kwargs = "kwargs.pkl" ----> Should consolidate with config.yaml? Plug-n-play - # pipeline = "pipeline.yaml" ---> Should consolidate with config.yaml? Plug-n-play - # pipeline_visualization = "pipeline_visualization.dot" - # run_visualization = "run_visualization.dot" - - -class RunManager: - def __init__(self, dir: str): - self.dir: Path = Path(dir) - - def list(self) -> list: - """List all runs""" - # import pandas as pd # TODO: reimplement json_normalize - - # output = [] - # for each_dir in self.dir.iterdir(): - # if not each_dir.is_dir(): - # continue - - # run_info = {"name": each_dir.name} - # with (each_dir / RunStructure.output).open("rb") as fi: - # run_info.update( - # pd.json_normalize({"output": pickle.load(fi)}).to_dict( - # orient="records" - # )[0] - # ) - # with (each_dir / RunStructure.input).open("rb") as fi: - # input_ = pickle.load(fi) - # input_ = {"input": {"args": input_["args"], **input_["kwargs"]}} - # run_info.update(pd.json_normalize(input_).to_dict(orient="records")[0]) - - # output.append(run_info) - - # return output - return [] - - def get(self): - pass - - def delete(self, name: str): - """Delete a run base a name - - Args: - name: the name of the run - """ - shutil.rmtree(self.dir / name) - - -class RunTracker: - """Define run-related methods to track the information in the run - - Args: - obj: the Function that contains necessary information to log info - which_progress: the name of the progress to store the run information. Can be - useful to track multiple progresses in the same run. Default to - "__progress__" which refers to the current progress. - config: the config of the run - """ - - def __init__(self, obj: Function, which_progress: str = "__progress__"): - self._obj = obj - self._context: Context = obj.context - - self._config: dict = {} - self._progress = f"{obj.fl.flow_name}|{obj.fl.run_id}|{which_progress}" - self._context.create_context(self._progress, exist_ok=True) - - if not obj.fl.prefix: - # root pipeline - self._context.set("name", obj.fl.flow_name, context=self._progress) - self._context.set("id", obj.fl.run_id, context=self._progress) - - def log_progress(self, name: str, **kwargs): - """Set the input and output of the step - - Args: - name: name of the step - kwargs: will be logged to the step progress as key, value - """ - value = self._context.get(name, default={}, context=self._progress) - value.update(kwargs) - self._context.set(name, value, context=self._progress) - - def logs(self, name: str | None = None) -> dict: - """Get the information of each step - - Args: - name: name of the pipeline or step. If None, get progress of all steps - - Returns: - input and output of the respective pipeline or step - """ - return self._context.get(name, context=self._progress) - - def steps(self) -> list[str]: - """Get the steps of the run - - Returns: - the steps of the run - """ - return list(self.logs(name=None).keys()) - - def input(self, name: str = ".") -> Any: - """Get the input of a pipeline - - Args: - name: name of the pipeline or step. - - Returns: - input of the respective pipeline - """ - return self.logs(name=name)["input"] - - def output(self, name: str = ".") -> Any: - """Get the output of a pipeline - - Args: - name: name of the pipeline or step. - - Returns: - output of the respective pipeline - """ - return self.logs(name=name)["output"] - - def persist(self): - """Persist the run result to a store""" - dir = storage.join(self._obj.config.store_result, self.id()) - with storage.open(storage.join(dir, "progress.pkl"), "wb") as fo: - pickle.dump(self.logs(name=None), fo) - with storage.open(storage.join(dir, "config.yml"), "w") as fo: - yaml.dump(self._config, fo) - - def id(self) -> str: - """Get the id of the run - - Returns: - the id of the run - """ - return self._context.get("id", context=self._progress) - - def load(self, run_path: str | Path): - """Load a run - - Args: - run_path: the path to the run - """ - run_path = Path(run_path) - with (run_path / "progress.pkl").open("rb") as fi: - progress = pickle.load(fi) - - for key, value in progress.items(): - self._context.set(key, value, context=self._progress) - - @property - def config(self) -> dict | None: - """Get the config of the run - - Returns: - the config of the run - """ - return self._config - - @config.setter - def config(self, config: dict): - """Set the config of the run - - Args: - config: the config of the run - """ - self._config = config diff --git a/libs/kotaemon/kotaemon/flow/safe.py b/libs/kotaemon/kotaemon/flow/safe.py deleted file mode 100644 index 18a7bd36..00000000 --- a/libs/kotaemon/kotaemon/flow/safe.py +++ /dev/null @@ -1,70 +0,0 @@ -"""Construct a flow declaratively in a safe manner.""" -import logging -from typing import Dict, Optional, Type - -from .base import Function -from .utils.modules import deserialize, import_dotted_string - -logger = logging.getLogger(__name__) - - -def load( - obj: dict, - /, - safe=True, - allowed_modules: Optional[Dict[str, Type]] = None, -) -> Function: - """Construct flow from exported dict - - Args: - obj: flow configuration exported with Flow.dump() - safe: if True, only allowed modules can be imported - modules: dict of allowed modules - - Returns: - Function: flow - """ - cls: Type["Function"] - if safe: - if allowed_modules is None: - raise ValueError("A dict of allowed modules not provided when safe=True") - if obj["function"] not in allowed_modules: - raise ValueError( - f"Module {obj['function']} not allowed. " - f"Allowed modules are {list(allowed_modules.keys())}" - ) - cls = allowed_modules[obj["function"]] - else: - cls = import_dotted_string( - obj["function"], safe=safe, allowed_modules=allowed_modules - ) - - params: dict = {} - for name, value in obj["params"].items(): - try: - params[name] = deserialize( - value, safe=safe, allowed_modules=allowed_modules - ) - except Exception as e: - logger.warn(e) - continue - - nodes: dict = { - key: load(value, safe=safe, allowed_modules=allowed_modules) - for key, value in obj["nodes"].items() - } - - func = cls(**params, **nodes) - func._ff_config.update(obj.get("configs", {})) - func._initialize() - - return func - - -def create( - obj: dict, - /, - safe=True, - allowed_modules: Optional[Dict[str, Type]] = None, -) -> Optional[Type[Function]]: - pass diff --git a/libs/kotaemon/kotaemon/flow/settings/__init__.py b/libs/kotaemon/kotaemon/flow/settings/__init__.py deleted file mode 100644 index 8844d0b5..00000000 --- a/libs/kotaemon/kotaemon/flow/settings/__init__.py +++ /dev/null @@ -1,22 +0,0 @@ -"""Settings access for TheFlow runtime.""" -from __future__ import annotations - -from typing import Any - - -class _SettingsProxy: - """Lazy proxy delegating to kotaemon_settings.get_settings().""" - - def __getattr__(self, item: str) -> Any: - from kotaemon_settings import get_settings - - settings = get_settings() - if hasattr(type(settings), item): - return getattr(settings, item) - extra = getattr(settings, "__pydantic_extra__", None) or {} - if item in extra: - return extra[item] - return getattr(settings, item) - - -settings = _SettingsProxy() diff --git a/libs/kotaemon/kotaemon/flow/settings/default.py b/libs/kotaemon/kotaemon/flow/settings/default.py deleted file mode 100644 index a7c6073a..00000000 --- a/libs/kotaemon/kotaemon/flow/settings/default.py +++ /dev/null @@ -1,31 +0,0 @@ -"""Default setting for important variables""" -from pathlib import Path - -from theflow.utils.paths import default_theflow_path, temp_path - -CONTEXT = { - "__type__": "theflow.context.Context", -} - -CACHE = { - "__type__": "theflow.cache.FileCache", - "path": str(Path(temp_path(), "cache")), -} - -STORAGE = { - "__type__": "theflow.storage.LocalStorage", - "prefix": str(default_theflow_path()), -} - -MIDDLEWARE = { - "default": [ - "theflow.middleware.TrackProgressMiddleware", - "theflow.middleware.CachingMiddleware", - "theflow.middleware.SkipComponentMiddleware", - ] -} - - -BASE_BACKEND = { - "__type__": "theflow.backends.Backend", -} diff --git a/libs/kotaemon/kotaemon/flow/storage/__init__.py b/libs/kotaemon/kotaemon/flow/storage/__init__.py deleted file mode 100644 index 220c75a5..00000000 --- a/libs/kotaemon/kotaemon/flow/storage/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -from ..settings import settings -from ..utils.modules import deserialize -from .local import LocalStorage - -# make this lazy (created during first request) -storage = deserialize(settings.STORAGE, safe=False) - -__all__ = ["LocalStorage", "storage"] diff --git a/libs/kotaemon/kotaemon/flow/storage/base.py b/libs/kotaemon/kotaemon/flow/storage/base.py deleted file mode 100644 index 0ad96547..00000000 --- a/libs/kotaemon/kotaemon/flow/storage/base.py +++ /dev/null @@ -1,33 +0,0 @@ -from abc import ABC, abstractmethod -from typing import Optional - - -class BaseStorage(ABC): - @abstractmethod - def __init__(self, *args, **kwargs): - ... - - @abstractmethod - def open(self, path: str, mode: str, encoding: Optional[str] = None): - """Open file in storage. Support context manager""" - ... - - @abstractmethod - def exists(self, path: str) -> bool: - """Check if a path exists in storage""" - ... - - @abstractmethod - def rm(self, path: str): - """Remove a path""" - ... - - @abstractmethod - def join(self, *paths: str) -> str: - """Join paths""" - ... - - @abstractmethod - def url(self, *paths: str) -> str: - """Get url of a path""" - ... diff --git a/libs/kotaemon/kotaemon/flow/storage/local.py b/libs/kotaemon/kotaemon/flow/storage/local.py deleted file mode 100644 index 64f8b46d..00000000 --- a/libs/kotaemon/kotaemon/flow/storage/local.py +++ /dev/null @@ -1,35 +0,0 @@ -from pathlib import Path -from typing import Optional - -from .base import BaseStorage - - -class LocalStorage(BaseStorage): - """Local on-disk file storage - - Args: - prefix: the prefix of the storage - """ - - def __init__(self, prefix: str): - self._prefix: Path = Path(prefix) - if not self._prefix.exists(): - self._prefix.mkdir(parents=True) - - def open(self, path: str, mode="rb", encoding: Optional[str] = None): - parent = (self._prefix / path).parent - if not parent.exists(): - parent.mkdir(parents=True) - return open(self._prefix / path, mode=mode, encoding=encoding) - - def exists(self, path: str) -> bool: - return (self._prefix / path).exists() - - def rm(self, path: str): - (self._prefix / path).unlink() - - def join(self, *paths: str) -> str: - return str(Path(*paths)) - - def url(self, *paths: str) -> str: - return str(self._prefix.joinpath(*paths)) diff --git a/libs/kotaemon/kotaemon/flow/utils/__init__.py b/libs/kotaemon/kotaemon/flow/utils/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/libs/kotaemon/kotaemon/flow/utils/documentation.py b/libs/kotaemon/kotaemon/flow/utils/documentation.py deleted file mode 100644 index 4df14e15..00000000 --- a/libs/kotaemon/kotaemon/flow/utils/documentation.py +++ /dev/null @@ -1,111 +0,0 @@ -"""Utility modules to extract documentation from the source Function.""" -from __future__ import annotations - -import importlib -import inspect -import pkgutil -import sys - -from ..base import Function, NodeAttr, ParamAttr - - -def get_function_documentation(func: type[Function]) -> dict: - """Return the documentation of the Function. - - Returns: - Dictionary description of the Function, suitable to be parsed. Sample: - { - "desc": "Description of the flow", - "nodes": { - "node_1": { - "desc": "Description of the node", - "type": "Default type of the node", - "input": "The input interface", - "output": "The output interface" - }, - }, - "params": { - "param_1": { - "desc": "Description of the parameter", - "type": "Default type of the parameter", - "default": "Default value of the parameter" - }, - } - } - """ - params, nodes = {}, {} - for name in dir(func): - attr = getattr(func, name) - if isinstance(attr, ParamAttr): - params[name] = { - "desc": attr._help, - "type": attr._type, - "default": attr._default, - "depends_on": attr._depends_on, - } - elif isinstance(attr, NodeAttr): - nodes[name] = { - "desc": attr._help, - "type": attr._default, - "input": attr._input, - "output": attr._output, - "depends_on": attr._depends_on, - } - - return { - "desc": func.__doc__ or "", - "params": params, - "nodes": nodes, - } - - -def get_functions_from_module(module_path: str, recursive: bool = True) -> dict: - """Get all Functions from module - - Args: - module_path: The path to the module - - Returns: - A dictionary of Functions, with the key being the name of the function and the - value being the Functions class itself. - """ - funcs: dict = {} - module = sys.modules.get(module_path) - if not ( - module - and (spec := getattr(module, "__spec__", None)) - and getattr(spec, "_initializing", False) is False - ): - module = importlib.import_module(module_path) - - for name, obj in inspect.getmembers(module): - if inspect.isclass(obj) and issubclass(obj, Function) and obj != Function: - if not obj.__module__.startswith(module_path): - # irrelevant import - continue - funcs[f"{obj.__module__}.{obj.__name__}"] = obj - - if recursive and "__path__" in dir(module): - for _, name, _ in pkgutil.iter_modules(module.__path__): - funcs.update( - get_functions_from_module(f"{module_path}.{name}", recursive=True) - ) - - return funcs - - -def get_function_documentation_from_module( - module_path: str, recursive: bool = True -) -> dict: - """Get all functions documenations from module - - Args: - module_path: The path to the module - recursive: Whether to recursively search for functions in submodules - - Returns: - A dictionary of functions, with the key being the name of the function and the - value being the function documentation - """ - funcs = get_functions_from_module(module_path, recursive=recursive) - return {name: get_function_documentation(func) for name, func in funcs.items()} diff --git a/libs/kotaemon/kotaemon/flow/utils/hashes.py b/libs/kotaemon/kotaemon/flow/utils/hashes.py deleted file mode 100644 index 188fad66..00000000 --- a/libs/kotaemon/kotaemon/flow/utils/hashes.py +++ /dev/null @@ -1,61 +0,0 @@ -from hashlib import md5 -from typing import Any - - -class naivehash: - """Hash a Python object - - Args: - hash_func: hash function to use. Default is md5 - """ - - def __init__(self, hash_func=None): - """Initialize the hash object""" - self.hash_func = hash_func() if hash_func is not None else md5() - - def update(self, obj: Any): - """Hash a Python object - - Args: - obj: Python object to be hashed - - Returns: - hash of the object - """ - type_ = f"{chr(0)}{type(obj)}{chr(0)}" - if isinstance(obj, (str, int, float, bool)) or obj is None: - self.hash_func.update(f"|{type_}|{obj}".encode()) - elif isinstance(obj, (tuple, list)): - self.update(f"|{type_}|") - for idx, item in enumerate(obj): - self.update(f"|{type_}{idx}|") - self.update(item) - elif isinstance(obj, set): - self.update(f"|{type_}|") - for idx, item in enumerate(sorted(obj)): - self.update(f"|{type_}{idx}|") - self.update(item) - elif isinstance(obj, dict): - self.update(f"|{type_}|") - for idx, key in enumerate(sorted(obj)): - self.update(f"|{type_}{idx}|") - self.update(key) - self.update(obj[key]) - else: - path = "" - path += str(obj.__module__) if hasattr(obj, "__module__") else "" - path += str(obj.__name__) if hasattr(obj, "__name__") else "" - self.update(f"|{type_}|{path}|") - - for idx, attr in enumerate(sorted(dir(obj))): - if attr.startswith("_"): - continue - self.update(f"|{type_}{idx}|") - self.update(attr) - # avoid self.update(getattr(obj, attr)) to avoid infinite recursion - self.update(str(getattr(obj, attr))) - - def __call__(self, obj: Any) -> str: - """Return the hash digest""" - self.update(obj) - return self.hash_func.hexdigest() diff --git a/libs/kotaemon/kotaemon/flow/utils/modules.py b/libs/kotaemon/kotaemon/flow/utils/modules.py deleted file mode 100644 index 795c2a38..00000000 --- a/libs/kotaemon/kotaemon/flow/utils/modules.py +++ /dev/null @@ -1,284 +0,0 @@ -import importlib -import inspect -import logging -import sys -from pathlib import Path -from typing import Any, Dict, Generic, Optional, Type, TypeVar - -logger = logging.getLogger(__name__) -NATIVE_TYPE = (dict, list, tuple, str, int, float, bool, type(None)) - - -def import_dotted_string( - dotted_string: str, /, safe=True, allowed_modules: Optional[Dict[str, Type]] = None -): - """Import a dotted string - - Args: - dotted_string: the dotted string to import - safe: if True, only allowed modules can be imported - allowed_modules: dict of allowed modules - - Returns: - the imported object - """ - if safe: - if allowed_modules is None: - raise ValueError("Must provide allowed_modules when safe=True") - - if dotted_string not in allowed_modules: - raise ValueError( - f"Module {dotted_string} is not allowed. " - f"Allowed modules are {list(allowed_modules.keys())}" - ) - - return allowed_modules[dotted_string] - - module_name, obj_name = dotted_string.rsplit(".", 1) - module = sys.modules.get(module_name) - - if not ( - module - and (spec := getattr(module, "__spec__", None)) - and getattr(spec, "_initializing", False) is False - ): - module = importlib.import_module(module_name) - - return getattr(module, obj_name) - - -def serialize_path(path: Path) -> dict: - """Serialize a Path object. - - For cross-platform compatibility, the type will be "pathlib.Path" rather than - "posixpath" or "ntpath". - """ - return {"__type__": "pathlib.Path", "path": str(path)} - - -SERIALIZE_BY_TYPES = { - Path: serialize_path, -} - - -def import_modules(*module_names: str) -> tuple: - """Import a module by string name, raise exception if not found - - Args: - module_name: the module name to import - - Returns: - the imported module - - Raises: - ImportError: if any of the module cannot be imported - """ - errors: list[str] = [] - modules = [] - for module_name in module_names: - try: - module = sys.modules.get(module_name) - - if not ( - module - and (spec := getattr(module, "__spec__", None)) - and getattr(spec, "_initializing", False) is False - ): - module = importlib.import_module(module_name) - modules.append(module) - except ImportError as e: - errors.append(module_name) - logger.warn(f"Cannot import module {module_name}: {e}") - - if errors: - raise ImportError(f"Cannot import modules: {', '.join(errors)}") - - return tuple(modules) - - -def serialize(value: Any) -> Any: - """Serialize a value to a JSON-serializable object""" - if isinstance(value, dict): - return {key: serialize(val) for key, val in value.items()} - - if isinstance(value, list): - return [serialize(val) for val in value] - - if isinstance(value, tuple): - return tuple(serialize(val) for val in value) - - if isinstance(value, NATIVE_TYPE): - return value - - if inspect.isfunction(value) or inspect.isclass(value): - if value.__name__ == "": - raise ValueError("Cannot serialize lambda functions") - return f"{{{{ {value.__module__}.{value.__name__} }}}}" - - if hasattr(value, "__persist_flow__"): - d = value.__persist_flow__() - return d - - for base in value.__class__.mro()[:-1]: - if base in SERIALIZE_BY_TYPES: - return SERIALIZE_BY_TYPES[base](value) - - if value.__module__ == "builtins": - return f"{{{{ {value.__module__}.{value.__name__} }}}}" - - if value.__module__ == "typing": - name = "" - if hasattr(value, "__name__"): - name = value.__name__ - elif hasattr(value, "_name"): - name = value._name - else: - raise ValueError(f"Cannot serialize {value}. Unknown name") - return f"{{{{ {value.__module__}.{name} }}}}" - - raise ValueError( - f"Cannot serialize {value}. Consider implementing __persist_flow__" - ) - - -def deserialize( - value: Any, /, safe=True, allowed_modules: Optional[Dict[str, Type]] = None -) -> Any: - """Deserialize a JSON-serializable object to a Python object - - Args: - value: the value to deserialize - safe: if True, only allowed modules can be imported - allowed_modules: dict of allowed modules - """ - if isinstance(value, str) and value.startswith("{{") and value.endswith("}}"): - return import_dotted_string( - value[2:-2].strip(), safe=safe, allowed_modules=allowed_modules - ) - - if isinstance(value, dict) and "__type__" in value: - cls = import_dotted_string( - value["__type__"], safe=safe, allowed_modules=allowed_modules - ) - params: dict = {} - for key, val in value.items(): - if key == "__type__": - continue - params[key] = deserialize(val, safe=safe, allowed_modules=allowed_modules) - return cls(**params) - - if isinstance(value, dict) and "__type__" not in value: - return { - key: deserialize(val, safe=safe, allowed_modules=allowed_modules) - for key, val in value.items() - } - - if isinstance(value, list): - return [ - deserialize(val, safe=safe, allowed_modules=allowed_modules) - for val in value - ] - - if isinstance(value, tuple): - return tuple( - deserialize(val, safe=safe, allowed_modules=allowed_modules) - for val in value - ) - - if isinstance(value, NATIVE_TYPE): - return value - - raise ValueError(f"Cannot deserialize type {type(value)} ({value})") - - -T = TypeVar("T") - - -class lazy(Generic[T]): - """Declare the init parameters to initialize an object - - This will declare an object and initialize it later, useful to set default - values of a class parameter. - """ - - def __init__(self, cls: Type[T], **params): - self._cls: Type[T] = cls - self._params: dict = params - - def __call__(self) -> T: - """Initialize the object""" - params = {} - for key, val in self._params.items(): - if isinstance(val, lazy): - params[key] = val() - else: - params[key] = val - - return self._cls(**params) - - def withx(self, **params) -> "lazy[T]": - """Continue declaring the object with additional parameters""" - return lazy(self._cls, **{**self._params, **params}) - - @classmethod - def from_serialized(cls, d: dict): - """Convert a dict-serialized object into an lazy object""" - target_cls = import_dotted_string(d.pop("__type__"), safe=False) - for key, value in d.items(): - if isinstance(value, dict) and "__type__" in value: - d[key] = cls.from_serialized(value) - return lazy(target_cls, **d) - - def __persist_flow__(self) -> dict: - """Express the object as a dict""" - params = {} - for key, value in self._params.items(): - params[key] = value.__persist_flow__() if isinstance(value, lazy) else value - - return {"__type__": f"{self._cls.__module__}.{self._cls.__name__}", **params} - - def __rshift__(self, other: "lazy[T]") -> Any: - """Chain two lazy objects together""" - from theflow.base import Function, SequentialFunction - - if not isinstance(other, lazy): - raise ValueError(f"Cannot chain lazy and non-lazy objects: {other}") - - if not issubclass(other._cls, Function) or not issubclass(self._cls, Function): - raise ValueError("Can only chain lazy Function") - - funcs = [] - if issubclass(self._cls, SequentialFunction): - funcs.extend(self._params.get("funcs", [])) - else: - funcs.append(self) - - if issubclass(other._cls, SequentialFunction): - funcs.extend(other._params.get("funcs", [])) - else: - funcs.append(other) - - return lazy(SequentialFunction, funcs=funcs) - - def __floordiv__(self, other: "lazy[T]") -> Any: - """Chain two lazy objects together""" - from theflow.base import ConcurrentFunction, Function - - if not isinstance(other, lazy): - raise ValueError(f"Cannot chain lazy and non-lazy objects: {other}") - - if not issubclass(other._cls, Function) or not issubclass(self._cls, Function): - raise ValueError("Can only chain lazy Function") - - funcs = [] - if issubclass(self._cls, ConcurrentFunction): - funcs.extend(self._params.get("funcs", [])) - else: - funcs.append(self) - - if issubclass(other._cls, ConcurrentFunction): - funcs.extend(other._params.get("funcs", [])) - else: - funcs.append(other) - - return lazy(ConcurrentFunction, funcs=funcs) diff --git a/libs/kotaemon/kotaemon/flow/utils/multiprocess.py b/libs/kotaemon/kotaemon/flow/utils/multiprocess.py deleted file mode 100644 index b36cf8d5..00000000 --- a/libs/kotaemon/kotaemon/flow/utils/multiprocess.py +++ /dev/null @@ -1,45 +0,0 @@ -import multiprocessing -import multiprocessing.managers -from typing import TYPE_CHECKING, Dict, List, cast - -if TYPE_CHECKING: - from ..base import Function - - -def _run_node(task): - obj: "Function" = task[0] - child_name: str = task[1] - params: Dict = task[2] - lock = task[3] - - with lock: - node = getattr(obj, child_name) - return node(**params) - - -def parallel(obj: "Function", child_name: str, tasks: List[Dict], **kwargs): - """Run a node in parallel with multiprocessing. - - This helper function allows accurately keeping track of the the number of time the - `child_name` node is called from the `obj` parent. - - Args: - obj (Function): Function object - child_name (str): Child name - tasks (List[Dict]): List of parameters for each task - kwargs: Keyword arguments for multiprocessing.Pool - """ - manager = None - try: - manager = multiprocessing.Manager() - obj._ff_childs_called = cast("dict", manager.dict(obj._ff_childs_called)) - lock = manager.Lock() - - tasks_mp = [(obj, child_name, task, lock) for task in tasks] - with multiprocessing.Pool(**kwargs) as pool: - yield from pool.imap(_run_node, tasks_mp) - finally: - if isinstance(obj._ff_childs_called, multiprocessing.managers.DictProxy): - obj._ff_childs_called = obj._ff_childs_called.copy() - if manager is not None: - manager.shutdown() diff --git a/libs/kotaemon/kotaemon/flow/utils/paths.py b/libs/kotaemon/kotaemon/flow/utils/paths.py deleted file mode 100644 index 53fd71dc..00000000 --- a/libs/kotaemon/kotaemon/flow/utils/paths.py +++ /dev/null @@ -1,183 +0,0 @@ -import re -from pathlib import Path -from typing import List, Optional, Union - -THEFLOW_DIR = ".theflow" - - -def project_root(loc: Optional[Union[str, Path]] = None) -> Path: - """Get the root directory of the project (contains .git/). Return cwd if .git/ - doesn't exist - - Args: - loc: the location to start searching for the root directory. If None, - assume the current working directory - - Returns: - the root directory of the project, or None if not found - """ - if loc is None: - loc = Path.cwd() - loc = Path(loc) - while loc != loc.parent: - if (loc / ".git").exists(): - return loc - loc = loc.parent - - return Path.cwd() - - -def get_theflow_path(loc: Optional[Union[str, Path]]) -> Optional[Path]: - """Get the theflow directory (contains .theflow/) - - Args: - loc: the location to start searching for the theflow directory. If None, - assume the current working directory - - Returns: - the theflow directory, or None if not found - """ - loc = Path.cwd() if loc is None else Path(loc) - while loc != loc.parent: - if (loc / THEFLOW_DIR).exists(): - return loc / THEFLOW_DIR - loc = loc.parent - - return None - - -def default_theflow_path( - loc: Union[None, str, Path] = None, create: bool = False -) -> Path: - """Get the theflow directory (.theflow/) or create it if not exists - - It travels up the directory tree until it finds the theflow directory. If not - found, it creates the theflow directory in the project root folder. If there - isn't project root folder, it creates the theflow directory in the current - working directory. - - Args: - loc: the location to start searching for the theflow directory. If None, - assume the current working directory - create: whether to create the theflow directory if not exists - - Returns: - the theflow directory - """ - if loc is None: - loc = Path.cwd() - loc = Path(loc) - flow_path = get_theflow_path(loc) - if flow_path is not None: - return flow_path - - flow_path = project_root(loc) / THEFLOW_DIR - if create: - flow_path.mkdir(exist_ok=True, parents=True) - - return flow_path - - -def temp_path() -> str: - """Get the default temporary directory - - Returns: - the default temporary directory - """ - import getpass - import os - import tempfile - - default: str = os.environ.get("THEFLOW_TEMP_PATH", "") - if not default: - try: - username = getpass.getuser() - except Exception: - username = "" - path = Path(tempfile.gettempdir(), f"theflow_{username}") - else: - path = Path(default) - path.mkdir(exist_ok=True, parents=True) - return str(path) - - -def is_name_matched(name: str, pattern: str) -> bool: - """Check if a name matches a pattern - - This method matches simple pattern with wildcard character "*". For example, the - pattern "a.*.b" matches "a.c.b" but not "a.b.c.b". - - Args: - name: the name to check - pattern: the pattern to match - - Returns: - True if the name matches the pattern, False otherwise - """ - pattern_parts: List[str] = [re.escape(part) for part in pattern.split("*")] - regex_pattern: str = r"^" + r"[^.]+".join(pattern_parts) + r"$" - return re.findall(regex_pattern, name) != [] - - -def is_parent_of_child(parent: str, child: str) -> bool: - """Check if a name is a direct parent of another name - - Example: - >> is_parent_of_child(".main.pipeline_A1", ".main.pipeline_A1.*") - True - >> is_parent_of_child(".main.pipeline_A1", ".main.pipeline_A1.pipeline_B1") - True - >> is_parent_of_child(".main.pipeline_A1", ".main.pipeline_A2") - False - - Args: - parent: the parent name - child: the child name. The child can be a wildcard pattern - - Returns: - True if the parent is a parent of the child, False otherwise - """ - parent, child = parent.strip("."), child.strip(".") - pattern = ".".join(child.split(".")[:-1]) - return is_name_matched(parent, pattern) - - -if __name__ == "__main__": - names = [ - "", - ".main", - ".main.pipeline_A1" ".main.pipeline_A1.pipeline_B1", - ".main.pipeline_A1.pipeline_B1.pipeline_C1", - ".main.pipeline_A1.pipeline_B1.pipeline_C1.step_a", - ".main.pipeline_A1.pipeline_B1.pipeline_C1.step_b", - ".main.pipeline_A1.pipeline_B1.pipeline_C1.step_c", - ".main.pipeline_A1.pipeline_B2", - ".main.pipeline_A1.pipeline_B2.pipeline_C2", - ".main.pipeline_A1.pipeline_B2.pipeline_C2.step_a", - ".main.pipeline_A1.pipeline_B2.pipeline_C2.step_b", - ".main.pipeline_A1.pipeline_B2.pipeline_C2.step_c", - ".main.pipeline_A2", - ".main.pipeline_A2.pipeline_B1", - ".main.pipeline_A2.pipeline_B1.pipeline_C1", - ".main.pipeline_A2.pipeline_B1.pipeline_C1.step_a", - ".main.pipeline_A2.pipeline_B1.pipeline_C1.step_b", - ".main.pipeline_A2.pipeline_B1.pipeline_C1.step_c", - ".main.pipeline_A2.pipeline_B2", - ".main.pipeline_A2.pipeline_B2.pipeline_C2", - ".main.pipeline_A2.pipeline_B2.pipeline_C2.step_a", - ".main.pipeline_A2.pipeline_B2.pipeline_C2.step_b", - ".main.pipeline_A2.pipeline_B2.pipeline_C2.step_c", - ".main.next_step", - ] - - # pattern = ".main.*.p*.*.step_a" - # pattern = ".main.*.*.pipeline_C1" - # pattern = ".main.pipeline_A2" - pattern = ".*.pipeline*" - # for name in names: - # if is_name_matched(name, pattern): - # print(name) - - for name in names: - if is_parent_of_child(name, pattern): - print(name, "is parent of", pattern) diff --git a/libs/kotaemon/kotaemon/flow/utils/pretties.py b/libs/kotaemon/kotaemon/flow/utils/pretties.py deleted file mode 100644 index 25806197..00000000 --- a/libs/kotaemon/kotaemon/flow/utils/pretties.py +++ /dev/null @@ -1,77 +0,0 @@ -import re - - -def reindent_docstring(docin: str) -> str: - """Remove beginning whitespace in a docstring - - Happens when the docstring is indented in method, function, class... but we want to - make it look nice in CLI. - - Args: - docin: the docstring to reindent - - Returns: - the reindented docstring - """ - if not docin: - return "" - - whitespaces = re.findall(r"\n[ \t]+", docin) - if whitespaces: - min_whitespace = min(whitespaces)[1:] - lines = [] - for line in docin.splitlines(): - if line.startswith(min_whitespace): - line = line[len(min_whitespace) :] - lines.append(line) - docin = "\n".join(lines).strip() - return docin - - -def flatten_dict(indict: dict) -> dict: - """Flatten nested dict into 1-level dict - - Example: - >>> flatten_dict({"a": {"b": 1}, "c": 2}) - {"a.b": 1, "c": 2} - - Args: - indict: input dict - - Returns: - flattened dict - """ - outdict = {} - for key, value in indict.items(): - if isinstance(value, dict): - for subkey, subvalue in flatten_dict(value).items(): - outdict[f"{key}.{subkey}"] = subvalue - else: - outdict[key] = value - return outdict - - -def unflatten_dict(indict: dict) -> dict: - """Unflatten 1-level dict into nested dict - - Example: - >>> unflatten_dict({"a.b": 1, "c": 2}) - {"a": {"b": 1}, "c": 2} - - Args: - indict: input dict - - Returns: - unflattened dict - """ - outdict: dict = {} - for key, value in indict.items(): - key = key.strip(".") - subkeys = key.split(".") - subdict = outdict - for subkey in subkeys[:-1]: - if subkey not in subdict: - subdict[subkey] = {} - subdict = subdict[subkey] - subdict[subkeys[-1]] = value - return outdict diff --git a/libs/kotaemon/kotaemon/flow/utils/typings.py b/libs/kotaemon/kotaemon/flow/utils/typings.py deleted file mode 100644 index 1fb7c797..00000000 --- a/libs/kotaemon/kotaemon/flow/utils/typings.py +++ /dev/null @@ -1,110 +0,0 @@ -"""Typing utilities - -These utilities support introspection of nodes and params through type annotations. -Consider these as experimental due to very naive implementation that for sure cannot -handle uncommon cases. -""" -import inspect -from typing import _GenericAlias # type: ignore -from typing import Any, Callable, Union, get_args, get_origin - - -def is_union_type(annotation) -> bool: - """Check if the annotation is a Union type""" - return annotation is Union or ( - isinstance(annotation, _GenericAlias) and annotation.__origin__ is Union - ) - - -def expand_types(annotation) -> list: - """Expand the type from source to target""" - result = [] - if is_union_type(annotation): - childs = get_args(annotation) - for child in childs: - result += expand_types(child) - return result - - origin = get_origin(annotation) - if origin: - result.append(origin) - else: - result.append(annotation) - return result - - -def is_compatible_with(type1, type2) -> bool: - """Check if the annotation type1 is at slightest compatible with type2 - - Slightest compatibility happens when there is at least 1 type in type1 that is - a subclass of at least 1 type in type2 - """ - type1s = expand_types(type1) - type2s = expand_types(type2) - if Any in type1s: - return True - if Any in type2s: - return True - - for each_1 in type1s: - for each_2 in type2s: - try: - if issubclass(each_1, each_2): - return True - except Exception: - continue - - return False - - -def input_signature( - func: Callable, ignore_bound: bool = True -) -> tuple[dict, bool, bool]: - """Get the input signature of a function or method - - Args: - func: the function or method to get the signature - ignore_bound: ignore the first argument if it is self or cls - - Returns: - - a dict of {argument_name: argument_type} - - a bool indicating if the function has *args - - a bool indicating if the function has **kwargs - """ - args = inspect.signature(func).parameters - type_annotation = {} - bounds = {"self", "cls"} - has_args, has_kwargs = False, False - for name, arg in args.items(): - if name in bounds and ignore_bound: - continue - if arg.kind == inspect.Parameter.VAR_POSITIONAL: - has_args = True - continue - if arg.kind == inspect.Parameter.VAR_KEYWORD: - has_kwargs = True - continue - if arg.annotation is not inspect.Parameter.empty: - type_annotation[name] = arg.annotation - continue - if arg.default is not inspect.Parameter.empty and arg.default is not None: - type_annotation[name] = type(arg.default) - continue - type_annotation[name] = Any - - return type_annotation, has_args, has_kwargs - - -def output_signature(func: Callable) -> Any: - """Get the output signature of a function or method - - Args: - func: the function or method to get the signature - - Returns: - the return type annotation - """ - annot = inspect.signature(func).return_annotation - if annot is inspect.Signature.empty: - annot = Any - return annot diff --git a/libs/kotaemon/kotaemon/flow/visualization.py b/libs/kotaemon/kotaemon/flow/visualization.py deleted file mode 100644 index 3e774617..00000000 --- a/libs/kotaemon/kotaemon/flow/visualization.py +++ /dev/null @@ -1,269 +0,0 @@ -from __future__ import annotations - -import ast -import inspect -from collections import defaultdict - -IGNORE = ( - ast.Call, - ast.arguments, - ast.arg, - ast.Assign, - ast.Name, - ast.Attribute, - ast.Load, - ast.keyword, - ast.Constant, - ast.Return, - ast.Dict, - ast.Subscript, - ast.JoinedStr, - ast.If, - ast.IfExp, -) - - -def trace_pipelne_run(cls) -> list: - """Trace the logic flow of a pipeline run - - Args: - cls: the pipeline class - - Returns: - list: the logic flow of the pipeline run (suitable for dot) - """ - tree = ast.parse(inspect.getsource(cls)) - analyzer = PipelineRunTracer() - analyzer.visit(tree) - - return analyzer.logic_flow - - -def get_ast_node_name(node) -> str: - """Get human-readable name of an ast node - - Args: - node: ast node - - Returns: - str: human-readable name of the ast node - """ - if isinstance(node, ast.Attribute): - return f"{get_ast_node_name(node.value)}.{node.attr}" - elif isinstance(node, ast.Name): - return node.id - elif isinstance(node, ast.FunctionDef): - return node.name - elif isinstance(node, ast.arg): - return node.arg - elif isinstance(node, ast.Constant): - return node.value - elif isinstance(node, ast.Call): - return get_ast_node_name(node.func) - elif isinstance(node, ast.JoinedStr): - text = "" - for value in node.values: - text += get_ast_node_name(value) - return text - elif isinstance(node, ast.FormattedValue): - return get_ast_node_name(node.value) - else: - return node - - -class PipelineRunTracer(ast.NodeVisitor): - def __init__(self): - self.logic_flow = [] - self.in_run = False - self._stacks: list[dict[str, list]] = [defaultdict(list)] - self._last_call = None - - def get_creator_of(self, name, default=None): - """Find the creator of a variable""" - for idx in range(len(self._stacks) - 1, -1, -1): - if name in self._stacks[idx]: - return self._stacks[idx][name] - - return default - - def set_creator_of(self, name, value): - """Set creator of a value""" - self._stacks[-1][name] = value - - def stack_begin(self): - stack: dict[str, list] = defaultdict(list) - self._stacks.append(stack) - return stack - - def stack_end(self): - return self._stacks.pop() - - def visit(self, node): - if self.in_run: - if not isinstance(node, IGNORE): - print(node) - return super().visit(node) - - def visit_Assign(self, node): - print(f"{self.indent()}Assign:") - - # TODO: seems Assign doesn't need stack - self.generic_visit(node) - for target in node.targets: - name = get_ast_node_name(target) - if self._last_call is not None: - self.set_creator_of(name, self._last_call) - self._last_call = None - - def visit_Name(self, node): - print(f"{self.indent()}Name: {get_ast_node_name(node)}") - - def indent(self): - return " " * len(self._stacks) - - def visit_FunctionDef(self, node): - if node.name != "run": - return - - self.in_run = True - print(f"{self.indent()}Function: {get_ast_node_name(node)}") - - return self.generic_visit(node) - - def visit_Call(self, node): - if not self.in_run: - return - - # determine function name - node_name = get_ast_node_name(node) - for kw in node.keywords: - if kw.arg == "_ff_name": - node_name = get_ast_node_name(kw.value) - break - - print(f"{self.indent()}Call: {node_name}") - - # determine args relation - for arg in node.args: - if isinstance(arg, ast.Call): - self.visit(arg) - self.logic_flow.append((self._last_call, node_name)) - self._last_call = None - elif isinstance(arg, ast.Name): - from_node = self.get_creator_of(arg.id, "__begin__") - if isinstance(from_node, str): - self.logic_flow.append((from_node, node_name)) - else: - for each in from_node: - self.logic_flow.append((each, node_name)) - else: - print("Else condition") - - # determine kwargs relation - for kw in node.keywords: - if kw.arg == "_ff_name": - continue - if isinstance(arg, ast.Call): - self.visit(arg) - self.logic_flow.append((self._last_call, node_name)) - self._last_call = None - elif isinstance(kw.value, ast.Name): - from_node = self.get_creator_of(kw.value.id, "__begin__") - if isinstance(from_node, str): - self.logic_flow.append((from_node, node_name)) - else: - for each in from_node: - self.logic_flow.append((each, node_name)) - else: - print("Else condition") - self._last_call = node_name - - def visit_Attribute(self, node): - print(f"{self.indent()}Attribute: {get_ast_node_name(node)}") - # return self.generic_visit(node) - - def visit_arg(self, node): - print(f"{self.indent()}Arg: {get_ast_node_name(node)}") - - def visit_arguments(self, node): - print(f"{self.indent()}Arguments:") - self.generic_visit(node) - - def visit_keyword(self, node): - print(f"{self.indent()}Keyword: {get_ast_node_name(node)}") - self.generic_visit(node) - - def visit_Load(self, node): - print(f"{self.indent()}Load: {get_ast_node_name(node)}") - return self.generic_visit(node) - - def visit_Constant(self, node): - print(f"{self.indent()}Constant: {get_ast_node_name(node)}") - return self.generic_visit(node) - - def visit_Return(self, node): - print(f"{self.indent()}Return: {get_ast_node_name(node)}") - self.generic_visit(node) - - def visit_Dict(self, node): - print(f"{self.indent()}Dict: {get_ast_node_name(node)}") - self.generic_visit(node) - - def visit_Subscript(self, node): - print(f"{self.indent()}Subscript: {get_ast_node_name(node)}") - self.generic_visit(node) - - def visit_If(self, node): - self.generic_visit(node.test) - self.stack_begin() - for each in node.body: - self.visit(each) - creator1 = self.stack_end() - - self.stack_begin() - for each in node.orelse: - self.visit(each) - creator2 = self.stack_end() - - creator = {} - creator.update(creator1) - for key, value in creator2.items(): - if key in creator: - value1 = [value] if not isinstance(value, list) else value - value2 = ( - [creator[key]] - if not isinstance(creator[key], list) - else creator[key] - ) - creator[key] = value1 + value2 - else: - creator[key] = value - self._stacks[-1].update(creator) - - def report(self): - for each_from, each_to in self.logic_flow: - print(f"{each_from} -> {each_to}") - - def visit_IfExp(self, node): - print(f"{self.indent()}IfExp: {get_ast_node_name(node)}") - - self.visit(node.test) - test = self._last_call - - last_call = [] - self.visit(node.body) - if self._last_call is not None: - if test is not None: - self.logic_flow.append((test, self._last_call)) - last_call.append(self._last_call) - self._last_call = None - - self.visit(node.orelse) - if self._last_call is not None: - if test is not None: - self.logic_flow.append((test, self._last_call)) - last_call.append(self._last_call) - self._last_call = None - - if last_call: - self._last_call = last_call