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

lionagi/__init__.py

@@ -27,7 +27,7 @@ from lionagi.core.action import func_to_tool
 from lionagi.core.report import Form, Report
 from lionagi.core.session.branch import Branch
 from lionagi.core.session.session import Session
-from lionagi.core.work.worker import work, Worker
+from lionagi.core.work.worker import work, Worker, worklink
 from lionagi.integrations.provider.services import Services
 from lionagi.integrations.chunker.chunk import chunk
 from lionagi.integrations.loader.load import load
@@ -41,6 +41,7 @@ __all__ = [
     "pile",
     "iModel",
     "work",
+    "worklink",
     "Worker",
     "Branch",
     "Session",

lionagi/core/generic/graph.py

@@ -44,6 +44,7 @@ class Graph(Node):
         condition: Condition | None = None,
         bundle=False,
         label=None,
+        edge_class=Edge,
         **kwargs,
     ):
         """Add an edge between two nodes in the graph."""
@@ -61,6 +62,7 @@ class Graph(Node):
             condition=condition,
             label=label,
             bundle=bundle,
+            edge_class=edge_class,
             **kwargs,
         )
 
@@ -184,19 +186,23 @@ class Graph(Node):
             node_info = node.to_dict()
             node_info.pop("ln_id")
             node_info.update({"class_name": node.class_name})
+            if hasattr(node, "name"):
+                node_info.update({"name": node.name})
             g.add_node(node.ln_id, **node_info)
 
         for _edge in self.internal_edges:
             edge_info = _edge.to_dict()
             edge_info.pop("ln_id")
             edge_info.update({"class_name": _edge.class_name})
+            if hasattr(_edge, "name"):
+                edge_info.update({"name": _edge.name})
             source_node_id = edge_info.pop("head")
             target_node_id = edge_info.pop("tail")
             g.add_edge(source_node_id, target_node_id, **edge_info)
 
         return g
 
-    def display(self, **kwargs):
+    def display(self, node_label="class_name", edge_label="label", draw_kwargs={}, **kwargs):
         """Display the graph using NetworkX and Matplotlib."""
         from lionagi.libs import SysUtil
 
@@ -217,10 +223,11 @@ class Graph(Node):
             node_size=500,
             node_color="orange",
             alpha=0.9,
-            labels=nx.get_node_attributes(g, "class_name"),
+            labels=nx.get_node_attributes(g, node_label),
+            **draw_kwargs
         )
 
-        labels = nx.get_edge_attributes(g, "label")
+        labels = nx.get_edge_attributes(g, edge_label)
         labels = {k: v for k, v in labels.items() if v}
 
         if labels:

lionagi/core/generic/node.py

@@ -10,6 +10,7 @@ modifying, and removing edges, and querying related nodes and connections.
 
 from pydantic import Field
 from pandas import Series
+from typing import Callable
 
 from lionagi.libs.ln_convert import to_list
 
@@ -122,6 +123,8 @@ class Node(Component, Relatable):
         condition: Condition | None = None,
         label: str | None = None,
         bundle: bool = False,
+        edge_class: Callable = Edge,
+        **kwargs
     ) -> None:
         """
         Establish directed relationship from this node to another.
@@ -141,12 +144,13 @@ class Node(Component, Relatable):
                 f"Invalid value for direction: {direction}, " "must be 'in' or 'out'"
             )
 
-        edge = Edge(
+        edge = edge_class(
             head=self if direction == "out" else node,
             tail=node if direction == "out" else self,
             condition=condition,
             bundle=bundle,
             label=label,
+            **kwargs
         )
 
         self.relations[direction].include(edge)

lionagi/core/report/base.py

@@ -26,6 +26,7 @@ from typing import Any, List, Dict
 import contextlib
 from lionagi.core.collections.abc import Component, Field
 from ..collections.util import to_list_type
+from lionagi.libs.ln_convert import to_str
 
 
 class BaseForm(Component):

lionagi/core/work/work_edge.py

@@ -0,0 +1,96 @@
+from typing import Callable
+from pydantic import Field, field_validator
+import inspect
+
+from lionagi.core.generic.edge import Edge
+from lionagi.core.collections.abc.concepts import Progressable
+
+from lionagi.core.work.worker import Worker
+
+
+class WorkEdge(Edge, Progressable):
+    """
+    Represents a directed edge between work tasks, responsible for transforming
+    the result of one task into parameters for the next task.
+
+    Attributes:
+        convert_function (Callable): Function to transform the result of the previous
+            work into parameters for the next work. This function must be decorated
+            with the `worklink` decorator.
+        convert_function_kwargs (dict): Additional parameters for the convert_function
+            other than "from_work" and "from_result".
+        associated_worker (Worker): The worker to which this WorkEdge belongs.
+    """
+    convert_function: Callable = Field(
+        ...,
+        description="Function to transform the result of the previous work into parameters for the next work."
+    )
+
+    convert_function_kwargs: dict = Field(
+        {},
+        description="parameters for the worklink function other than \"from_work\" and \"from_result\""
+    )
+
+    associated_worker: Worker = Field(
+        ...,
+        description="The worker to which this WorkEdge belongs."
+    )
+
+    @field_validator("convert_function", mode="before")
+    def _validate_convert_funuction(cls, func):
+        """
+        Validates that the convert_function is decorated with the worklink decorator.
+
+        Args:
+            func (Callable): The function to validate.
+
+        Returns:
+            Callable: The validated function.
+
+        Raises:
+            ValueError: If the function is not decorated with the worklink decorator.
+        """
+        try:
+            getattr(func, "_worklink_decorator_params")
+            return func
+        except:
+            raise ValueError("convert_function must be a worklink decorated function")
+
+    @property
+    def name(self):
+        """
+        Returns the name of the convert_function.
+
+        Returns:
+            str: The name of the convert_function.
+        """
+        return self.convert_function.__name__
+
+    async def forward(self, task):
+        """
+        Transforms the result of the current work into parameters for the next work
+        and schedules the next work task.
+
+        Args:
+            task (Task): The task to process.
+
+        Returns:
+            Work: The next work task to be executed.
+
+        Raises:
+            StopIteration: If the task has no available steps left to proceed.
+        """
+        if task.available_steps == 0:
+            task.status_note = ("Task stopped proceeding further as all available steps have been used up, "
+                                "but the task has not yet reached completion.")
+            return
+        func_signature = inspect.signature(self.convert_function)
+        kwargs = self.convert_function_kwargs.copy()
+        if "from_work" in func_signature.parameters:
+            kwargs = {"from_work": task.current_work} | kwargs
+        if "from_result" in func_signature.parameters:
+            kwargs = {"from_result": task.current_work.result} | kwargs
+
+        self.convert_function.auto_schedule = True
+        next_work = await self.convert_function(self=self.associated_worker, **kwargs)
+        return next_work

lionagi/core/work/work_function.py

@@ -13,14 +13,12 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 """
+import asyncio
 import logging
 
 from lionagi.libs.ln_func_call import rcall
-from lionagi.libs.ln_parse import ParseUtil
 from lionagi.core.work.worklog import WorkLog
 
-from lionagi.core.report.form import Form
-
 
 class WorkFunction:
     """
@@ -35,7 +33,7 @@ class WorkFunction:
     """
 
     def __init__(
-        self, assignment, function, retry_kwargs=None, guidance=None, capacity=10
+        self, assignment, function, retry_kwargs=None, guidance=None, capacity=10, refresh_time=1
     ):
         """
         Initializes a WorkFunction instance.
@@ -47,12 +45,15 @@ class WorkFunction:
                 Defaults to None.
             guidance (str, optional): The guidance or documentation for the function.
                 Defaults to None.
-            capacity (int, optional): The capacity of the work log. Defaults to 10.
+            capacity (int, optional): The capacity of the work queue batch processing.
+                Defaults to 10.
+            refresh_time (int, optional): The time interval to refresh the work log queue.
+                Defaults to 1.
         """
         self.assignment = assignment
         self.function = function
         self.retry_kwargs = retry_kwargs or {}
-        self.worklog = WorkLog(capacity)
+        self.worklog = WorkLog(capacity, refresh_time=refresh_time)
         self.guidance = guidance or self.function.__doc__
 
     @property
@@ -65,6 +66,16 @@ class WorkFunction:
         """
         return self.function.__name__
 
+    @property
+    def execution_mode(self):
+        """
+         Gets the execution mode of the work function's queue.
+
+         Returns:
+             bool: The execution mode of the work function's queue.
+         """
+        return self.worklog.queue.execution_mode
+
     def is_progressable(self):
         """
         Checks if the work function is progressable.
@@ -99,3 +110,15 @@ class WorkFunction:
         Process the first capacity_size works in the work queue.
         """
         await self.worklog.queue.process()
+
+    async def execute(self):
+        """
+        Starts the execution of the work function's queue.
+        """
+        asyncio.create_task(self.worklog.queue.execute())
+
+    async def stop(self):
+        """
+        Stops the execution of the work function's queue.
+        """
+        await self.worklog.stop()

lionagi/core/work/work_function_node.py

@@ -0,0 +1,44 @@
+from pydantic import Field
+import asyncio
+
+from lionagi.core.generic.node import Node
+from lionagi.core.work.work_function import WorkFunction
+from lionagi.core.collections import Exchange
+from lionagi.core.mail.mail import Mail
+
+
+class WorkFunctionNode(WorkFunction, Node):
+    """
+    A class representing a work function node, combining the functionality
+    of a WorkFunction and a Node.
+
+    Attributes:
+        assignment (str): The assignment description of the work function.
+        function (Callable): The function to be performed.
+        retry_kwargs (dict): The retry arguments for the function.
+        guidance (str): The guidance or documentation for the function.
+        capacity (int): The capacity of the work queue batch processing.
+        refresh_time (int): The time interval to refresh the work log queue.
+    """
+
+    def __init__(self, assignment, function, retry_kwargs=None, guidance=None, capacity=10, refresh_time=1, **kwargs):
+        """
+        Initializes a WorkFunctionNode instance.
+
+        Args:
+            assignment (str): The assignment description of the work function.
+            function (Callable): The function to be performed.
+            retry_kwargs (dict, optional): The retry arguments for the function. Defaults to None.
+            guidance (str, optional): The guidance or documentation for the function. Defaults to None.
+            capacity (int, optional): The capacity of the work queue batch processing. Defaults to 10.
+            refresh_time (int, optional): The time interval to refresh the work log queue. Defaults to 1.
+            **kwargs: Additional keyword arguments for the Node initialization.
+        """
+        Node.__init__(self, **kwargs)
+        WorkFunction.__init__(self,
+                              assignment=assignment,
+                              function=function,
+                              retry_kwargs=retry_kwargs,
+                              guidance=guidance,
+                              capacity=capacity,
+                              refresh_time=refresh_time)

lionagi/core/work/work_queue.py

@@ -25,19 +25,33 @@ class WorkQueue:
     Attributes:
         capacity (int): The maximum number of tasks the queue can handle.
         queue (asyncio.Queue): The queue holding the tasks.
-        _stop_event (asyncio.Event): Event to signal stopping `execute` of the queue.
+        _stop_event (asyncio.Event): Event to signal stopping the execution of the queue.
         available_capacity (int): The remaining number of tasks the queue can handle.
         execution_mode (bool): If `execute` is running.
+        refresh_time (int): The time interval between task processing.
     """
 
-    def __init__(self, capacity=5):
+    def __init__(self, capacity=5, refresh_time=1):
+        """
+        Initializes a new instance of WorkQueue.
+
+        Args:
+            capacity (int): The maximum number of tasks the queue can handle.
+            refresh_time (int): The time interval between task processing.
+
+        Raises:
+            ValueError: If capacity is less than 0 or refresh_time is negative.
+        """
         if capacity < 0:
             raise ValueError("initial capacity must be >= 0")
+        if refresh_time < 0:
+            raise ValueError("refresh time for execution can not be negative")
         self.capacity = capacity
         self.queue = asyncio.Queue()
         self._stop_event = asyncio.Event()
         self.available_capacity = capacity
         self.execution_mode = False
+        self.refresh_time = refresh_time
 
     async def enqueue(self, work) -> None:
         """Enqueue a work item."""
@@ -60,20 +74,6 @@ class WorkQueue:
         """Return whether the queue has been stopped."""
         return self._stop_event.is_set()
 
-    # async def process(self):
-        # async def _parse_work(work):
-        #     async with self.semaphore:
-        #         await work.perform()
-        #
-        # tasks = set()
-        # while self.queue.qsize() > 0:
-        #     next = await self.dequeue()
-        #     next.status = WorkStatus.IN_PROGRESS
-        #     task = asyncio.create_task(_parse_work(next))
-        #     tasks.add(task)
-        #
-        # await asyncio.wait(tasks)
-
     async def process(self) -> None:
         """Process the work items in the queue."""
         tasks = set()
@@ -88,7 +88,7 @@ class WorkQueue:
             await asyncio.wait(tasks)
             self.available_capacity = self.capacity
 
-    async def execute(self, refresh_time=1):
+    async def execute(self):
         """
             Continuously executes the process method at a specified refresh interval.
 
@@ -97,7 +97,9 @@ class WorkQueue:
                     successive calls to `process`. Defaults to 1.
         """
         self.execution_mode = True
+        self._stop_event.clear()
+
         while not self.stopped:
             await self.process()
-            await asyncio.sleep(refresh_time)
+            await asyncio.sleep(self.refresh_time)
         self.execution_mode = False

lionagi/core/work/work_task.py

@@ -0,0 +1,155 @@
+import inspect
+from typing import Any, Callable
+from pydantic import Field, field_validator, model_validator
+
+from lionagi.core.collections.abc.component import Component
+from lionagi.core.work.work import WorkStatus, Work
+from collections.abc import Coroutine
+
+
+class WorkTask(Component):
+    """
+    A class representing a work task that can be processed in multiple steps.
+
+    Attributes:
+        name (str | None): The name of the task.
+        status (WorkStatus): The current status of the task.
+        status_note (str): A note for the task's current status.
+        work_history (list[Work]): A list of works processed in this task.
+        max_steps (int | None): The maximum number of works allowed in this task.
+        current_work (Work | None): The current work in progress.
+        post_processing (Callable | None): The post-processing function to be executed after the entire task is successfully completed.
+    """
+    name: str | None = Field(
+        None,
+        description="Name of the task"
+    )
+
+    status: WorkStatus = Field(
+        WorkStatus.PENDING,
+        description="The current status of the task"
+    )
+
+    status_note: str = Field(
+        None,
+        description="Note for tasks current status"
+    )
+
+    work_history: list[Work] = Field(
+        [],
+        description="List of works processed"
+    )
+
+    max_steps: int | None = Field(
+        10,
+        description="Maximum number of works allowed"
+    )
+
+    current_work: Work | None = Field(
+        None,
+        description="The current work in progress"
+    )
+
+    post_processing: Callable | None = Field(
+        None,
+        description="The post-processing function to be executed after the entire task has been successfully completed."
+    )
+
+    @field_validator("max_steps", mode="before")
+    def _validate_max_steps(cls, value):
+        """
+        Validates that max_steps is a positive integer.
+
+        Args:
+            value (int): The value to validate.
+
+        Returns:
+            int: The validated value.
+
+        Raises:
+            ValueError: If value is not a positive integer.
+        """
+        if value <= 0:
+            raise ValueError("Invalid value: max_steps must be a positive integer.")
+        return value
+
+    @field_validator("post_processing", mode="before")
+    def _validate_prost_processing(cls, value):
+        """
+        Validates that post_processing is an asynchronous function.
+
+        Args:
+            value (Callable): The value to validate.
+
+        Returns:
+            Callable: The validated value.
+
+        Raises:
+            ValueError: If value is not an asynchronous function.
+        """
+        if value is not None and not inspect.iscoroutinefunction((value)):
+            raise ValueError("post_processing must be a async function")
+        return value
+
+    @property
+    def available_steps(self):
+        """
+        Calculates the number of available steps left in the task.
+
+        Returns:
+            int: The number of available steps.
+        """
+        return max(0, self.max_steps - len(self.work_history))
+
+    def clone(self):
+        """
+        Creates a clone of the current WorkTask instance.
+
+        Returns:
+            WorkTask: A new instance of WorkTask with the same attributes.
+        """
+        new_worktask = WorkTask(name=self.name, status=self.status, max_steps=self.max_steps, current_work=self.current_work)
+        for work in self.work_history:
+            new_worktask.work_history.append(work)
+        return new_worktask
+
+    async def process(self, current_func_node):
+        """
+        Processes the current work function node.
+
+        Args:
+            current_func_node (WorkFunctionNode): The current function node being processed.
+
+        Returns:
+            str | list[WorkTask]: Returns "COMPLETED", "FAILED", or a list of new WorkTask instances if there are next works to process.
+        """
+        if self.current_work.status == WorkStatus.FAILED:
+            self.status = WorkStatus.FAILED
+            self.status_note = f"Work {self.current_work.ln_id} failed. Error: {self.current_work.error}"
+            self.current_work = None
+            return "FAILED"
+        elif self.current_work.status == WorkStatus.COMPLETED:
+            next_works = []
+            if not current_func_node.relations["out"].is_empty():
+                for workedge in current_func_node.relations["out"]:
+                    next_work = await workedge.forward(self)
+                    if next_work is not None:
+                        next_works.append(next_work)
+            if len(next_works) == 0:
+                self.current_work = None
+                self.status = WorkStatus.COMPLETED
+                if self.post_processing:
+                    await self.post_processing(self)
+                return "COMPLETED"
+            else:
+                return_tasks = []
+                for i in reversed(range(len(next_works))):
+                    if i == 0:
+                        self.current_work = next_works[i]
+                        self.work_history.append(next_works[i])
+                    else:
+                        clone_task = self.clone()
+                        clone_task.current_work = next_works[i]
+                        clone_task.work_history.append(next_works[i])
+                        return_tasks.append(clone_task)
+                return return_tasks