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

lionagi/core/work/worker.py

@@ -16,11 +16,12 @@ limitations under the License.
 
 from abc import ABC
 from functools import wraps
-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
+from lionagi.core.collections.abc import get_lion_id
 
 
 class Worker(ABC):
@@ -30,19 +31,30 @@ class Worker(ABC):
     Attributes:
         name (str): The name of the worker.
         work_functions (dict[str, WorkFunction]): Dictionary mapping assignments to WorkFunction objects.
+        forms (dict[str, Form]): Dictionary mapping form identifier to Form objects.
+        default_form (str|None): The default form to be used by the worker.
     """
 
     name: str = "Worker"
-    work_functions: dict[str, WorkFunction] = {}
 
-    def __init__(self) -> None:
-        self.stopped = False
+    def __init__(self, forms=None, default_form=None) -> None:
+        """
+        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):
         """
         Stops the worker and all associated work functions.
         """
-        self.stopped = True
+        # self.stopped = True
         _logging.info(f"Stopping worker {self.name}")
         non_stopped_ = []
 
@@ -69,36 +81,73 @@ class Worker(ABC):
             and not self.stopped
         )
 
-    async def process(self, refresh_time=1):
+    async def change_default_form(self, form_key):
         """
-        Processes all work functions periodically.
+        Changes the default form to the specified form key.
 
         Args:
-            refresh_time (int): Time interval between each process cycle.
+            form_key (str): The key of the form to set as the default.
+
+        Raises:
+            ValueError: If the form key does not exist in the forms dictionary.
+
+        """
+        if form_key not in self.forms.keys():
+            raise ValueError(f"Unable to change default form. Key {form_key} does not exist.")
+        self.default_form = self.forms[form_key]
+
+    def _get_decorated_functions(self, decorator_attr, name_only=True):
         """
-        while await self.is_progressable():
-            await pcall([i.process(refresh_time) for i in self.work_functions.values()])
-            await asyncio.sleep(refresh_time)
+        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.
 
-    # TODO: Implement process method
+        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}")
 
-    # 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)
+    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 _wrapper(
+    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,
     ):
         """
@@ -107,47 +156,90 @@ class Worker(ABC):
         Args:
             func (Callable): The function to be executed.
             assignment (str): The assignment description.
-            capacity (int): Capacity for the work log.
+            form_param_key (str): The key to identify the form parameter in
+                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 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 or func.__doc__,
+                guidance=guidance or function.__doc__,
                 capacity=capacity,
+                refresh_time=refresh_time
             )
 
-        work_func: WorkFunction = self.work_functions[func.__name__]
-        task = asyncio.create_task(work_func.perform(self, *args, **kwargs))
-        work = Work(async_task=task)
+        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(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:
+                bound_args = func_signature.bind(None, *args, **kwargs)
+            else:
+                bound_args = func_signature.bind(*args, **kwargs)
+            bound_args.apply_defaults()
+            arguments = bound_args.arguments
+
+            form_key = arguments.get(form_param_key)
+            try:
+                form_key = get_lion_id(form_key)
+            except:
+                pass
+            form = self.forms.get(form_key) or self.default_form
+
+            if form:
+                subform = form.__class__(assignment=work_func.assignment, task=work_func.guidance)
+                for k in subform.input_fields:
+                    v = getattr(form, k, None)
+                    setattr(subform, k, v)
+                subform.origin = form
+                kwargs = {"form": subform} | kwargs
+            else:
+                raise ValueError(f"Cannot locate form in Worker's forms and default_form is not available.")
+
+        task = work_func.perform(self, *args, **kwargs)
+        work = Work(async_task=task, async_task_name=work_func.name)
         await work_func.worklog.append(work)
-        return True
+        return work
 
 
 def work(
     assignment=None,
+    form_param_key=None,
     capacity=10,
     guidance=None,
     retry_kwargs=None,
-    refresh_time=1,
     timeout=10,
+    refresh_time=1
 ):
     """
     Decorator to mark a method as a work function.
 
     Args:
         assignment (str): The assignment description of the work function.
-        capacity (int): Capacity for the work log.
+        form_param_key (str): The key to identify the form parameter in
+                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 queue batch processing.
         guidance (str): Guidance or documentation for the work function.
         retry_kwargs (dict): Retry arguments for the work function.
-        refresh_time (int): Time interval between each process cycle.
         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):
@@ -157,22 +249,118 @@ def work(
             *args,
             func=func,
             assignment=assignment,
+            form_param_key=form_param_key,
             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):
         """
@@ -56,15 +58,11 @@ class WorkLog(Progressable):
 
     async def forward(self):
         """
-        Forwards pending work items to the queue if capacity allows.
+        Forwards pending work items to the queue.
         """
-        if not self.queue.available_capacity:
-            return
-        else:
-            while len(self.pending) > 0 and self.queue.available_capacity:
-                work: Work = self.pile[self.pending.popleft()]
-                work.status = WorkStatus.IN_PROGRESS
-                await self.queue.enqueue(work)
+        while len(self.pending) > 0:
+            work: Work = self.pile[self.pending.popleft()]
+            await self.queue.enqueue(work)
 
     async def stop(self):
         """