lionagi 0.2.6__py3-none-any.whl → 0.2.8__py3-none-any.whl

lionagi/core/collections/pile.py

@@ -1,50 +1,43 @@
-"""
-Copyright 2024 HaiyangLi
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-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.
-"""
-
-"""
-This module defines the Pile class, a versatile container for managing
-collections of Element objects. It supports structured access and
-manipulation, including retrieval, addition, and deletion of elements.
-"""
+from __future__ import annotations
 
+import asyncio
 from collections.abc import Iterable
-from typing import TypeVar, Type, Any, Generic
+from functools import wraps
+from typing import Any, AsyncIterator, Callable, Generic, Type, TypeVar
 
 from pydantic import Field, field_validator
 
-from lionagi.libs.ln_convert import is_same_dtype, to_df
-from lionagi.libs.ln_func_call import bcall, alcall, CallDecorator as cd
+from lionagi.libs.ln_convert import is_same_dtype, to_df, to_dict
+from lionagi.libs.ln_func_call import CallDecorator as cd
+from lionagi.libs.ln_func_call import alcall
+
 from .abc import (
-    Element,
-    Record,
     Component,
-    Ordering,
+    Element,
+    ItemNotFoundError,
     LionIDable,
-    get_lion_id,
-    LionValueError,
     LionTypeError,
-    ItemNotFoundError,
+    LionValueError,
     ModelLimitExceededError,
+    Ordering,
+    Record,
+    get_lion_id,
 )
 from .model import iModel
-from .util import to_list_type, _validate_order
+from .util import _validate_order, to_list_type
 
 T = TypeVar("T")
 
 
+def async_synchronized(func: Callable):
+    @wraps(func)
+    async def wrapper(self, *args, **kwargs):
+        async with self.async_lock:
+            return await func(self, *args, **kwargs)
+
+    return wrapper
+
+
 class Pile(Element, Record, Generic[T]):
     """
     Collection class for managing Element objects.
@@ -70,6 +63,14 @@ class Pile(Element, Record, Generic[T]):
     query_response: list = []
     tools: dict = {}
 
+    def __pydantic_extra__(self) -> dict[str, Any]:
+        return {
+            "_async": Field(default_factory=asyncio.Lock),
+        }
+
+    def __pydantic_private__(self) -> dict[str, Any]:
+        return self.__pydantic_extra__()
+
     def __init__(
         self,
         items=None,
@@ -442,6 +443,103 @@ class Pile(Element, Record, Generic[T]):
     def __radd__(self, other: T) -> "Pile":
         return other + self
 
+    def __ior__(self, other: Any | Pile) -> Pile:
+        if not isinstance(other, Pile):
+            raise LionTypeError(
+                "Invalid type for Pile operation.",
+                expected_type=Pile,
+                actual_type=type(other),
+            )
+        other = self._validate_pile(list(other))
+        self.include(other)
+        return self
+
+    def __or__(self, other: Any | Pile) -> Pile:
+        if not isinstance(other, Pile):
+            raise LionTypeError(
+                "Invalid type for Pile operation.",
+                expected_type=Pile,
+                actual_type=type(other),
+            )
+
+        result = self.__class__(
+            items=self.values(),
+            item_type=self.item_type,
+            order=self.order,
+        )
+        result.include(list(other))
+        return result
+
+    def __ixor__(self, other: Any | Pile) -> Pile:
+        if not isinstance(other, Pile):
+            raise LionTypeError(
+                "Invalid type for Pile operation.",
+                expected_type=Pile,
+                actual_type=type(other),
+            )
+
+        to_exclude = []
+        for i in other:
+            if i in self:
+                to_exclude.append(i)
+
+        other = [i for i in other if i not in to_exclude]
+        self.exclude(to_exclude)
+        self.include(other)
+        return self
+
+    def __xor__(self, other: Any | Pile) -> Pile:
+        if not isinstance(other, Pile):
+            raise LionTypeError(
+                "Invalid type for Pile operation.",
+                expected_type=Pile,
+                actual_type=type(other),
+            )
+
+        to_exclude = []
+        for i in other:
+            if i in self:
+                to_exclude.append(i)
+
+        values = [i for i in self if i not in to_exclude] + [
+            i for i in other if i not in to_exclude
+        ]
+
+        result = self.__class__(
+            items=values,
+            item_type=self.item_type,
+        )
+        return result
+
+    def __iand__(self, other: Any) -> Pile:
+        if not isinstance(other, Pile):
+            raise LionTypeError(
+                "Invalid type for Pile operation.",
+                expected_type=Pile,
+                actual_type=type(other),
+            )
+
+        to_exclude = []
+        for i in self.values():
+            if i not in other:
+                to_exclude.append(i)
+        self.exclude(to_exclude)
+        return self
+
+    def __and__(self, other: Any | Pile) -> Pile:
+        if not isinstance(other, Pile):
+            raise LionTypeError(
+                "Invalid type for Pile operation.",
+                expected_type=Pile,
+                actual_type=type(other),
+            )
+
+        values = [i for i in self if i in other]
+        return self.__class__(
+            items=values,
+            item_type=self.item_type,
+        )
+
     def size(self) -> int:
         """Return the total size of the pile."""
         return sum([len(i) for i in self])
@@ -666,7 +764,7 @@ class Pile(Element, Record, Generic[T]):
         else:
             raise ValueError("Invalid index type")
 
-    async def query_pile(self, query, engine_kwargs={}, **kwargs):
+    async def query_pile(self, query, engine_kwargs={}, return_dict=False, **kwargs):
         """
         Query the pile using the created query engine.
 
@@ -682,9 +780,11 @@ class Pile(Element, Record, Generic[T]):
             self.create_query_engine(**engine_kwargs)
         response = await self.engines["query"].aquery(query, **kwargs)
         self.query_response.append(response)
+        if return_dict:
+            return to_dict(response)
         return str(response)
 
-    async def chat_pile(self, query, engine_kwargs={}, **kwargs):
+    async def chat_pile(self, query, engine_kwargs={}, return_dict=False, **kwargs):
         """
         Chat with the pile using the created chat engine.
 
@@ -700,6 +800,8 @@ class Pile(Element, Record, Generic[T]):
             self.create_chat_engine(**engine_kwargs)
         response = await self.engines["chat"].achat(query, **kwargs)
         self.query_response.append(response)
+        if return_dict:
+            return to_dict(response)
         return str(response)
 
     async def embed_pile(
@@ -792,6 +894,7 @@ class Pile(Element, Record, Generic[T]):
         name=None,
         guidance=None,
         query_description=None,
+        return_dict=False,
         **kwargs,
     ):
         """
@@ -830,10 +933,10 @@ class Pile(Element, Record, Generic[T]):
 
         async def query(query: str):
             if query_type == "query":
-                return await self.query_pile(query, **kwargs)
+                return await self.query_pile(query, return_dict=return_dict, **kwargs)
 
             elif query_type == "chat":
-                return await self.chat_pile(query, **kwargs)
+                return await self.chat_pile(query, return_dict=return_dict, **kwargs)
 
         name = name or "query"
         tool = func_to_tool(query)[0]
@@ -872,6 +975,185 @@ class Pile(Element, Record, Generic[T]):
         """
         return self.to_df().__repr__()
 
+    def __getstate__(self):
+        """Prepare the Pile instance for pickling."""
+        state = self.__dict__.copy()
+        state["_async_lock"] = None
+        return state
+
+    def __setstate__(self, state):
+        """Restore the Pile instance after unpickling."""
+        self.__dict__.update(state)
+        self._async_lock = asyncio.Lock()
+
+    @property
+    def async_lock(self):
+        """Ensure the async lock is always available, even during unpickling"""
+        if not hasattr(self, "_async_lock") or self._async_lock is None:
+            self._async_lock = asyncio.Lock()
+        return self._async_lock
+
+    # Async Interface methods
+    @async_synchronized
+    async def asetitem(
+        self,
+        key: Any,
+        item: T | Iterable[T],
+        /,
+    ) -> None:
+        """Asynchronously set an item or items in the Pile.
+
+        Args:
+            key: The key to set. Can be an integer index, a string ID, or a
+                slice.
+            item: The item or items to set. Must be of type T or an iterable
+                of T for slices.
+
+        Raises:
+            TypeError: If the item type is not allowed.
+            KeyError: If the key is invalid.
+            ValueError: If trying to set multiple items with a non-slice key.
+        """
+        self._setitem(key, item)
+
+    @async_synchronized
+    async def apop(
+        self,
+        key: Any,
+        default: Any = ...,
+        /,
+    ):
+        """Asynchronously remove and return an item or items from the Pile.
+
+        Args:
+            key: The key of the item(s) to remove. Can be an integer index,
+                a string ID, or a slice.
+            default: The value to return if the key is not found. Defaults to
+                ....
+
+        Returns:
+            The removed item(s), or the default value if not found.
+
+        Raises:
+            KeyError: If the key is not found and no default is provided.
+        """
+        return self._pop(key, default)
+
+    @async_synchronized
+    async def aremove(
+        self,
+        item: T,
+        /,
+    ) -> None:
+        """Asynchronously remove a specific item from the Pile.
+
+        Args:
+            item: The item to remove.
+
+        Raises:
+            ValueError: If the item is not found in the Pile.
+        """
+        self._remove(item)
+
+    @async_synchronized
+    async def ainclude(
+        self,
+        item: T | Iterable[T],
+        /,
+    ) -> None:
+        """Asynchronously include item(s) in the Pile if not already present.
+
+        Args:
+            item: Item or iterable of items to include.
+
+        Raises:
+            TypeError: If the item(s) are not of allowed types.
+        """
+        self._include(item)
+        if item not in self:
+            raise LionTypeError(f"Item {item} is not of allowed types")
+
+    @async_synchronized
+    async def aexclude(
+        self,
+        item: T | Iterable[T],
+        /,
+    ) -> None:
+        """Asynchronously exclude item(s) from the Pile if present.
+
+        Args:
+            item: Item or iterable of items to exclude.
+
+        Note:
+            This method does not raise an error if an item is not found.
+        """
+        self._exclude(item)
+
+    @async_synchronized
+    async def aclear(self) -> None:
+        self._clear()
+
+    @async_synchronized
+    async def aupdate(
+        self,
+        other: Any,
+        /,
+    ) -> None:
+        self._update(other)
+
+    @async_synchronized
+    async def aget(
+        self,
+        key: Any,
+        default=...,
+        /,
+    ) -> list | Any | T:
+        return self._get(key, default)
+
+    async def __aiter__(self) -> AsyncIterator[T]:
+        """Return an asynchronous iterator over the items in the Pile.
+
+        This method creates a snapshot of the current order to prevent
+        issues with concurrent modifications during iteration.
+
+        Yields:
+            Items in the Pile in their current order.
+
+        Note:
+            This method yields control to the event loop after each item,
+            allowing other async operations to run between iterations.
+        """
+
+        async with self.async_lock:
+            current_order = list(self.order)
+
+        for key in current_order:
+            yield self.pile_[key]
+            await asyncio.sleep(0)  # Yield control to the event loop
+
+    async def __anext__(self) -> T:
+        """Asynchronously return the next item in the Pile."""
+        try:
+            return await anext(self.AsyncPileIterator(self))
+        except StopAsyncIteration:
+            raise StopAsyncIteration("End of pile")
+
+    class AsyncPileIterator:
+        def __init__(self, pile: Pile):
+            self.pile = pile
+            self.index = 0
+
+        def __aiter__(self) -> AsyncIterator[T]:
+            return self
+
+        async def __anext__(self) -> T:
+            if self.index >= len(self.pile):
+                raise StopAsyncIteration
+            item = self.pile[self.pile.order[self.index]]
+            self.index += 1
+            await asyncio.sleep(0)  # Yield control to the event loop
+            return item
+
 
 def pile(
     items: Iterable[T] | None = None,

lionagi/core/collections/progression.py

@@ -1,23 +1,10 @@
-"""
-Copyright 2024 HaiyangLi
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
+import contextlib
 
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-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.
-"""
+from pydantic import Field, field_validator
 
-import contextlib
 from lionagi.libs import SysUtil
-from pydantic import Field, field_validator
-from .abc import Ordering, get_lion_id, ItemNotFoundError, LionIDable, Element
+
+from .abc import Element, ItemNotFoundError, LionIDable, Ordering, get_lion_id
 from .util import _validate_order
 
 

lionagi/core/collections/util.py

@@ -1,7 +1,7 @@
-from collections.abc import Mapping, Generator
 from collections import deque
+from collections.abc import Generator, Mapping
 
-from .abc import LionTypeError, Record, Ordering, Component, get_lion_id, Element
+from .abc import Component, Element, LionTypeError, Ordering, Record, get_lion_id
 
 
 def to_list_type(value):

lionagi/core/director/README.md

@@ -0,0 +1 @@
+director is an agent specialized in writing directives

lionagi/core/director/direct.py

@@ -1,21 +1,5 @@
-"""
-Copyright 2024 HaiyangLi
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-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.
-"""
-
-from lionagi.core.unit.unit import Unit
 from lionagi.core.session.branch import Branch
+from lionagi.core.unit.unit import Unit
 
 
 async def chat(

lionagi/core/engine/branch_engine.py

@@ -1,13 +1,13 @@
 import contextlib
-from lionagi.libs import convert, AsyncUtil, ParseUtil
-from lionagi.core.generic.edge import Edge
+
 from lionagi.core.action import ActionNode
-from lionagi.core.mail.mail import Mail
-from lionagi.core.message import System, Instruction
 from lionagi.core.collections import Pile, Progression
-
-from lionagi.core.session.branch import Branch
 from lionagi.core.executor.base_executor import BaseExecutor
+from lionagi.core.generic.edge import Edge
+from lionagi.core.mail.mail import Mail
+from lionagi.core.message import Instruction, System
+from lionagi.core.session.branch import Branch
+from lionagi.libs import AsyncUtil, ParseUtil, convert
 
 
 class BranchExecutor(Branch, BaseExecutor):

lionagi/core/engine/instruction_map_engine.py

@@ -1,12 +1,12 @@
 import asyncio
+
 from pydantic import Field
 
+from lionagi.core.collections import Exchange, Pile, pile, progression
+from lionagi.core.engine.branch_engine import BranchExecutor
+from lionagi.core.executor.base_executor import BaseExecutor
 from lionagi.core.mail.mail import Mail, Package
-from lionagi.core.collections import Exchange
 from lionagi.core.mail.mail_manager import MailManager
-from lionagi.core.executor.base_executor import BaseExecutor
-from lionagi.core.engine.branch_engine import BranchExecutor
-from lionagi.core.collections import progression, pile, Pile
 
 
 class InstructionMapEngine(BaseExecutor):

lionagi/core/engine/script_engine.py

@@ -1,22 +1,8 @@
-"""
-Copyright 2024 HaiyangLi
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-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 ast
 from functools import lru_cache
+
 from lionagi.libs import AsyncUtil
+
 from ..evaluator.base_evaluator import BaseEvaluator
 from .sandbox_ import SandboxTransformer
 

lionagi/core/executor/base_executor.py

@@ -3,8 +3,8 @@ from typing import Any
 
 from pydantic import Field
 
-from lionagi.core.collections.abc import Element, Progressable, Executable
 from lionagi.core.collections import Exchange
+from lionagi.core.collections.abc import Element, Executable, Progressable
 from lionagi.core.mail.mail import Mail, Package
 
 

lionagi/core/executor/graph_executor.py

@@ -1,16 +1,13 @@
 from collections import deque
 
-from lionagi.libs import AsyncUtil, convert
-
-from lionagi.core.generic.node import Node
-from lionagi.core.generic.edge import Edge
+from lionagi.core.action import ActionNode, DirectiveSelection, Tool
+from lionagi.core.collections.progression import progression
 from lionagi.core.executor.base_executor import BaseExecutor
-
-from lionagi.core.action import Tool, DirectiveSelection, ActionNode
-
-from lionagi.core.mail import Mail
+from lionagi.core.generic.edge import Edge
 from lionagi.core.generic.graph import Graph
-from lionagi.core.collections.progression import progression
+from lionagi.core.generic.node import Node
+from lionagi.core.mail import Mail
+from lionagi.libs import AsyncUtil, convert
 
 
 class GraphExecutor(BaseExecutor, Graph):

lionagi/core/executor/neo4j_executor.py

@@ -1,18 +1,16 @@
-from collections import deque
 import json
+from collections import deque
 from typing import Callable
 
-from lionagi.core.executor.base_executor import BaseExecutor
-from lionagi.integrations.storage.neo4j import Neo4j
-from lionagi.integrations.storage.storage_util import ParseNode
+from lionagi.core.action import ActionNode, DirectiveSelection, Tool
 from lionagi.core.agent.base_agent import BaseAgent
+from lionagi.core.collections.progression import progression
 from lionagi.core.engine.instruction_map_engine import InstructionMapEngine
-
-from lionagi.core.mail import Mail
-from lionagi.core.action import Tool, DirectiveSelection, ActionNode
+from lionagi.core.executor.base_executor import BaseExecutor
 from lionagi.core.generic.edge import Edge
-from lionagi.core.collections.progression import progression
-
+from lionagi.core.mail import Mail
+from lionagi.integrations.storage.neo4j import Neo4j
+from lionagi.integrations.storage.storage_util import ParseNode
 from lionagi.libs import AsyncUtil
 
 

lionagi/core/generic/__init__.py

@@ -1,7 +1,6 @@
 from .edge import Edge
-from .node import Node
 from .graph import Graph
+from .node import Node
 from .tree import Tree
 
-
 __all__ = ["Edge", "Node", "Graph", "Tree"]

lionagi/core/generic/edge.py

@@ -1,6 +1,8 @@
-from pydantic import Field, field_validator
 from typing import Any
-from lionagi.core.collections.abc import Component, get_lion_id, LionIDable, Condition
+
+from pydantic import Field, field_validator
+
+from lionagi.core.collections.abc import Component, Condition, LionIDable, get_lion_id
 from lionagi.core.generic.edge_condition import EdgeCondition
 
 
@@ -70,7 +72,8 @@ class Edge(Component):
         if self.condition is None:
             return
 
-        import inspect, sys
+        import inspect
+        import sys
 
         def new_getfile(object, _old_getfile=inspect.getfile):
             if not inspect.isclass(object):
@@ -96,6 +99,7 @@ class Edge(Component):
         inspect.getfile = new_getfile
 
         import inspect
+
         from IPython.core.magics.code import extract_symbols
 
         obj = self.condition.__class__

lionagi/core/generic/edge_condition.py

@@ -1,7 +1,8 @@
 from typing import Any
-from pydantic import Field
+
+from pydantic import BaseModel, Field
+
 from lionagi.core.collections.abc import Condition
-from pydantic import BaseModel
 
 
 class EdgeCondition(Condition, BaseModel):