lionagi 0.2.1__py3-none-any.whl → 0.2.2__py3-none-any.whl

lionagi/core/work/worker.py

@@ -16,11 +16,8 @@ limitations under the License.
 
 from abc import ABC
 from functools import wraps
-from typing import Callable
-import asyncio
 import inspect
 from lionagi import logging as _logging
-from lionagi.libs.ln_func_call import pcall
 from lionagi.core.work.work_function import WorkFunction
 from lionagi.core.work.work import Work
 from lionagi.core.report.form import Form
@@ -39,12 +36,19 @@ class Worker(ABC):
     """
 
     name: str = "Worker"
-    work_functions: dict[str, WorkFunction] = {}
 
     def __init__(self, forms=None, default_form=None) -> None:
-        # self.stopped = False
+        """
+        Initializes a new instance of Worker.
+
+        Args:
+            forms (dict[str, Form], optional): Dictionary mapping form identifier to Form objects.
+            default_form (str|None, optional): The default form to be used by the worker.
+        """
+        self.work_functions: dict[str, WorkFunction] = {}
         self.forms: dict[str, Form] = forms or {}
         self.default_form = default_form
+        self._validate_worklink_functions()
 
     async def stop(self):
         """
@@ -92,37 +96,58 @@ class Worker(ABC):
             raise ValueError(f"Unable to change default form. Key {form_key} does not exist.")
         self.default_form = self.forms[form_key]
 
-    # async def process(self, refresh_time=1):
-    #     """
-    #     Processes all work functions periodically.
-    #
-    #     Args:
-    #         refresh_time (int): Time interval between each process cycle.
-    #     """
-    #     while await self.is_progressable():
-    #         await pcall([i.process(refresh_time) for i in self.work_functions.values()])
-    #         await asyncio.sleep(refresh_time)
-
-    # TODO: Implement process method
-
-    # async def process(self, refresh_time=1):
-    #     while not self.stopped:
-    #         tasks = [
-    #             asyncio.create_task(func.process(refresh_time=refresh_time))
-    #             for func in self.work_functions.values()
-    #         ]
-    #         await asyncio.wait(tasks)
-    #         await asyncio.sleep(refresh_time)
-
-    async def _wrapper(
+    def _get_decorated_functions(self, decorator_attr, name_only=True):
+        """
+        Retrieves decorated functions based on the specified decorator attribute.
+
+        Args:
+            decorator_attr (str): The attribute name of the decorator.
+            name_only (bool, optional): Whether to return only the function names. Defaults to True.
+
+        Returns:
+            list: List of decorated function names or tuples containing function details.
+        """
+        decorated_functions = []
+        for name, func in inspect.getmembers(self.__class__, predicate=inspect.isfunction):
+            if hasattr(func, decorator_attr):
+                if name_only:
+                    decorated_functions.append(name)
+                else:
+                    decorator_params = getattr(func, decorator_attr)
+                    decorated_functions.append((name, func, decorator_params))
+        return decorated_functions
+
+    def _validate_worklink_functions(self):
+        """
+        Validates worklink functions to ensure they have the required parameters.
+        """
+        worklink_decorated_function = self._get_decorated_functions(decorator_attr="_worklink_decorator_params", name_only=False)
+        for func_name, func, _ in worklink_decorated_function:
+            func_signature = inspect.signature(func)
+            if "from_work" not in func_signature.parameters and "from_result" not in func_signature.parameters:
+                raise ValueError(f"Either \"from_work\" or \"from_result\" must be a parameter in function {func_name}")
+
+    def construct_all_work_functions(self):
+        """
+        Constructs all work functions for the worker.
+        """
+        if getattr(self, "work_functions", None) is None:
+            self.work_functions = {}
+        work_decorated_function = self._get_decorated_functions(decorator_attr="_work_decorator_params", name_only=False)
+        for func_name, func, dec_params in work_decorated_function:
+            if func_name not in self.work_functions:
+                self.work_functions[func_name] = WorkFunction(**dec_params)
+
+    async def _work_wrapper(
         self,
         *args,
-        func=None,
+        function=None,
         assignment=None,
         form_param_key=None,
         capacity=None,
         retry_kwargs=None,
         guidance=None,
+        refresh_time=1,
         **kwargs,
     ):
         """
@@ -135,27 +160,28 @@ class Worker(ABC):
                 the function's signature. This parameter is used to locate and fill
                 the appropriate form according to the assignment. Raises an error
                 if the form parameter key is not found in the function's signature.
-            capacity (int): Capacity for the work log.
+            capacity (int): Capacity for the work queue batch processing.
             retry_kwargs (dict): Retry arguments for the function.
             guidance (str): Guidance or documentation for the function.
         """
         if getattr(self, "work_functions", None) is None:
             self.work_functions = {}
 
-        if func.__name__ not in self.work_functions:
-            self.work_functions[func.__name__] = WorkFunction(
+        if function.__name__ not in self.work_functions:
+            self.work_functions[function.__name__] = WorkFunction(
                 assignment=assignment,
-                function=func,
+                function=function,
                 retry_kwargs=retry_kwargs or {},
-                guidance=guidance,
+                guidance=guidance or function.__doc__,
                 capacity=capacity,
+                refresh_time=refresh_time
             )
 
-        work_func: WorkFunction = self.work_functions[func.__name__]
+        work_func: WorkFunction = self.work_functions[function.__name__]
 
         # locate form that should be filled according to the assignment
         if form_param_key:
-            func_signature = inspect.signature(func)
+            func_signature = inspect.signature(function)
             if form_param_key not in func_signature.parameters:
                 raise KeyError(f"Failed to locate form. \"{form_param_key}\" is not defined in the function.")
             if "self" in func_signature.parameters:
@@ -195,6 +221,7 @@ def work(
     guidance=None,
     retry_kwargs=None,
     timeout=10,
+    refresh_time=1
 ):
     """
     Decorator to mark a method as a work function.
@@ -205,10 +232,14 @@ def work(
                 the function's signature. This parameter is used to locate and fill
                 the appropriate form according to the assignment. Raises an error
                 if the form parameter key is not found in the function's signature.
-        capacity (int): Capacity for the work log.
+        capacity (int): Capacity for the work queue batch processing.
         guidance (str): Guidance or documentation for the work function.
         retry_kwargs (dict): Retry arguments for the work function.
         timeout (int): Timeout for the work function.
+        refresh_time (int, optional): Refresh time for the work log queue.
+
+    Returns:
+        Callable: The decorated function.
     """
 
     def decorator(func):
@@ -222,20 +253,114 @@ def work(
             capacity=capacity,
             retry_kwargs=retry_kwargs,
             guidance=guidance,
+            refresh_time=refresh_time,
             **kwargs,
         ):
+            if not inspect.iscoroutinefunction(func):
+                raise TypeError(f"{func.__name__} must be an asynchronous function")
             retry_kwargs = retry_kwargs or {}
             retry_kwargs["timeout"] = retry_kwargs.get("timeout", timeout)
-            return await self._wrapper(
+            return await self._work_wrapper(
                 *args,
-                func=func,
+                function=func,
                 assignment=assignment,
                 form_param_key=form_param_key,
                 capacity=capacity,
                 retry_kwargs=retry_kwargs,
                 guidance=guidance,
+                refresh_time=refresh_time,
                 **kwargs,
             )
+        wrapper._work_decorator_params = {"assignment": assignment,
+                                          "function": func,
+                                          "retry_kwargs": retry_kwargs,
+                                          "guidance": guidance,
+                                          "capacity": capacity,
+                                          "refresh_time": refresh_time}
+
+        return wrapper
+
+    return decorator
+
+
+def worklink(
+        from_: str,
+        to_: str,
+        auto_schedule: bool = True
+):
+    """
+    Decorator to create a link between two work functions.
+
+    Args:
+        from_ (str): The name of the source work function.
+        to_ (str): The name of the target work function.
+        auto_schedule (bool, optional): Whether to automatically schedule the next work function. Defaults to True.
+
+    Returns:
+        Callable: The decorated function.
+    """
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(
+            self: Worker,
+            *args,
+            func=func,
+            from_=from_,
+            to_=to_,
+            **kwargs
+        ):
+            if not inspect.iscoroutinefunction(func):
+                raise TypeError(f"{func.__name__} must be an asynchronous function")
+
+            work_funcs = self._get_decorated_functions(decorator_attr="_work_decorator_params")
+            if from_ not in work_funcs or to_ not in work_funcs:
+                raise ValueError("Invalid link. 'from_' and 'to_' must be the name of work decorated functions.")
+
+            func_signature = inspect.signature(func)
+            if "from_work" not in func_signature.parameters and "from_result" not in func_signature.parameters:
+                raise ValueError(f"Either \"from_work\" or \"from_result\" must be a parameter in function {func.__name__}")
+
+            if "self" in func_signature.parameters:
+                bound_args = func_signature.bind(None, *args, **kwargs)
+            else:
+                bound_args = func_signature.bind(*args, **kwargs)
+            bound_args.apply_defaults()
+            arguments = bound_args.arguments
+            if "kwargs" in arguments:
+                arguments.update(arguments.pop("kwargs"))
+
+            if from_work := arguments.get("from_work"):
+                if not isinstance(from_work, Work):
+                    raise ValueError("Invalid type for from_work. Only work objects are accepted.")
+                if from_work.async_task_name != from_:
+                    raise ValueError(f"Invalid work object in from_work. "
+                                     f"async_task_name \"{from_work.async_task_name}\" does not match from_ \"{from_}\"")
+
+            next_params = await func(self, *args, **kwargs)
+            to_work_func = getattr(self, to_)
+            if next_params is None:
+                return
+            if isinstance(next_params, list):
+                if wrapper.auto_schedule:
+                    return await to_work_func(*next_params)
+            elif isinstance(next_params, dict):
+                if wrapper.auto_schedule:
+                    return await to_work_func(**next_params)
+            elif isinstance(next_params, tuple) and len(next_params) == 2:
+                if isinstance(next_params[0], list) and isinstance(next_params[1], dict):
+                    if wrapper.auto_schedule:
+                        return await to_work_func(*next_params[0], **next_params[1])
+                else:
+                    raise TypeError(f"Invalid return type {func.__name__}")
+            else:
+                raise TypeError(f"Invalid return type {func.__name__}")
+
+            return next_params
+
+        wrapper.auto_schedule = auto_schedule
+        wrapper._worklink_decorator_params = {"func": func,
+                                              "from_": from_,
+                                              "to_": to_}
 
         return wrapper
 

lionagi/core/work/worker_engine.py

@@ -0,0 +1,179 @@
+import asyncio
+from lionagi.core.work.work import WorkStatus
+from lionagi.core.work.worker import Worker
+from lionagi.core.work.work_task import WorkTask
+from lionagi.core.work.work_edge import WorkEdge
+from lionagi.core.work.work_function_node import WorkFunctionNode
+
+from lionagi.core.collections.pile import pile
+from lionagi.core.generic.graph import Graph
+
+
+class WorkerEngine:
+    """
+    A class representing an engine that manages and executes work tasks for a worker.
+
+    Attributes:
+        worker (Worker): The worker instance that the engine manages.
+        tasks (Pile): A pile of tasks to be executed.
+        active_tasks (Pile): A pile of currently active tasks.
+        failed_tasks (Pile): A pile of tasks that have failed.
+        worker_graph (Graph): A graph representing the relationships between work functions.
+        refresh_time (int): The time interval for refreshing the work log queue.
+        _stop_event (asyncio.Event): An event to signal stopping the execution.
+    """
+
+    def __init__(self, worker: Worker, refresh_time=1):
+        """
+        Initializes a new instance of WorkerEngine.
+
+        Args:
+            worker (Worker): The worker instance to be managed.
+            refresh_time (int): The time interval for refreshing the work log queue.
+        """
+        self.worker = worker
+        self.tasks = pile()
+        self.active_tasks = pile()
+        self.failed_tasks = pile()
+        self.worker_graph = Graph()
+        self._construct_work_functions()
+        self._construct_workedges()
+        self.refresh_time = refresh_time
+        self._stop_event = asyncio.Event()
+
+    async def add_task(self, *args, task_function: str, task_name=None, task_max_steps=10, task_post_processing=None, **kwargs):
+        """
+        Adds a new task to the task queue.
+
+        Args:
+            task_function (str): The name of the task function to execute.
+            task_name (str, optional): The name of the task.
+            task_max_steps (int, optional): The maximum number of steps for the task.
+            task_post_processing (Callable, optional): The post-processing function for the task.
+            *args: Positional arguments for the task function.
+            **kwargs: Keyword arguments for the task function.
+
+        Returns:
+            WorkTask: The newly created task.
+        """
+        task = WorkTask(name=task_name, max_steps=task_max_steps, post_processing=task_post_processing)
+        self.tasks.append(task)
+        function = getattr(self.worker, task_function)
+        work = await function(*args, **kwargs)
+        task.current_work = work
+        task.work_history.append(work)
+        return task
+
+    async def activate_work_queues(self):
+        """
+                Activates the work queues for all work functions.
+        """
+        for work_function in self.worker.work_functions.values():
+            if not work_function.worklog.queue.execution_mode:
+                asyncio.create_task(work_function.worklog.queue.execute())
+
+    async def stop(self):
+        """
+        Stops the execution of tasks.
+        """
+        self._stop_event.set()
+
+    @property
+    def stopped(self) -> bool:
+        """
+        Checks if the execution has been stopped.
+
+        Returns:
+            bool: True if stopped, otherwise False.
+        """
+        return self._stop_event.is_set()
+
+    async def process(self, task: WorkTask):
+        """
+        Processes a single task.
+
+        Args:
+            task (WorkTask): The task to be processed.
+        """
+        current_work_func_name = task.current_work.async_task_name
+        current_work_func = self.worker.work_functions[current_work_func_name]
+        updated_task = await task.process(current_work_func)
+        if updated_task == "COMPLETED":
+            self.active_tasks.pop(task)
+        elif updated_task == "FAILED":
+            self.active_tasks.pop(task)
+            self.failed_tasks.append(task)
+        elif isinstance(updated_task, list):
+            self.tasks.include(updated_task)
+            self.active_tasks.include(updated_task)
+        else:
+            await asyncio.sleep(self.refresh_time)
+
+    async def execute(self, stop_queue=True):
+        """
+        Executes all tasks in the task queue.
+
+        Args:
+            stop_queue (bool, optional): Whether to stop the queue after execution. Defaults to True.
+        """
+        for task in self.tasks:
+            if task.status == WorkStatus.PENDING:
+                task.status = WorkStatus.IN_PROGRESS
+                self.active_tasks.append(task)
+
+        await self.activate_work_queues()
+        self._stop_event.clear()
+
+        while len(self.active_tasks) > 0 and not self.stopped:
+            for work_function in self.worker.work_functions.values():
+                if len(work_function.worklog.pending) > 0:
+                    await work_function.worklog.forward()
+            tasks = list(self.active_tasks)
+            await asyncio.gather(*[self.process(task) for task in tasks])
+            await asyncio.sleep(self.refresh_time)
+
+        if stop_queue:
+            await self.stop()
+            await self.worker.stop()
+
+    async def execute_lasting(self):
+        """
+        Executes tasks continuously until stopped.
+        """
+        self._stop_event.clear()
+        
+        async def execute_lasting_inner():
+            while not self.stopped:
+                await self.execute(stop_queue=False)
+                await asyncio.sleep(self.refresh_time)
+        asyncio.create_task(execute_lasting_inner())
+
+    def _construct_work_functions(self):
+        """
+        Constructs work functions for the worker.
+        """
+        if getattr(self.worker, "work_functions", None) is None:
+            self.worker.work_functions = {}
+        work_decorated_function = self.worker._get_decorated_functions(decorator_attr="_work_decorator_params",
+                                                                       name_only=False)
+        for func_name, func, dec_params in work_decorated_function:
+            if func_name not in self.worker.work_functions:
+                self.worker.work_functions[func_name] = WorkFunctionNode(**dec_params)
+                self.worker_graph.add_node(self.worker.work_functions[func_name])
+            else:
+                if not isinstance(self.worker.work_functions[func_name], WorkFunctionNode):
+                    raise TypeError(f"WorkFunction {func_name} already exists but is not a WorkFunctionNode. "
+                                    f"If you would like to use it in WorkerEngine, please convert it to a "
+                                    f"WorkFunctionNode, or initiate a new worker, or pop it from work_function dict")
+
+    def _construct_workedges(self):
+        """
+        Constructs work edges for the worker graph.
+        """
+        worklink_decorated_function = self.worker._get_decorated_functions(decorator_attr="_worklink_decorator_params",
+                                                                           name_only=False)
+
+        for func_name, func, dec_params in worklink_decorated_function:
+            head = self.worker.work_functions[dec_params["from_"]]
+            tail = self.worker.work_functions[dec_params["to_"]]
+            self.worker_graph.add_edge(head=head, tail=tail, convert_function=func, associated_worker=self.worker, edge_class=WorkEdge)

lionagi/core/work/worklog.py

@@ -30,19 +30,21 @@ class WorkLog(Progressable):
         queue (WorkQueue): A queue to manage the execution of work items.
     """
 
-    def __init__(self, capacity=10, workpile=None):
+    def __init__(self, capacity=10, workpile=None, refresh_time=1):
         """
         Initializes a new instance of WorkLog.
 
         Args:
-            capacity (int): The capacity of the work queue.
+            capacity (int): The capacity of the work queue batch processing.
             workpile (Pile, optional): An optional pile of initial work items.
+            refresh_time (int, optional): The time interval to refresh the work log queue.
+                Defaults to 1.
         """
         self.pile = (
             workpile if workpile and isinstance(workpile, Pile) else pile({}, Work)
         )
         self.pending = progression(workpile) if workpile else progression()
-        self.queue = WorkQueue(capacity=capacity)
+        self.queue = WorkQueue(capacity=capacity, refresh_time=refresh_time)
 
     async def append(self, work: Work):
         """