lionagi 0.13.7__py3-none-any.whl → 0.14.0__py3-none-any.whl

lionagi/libs/concurrency/task.py

@@ -0,0 +1,109 @@
+"""Task group implementation for structured concurrency."""
+
+from collections.abc import Awaitable, Callable
+from types import TracebackType
+from typing import Any, Optional, TypeVar
+
+import anyio
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+class TaskGroup:
+    """A group of tasks that are treated as a unit."""
+
+    def __init__(self):
+        """Initialize a new task group."""
+        self._task_group = None
+
+    async def start_soon(
+        self,
+        func: Callable[..., Awaitable[Any]],
+        *args: Any,
+        name: str | None = None,
+    ) -> None:
+        """Start a new task in this task group.
+
+        Args:
+            func: The coroutine function to call
+            *args: Positional arguments to pass to the function
+            name: Optional name for the task
+
+        Note:
+            This method does not wait for the task to initialize.
+        """
+        if self._task_group is None:
+            raise RuntimeError("Task group is not active")
+        self._task_group.start_soon(func, *args, name=name)
+
+    async def start(
+        self,
+        func: Callable[..., Awaitable[R]],
+        *args: Any,
+        name: str | None = None,
+    ) -> R:
+        """Start a new task and wait for it to initialize.
+
+        Args:
+            func: The coroutine function to call
+            *args: Positional arguments to pass to the function
+            name: Optional name for the task
+
+        Returns:
+            The value passed to task_status.started()
+
+        Note:
+            The function must accept a task_status keyword argument and call
+            task_status.started() once initialization is complete.
+        """
+        if self._task_group is None:
+            raise RuntimeError("Task group is not active")
+        return await self._task_group.start(func, *args, name=name)
+
+    async def __aenter__(self) -> "TaskGroup":
+        """Enter the task group context.
+
+        Returns:
+            The task group instance.
+        """
+        task_group = anyio.create_task_group()
+        self._task_group = await task_group.__aenter__()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> bool:
+        """Exit the task group context.
+
+        This will wait for all tasks in the group to complete.
+        If any task raised an exception, it will be propagated.
+        If multiple tasks raised exceptions, they will be combined into an ExceptionGroup.
+
+        Returns:
+            True if the exception was handled, False otherwise.
+        """
+        if self._task_group is None:
+            return False
+
+        try:
+            return await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
+        finally:
+            self._task_group = None
+
+
+def create_task_group() -> TaskGroup:
+    """Create a new task group.
+
+    Returns:
+        A new task group instance.
+
+    Example:
+        async with create_task_group() as tg:
+            await tg.start_soon(task1)
+            await tg.start_soon(task2)
+    """
+    return TaskGroup()

lionagi/operations/builder.py

@@ -75,6 +75,7 @@ class OperationGraphBuilder:
         operation: BranchOperations,
         node_id: str | None = None,
         depends_on: list[str] | None = None,
+        inherit_context: bool = False,
         **parameters,
     ) -> str:
         """
@@ -84,6 +85,8 @@ class OperationGraphBuilder:
             operation: The branch operation
             node_id: Optional ID reference for this node
             depends_on: List of node IDs this depends on
+            inherit_context: If True and has dependencies, inherit conversation
+                           context from primary (first) dependency
             **parameters: Operation parameters
 
         Returns:
@@ -92,6 +95,11 @@ class OperationGraphBuilder:
         # Create operation node
         node = Operation(operation=operation, parameters=parameters)
 
+        # Store context inheritance strategy
+        if inherit_context and depends_on:
+            node.metadata["inherit_context"] = True
+            node.metadata["primary_dependency"] = depends_on[0]
+
         self.graph.add_node(node)
         self._operations[node.id] = node
 
@@ -126,6 +134,8 @@ class OperationGraphBuilder:
         source_node_id: str,
         operation: BranchOperations,
         strategy: ExpansionStrategy = ExpansionStrategy.CONCURRENT,
+        inherit_context: bool = False,
+        chain_context: bool = False,
         **shared_params,
     ) -> list[str]:
         """
@@ -139,6 +149,9 @@ class OperationGraphBuilder:
             source_node_id: ID of node that produced these items
             operation: Operation to apply to each item
             strategy: How to organize the expanded operations
+            inherit_context: If True, expanded operations inherit context from source
+            chain_context: If True and strategy is SEQUENTIAL, each operation
+                          inherits from the previous (only applies to SEQUENTIAL)
             **shared_params: Shared parameters for all operations
 
         Returns:
@@ -171,6 +184,21 @@ class OperationGraphBuilder:
                 },
             )
 
+            # Handle context inheritance for expanded operations
+            if inherit_context:
+                if (
+                    chain_context
+                    and strategy == ExpansionStrategy.SEQUENTIAL
+                    and i > 0
+                ):
+                    # Chain context: inherit from previous expanded operation
+                    node.metadata["inherit_context"] = True
+                    node.metadata["primary_dependency"] = new_node_ids[i - 1]
+                else:
+                    # Inherit from source node
+                    node.metadata["inherit_context"] = True
+                    node.metadata["primary_dependency"] = source_node_id
+
             self.graph.add_node(node)
             self._operations[node.id] = node
             new_node_ids.append(node.id)
@@ -195,7 +223,10 @@ class OperationGraphBuilder:
     def add_aggregation(
         self,
         operation: BranchOperations,
+        node_id: str | None = None,
         source_node_ids: list[str] | None = None,
+        inherit_context: bool = False,
+        inherit_from_source: int = 0,
         **parameters,
     ) -> str:
         """
@@ -203,7 +234,10 @@ class OperationGraphBuilder:
 
         Args:
             operation: Aggregation operation
+            node_id: Optional ID reference for this node
             source_node_ids: Nodes to aggregate from (defaults to current heads)
+            inherit_context: If True, inherit conversation context from one source
+            inherit_from_source: Index of source to inherit context from (default: 0)
             **parameters: Operation parameters
 
         Returns:
@@ -226,6 +260,18 @@ class OperationGraphBuilder:
             metadata={"aggregation": True},
         )
 
+        # Store node reference if provided
+        if node_id:
+            node.metadata["reference_id"] = node_id
+
+        # Store context inheritance for aggregations
+        if inherit_context and sources:
+            node.metadata["inherit_context"] = True
+            # Use the specified source index (bounded by available sources)
+            source_idx = min(inherit_from_source, len(sources) - 1)
+            node.metadata["primary_dependency"] = sources[source_idx]
+            node.metadata["inherit_from_source"] = source_idx
+
         self.graph.add_node(node)
         self._operations[node.id] = node