Mirrors/kotaemon
1
0
Fork 0
mirror of https://github.com/Cinnamon/kotaemon.git synced 2026-08-29 10:08:59 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactor: remove redundant folders

Browse Source
This commit is contained in:
phv2312
2026-06-01 20:15:23 +07:00
parent 6fc91dc22f
commit 9490e8d020
34 changed files with 0 additions and 4775 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
libs/kotaemon/kotaemon/flow/__init__.py
View File

@@ -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",
]

4
libs/kotaemon/kotaemon/flow/backends/__init__.py
View File

@@ -1,4 +0,0 @@
from .base import Backend
from .http_sync import HttpSyncBackend
__all__ = ["Backend", "HttpSyncBackend"]

154
libs/kotaemon/kotaemon/flow/backends/base.py
View File

@@ -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

173
libs/kotaemon/kotaemon/flow/backends/http_sync.py
View File

@@ -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

1757
libs/kotaemon/kotaemon/flow/base.py
View File

File diff suppressed because it is too large Load Diff

6
libs/kotaemon/kotaemon/flow/cache/__init__.py vendored
View File

@@ -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"]

159
libs/kotaemon/kotaemon/flow/cache/base.py vendored
View File

@@ -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
"""
...

92
libs/kotaemon/kotaemon/flow/cache/filebased.py vendored
View File

@@ -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

90
libs/kotaemon/kotaemon/flow/cache/memcached.py vendored
View File

@@ -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)

102
libs/kotaemon/kotaemon/flow/cache/memory.py vendored
View File

@@ -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]

15
libs/kotaemon/kotaemon/flow/callbacks.py
View File

@@ -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__}"

0
libs/kotaemon/kotaemon/flow/cli.py
View File

190
libs/kotaemon/kotaemon/flow/config.py
View File

@@ -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

166
libs/kotaemon/kotaemon/flow/context.py
View File

@@ -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

97
libs/kotaemon/kotaemon/flow/debug.py
View File

@@ -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

10
libs/kotaemon/kotaemon/flow/exceptions.py
View File

@@ -1,10 +0,0 @@
class InvalidAttrDefinition(AttributeError):
pass
class CyclicDependencyError(Exception):
pass
class CyclicPipelineError(Exception):
pass

215
libs/kotaemon/kotaemon/flow/middleware.py
View File

@@ -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)

0
libs/kotaemon/kotaemon/flow/runs/__init__.py
View File

193
libs/kotaemon/kotaemon/flow/runs/base.py
View File

@@ -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

70
libs/kotaemon/kotaemon/flow/safe.py
View File

@@ -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

22
libs/kotaemon/kotaemon/flow/settings/__init__.py
View File

@@ -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()

31
libs/kotaemon/kotaemon/flow/settings/default.py
View File

@@ -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",
}

8
libs/kotaemon/kotaemon/flow/storage/__init__.py
View File

@@ -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"]

33
libs/kotaemon/kotaemon/flow/storage/base.py
View File

@@ -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"""
...

35
libs/kotaemon/kotaemon/flow/storage/local.py
View File

@@ -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))

0
libs/kotaemon/kotaemon/flow/utils/__init__.py
View File

111
libs/kotaemon/kotaemon/flow/utils/documentation.py
View File

@@ -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()}

61
libs/kotaemon/kotaemon/flow/utils/hashes.py
View File

@@ -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()

284
libs/kotaemon/kotaemon/flow/utils/modules.py
View File

@@ -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__ == "<lambda>":
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)

45
libs/kotaemon/kotaemon/flow/utils/multiprocess.py
View File

@@ -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()

183
libs/kotaemon/kotaemon/flow/utils/paths.py
View File

@@ -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)

77
libs/kotaemon/kotaemon/flow/utils/pretties.py
View File

@@ -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

110
libs/kotaemon/kotaemon/flow/utils/typings.py
View File

@@ -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

269
libs/kotaemon/kotaemon/flow/visualization.py
View File

@@ -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