lionagi 0.2.11__py3-none-any.whl → 0.3.1__py3-none-any.whl

lionagi/core/generic/graph.py

@@ -24,7 +24,11 @@ class Graph(Node):
     def internal_edges(self) -> Pile[Edge]:
         """Return a pile of all edges in the graph."""
         return pile(
-            {edge.ln_id: edge for node in self.internal_nodes for edge in node.edges},
+            {
+                edge.ln_id: edge
+                for node in self.internal_nodes
+                for edge in node.edges
+            },
             Edge,
         )
 
@@ -70,7 +74,9 @@ class Graph(Node):
         edge = edge if isinstance(edge, list) else [edge]
         for i in edge:
             if i not in self.internal_edges:
-                raise ItemNotFoundError(f"Edge {i} does not exist in structure.")
+                raise ItemNotFoundError(
+                    f"Edge {i} does not exist in structure."
+                )
             with contextlib.suppress(ItemNotFoundError):
                 self._remove_edge(i)
 
@@ -105,7 +111,8 @@ class Graph(Node):
                     [
                         edge
                         for edge in edges
-                        if edge.label in to_list(label, dropna=True, flatten=True)
+                        if edge.label
+                        in to_list(label, dropna=True, flatten=True)
                     ]
                 )
                 if edges
@@ -124,7 +131,9 @@ class Graph(Node):
     def _remove_edge(self, edge: Edge | str) -> bool:
         """Remove a specific edge from the graph."""
         if edge not in self.internal_edges:
-            raise ItemNotFoundError(f"Edge {edge} does not exist in structure.")
+            raise ItemNotFoundError(
+                f"Edge {edge} does not exist in structure."
+            )
 
         edge = self.internal_edges[edge]
         head: Node = self.internal_nodes[edge.head]
@@ -139,7 +148,8 @@ class Graph(Node):
             [
                 node
                 for node in self.internal_nodes
-                if node.relations["in"].is_empty() and not isinstance(node, Actionable)
+                if node.relations["in"].is_empty()
+                and not isinstance(node, Actionable)
             ]
         )
 
@@ -147,7 +157,9 @@ class Graph(Node):
         """Check if the graph is acyclic (contains no cycles)."""
         node_ids = list(self.internal_nodes.keys())
         check_deque = deque(node_ids)
-        check_dict = {key: 0 for key in node_ids}  # 0: not visited, 1: temp, 2: perm
+        check_dict = {
+            key: 0 for key in node_ids
+        }  # 0: not visited, 1: temp, 2: perm
 
         def visit(key):
             if check_dict[key] == 2:
@@ -202,7 +214,11 @@ class Graph(Node):
         return g
 
     def display(
-        self, node_label="class_name", edge_label="label", draw_kwargs={}, **kwargs
+        self,
+        node_label="class_name",
+        edge_label="label",
+        draw_kwargs={},
+        **kwargs,
     ):
         """Display the graph using NetworkX and Matplotlib."""
         from lionagi.libs import SysUtil

lionagi/core/generic/node.py

@@ -8,7 +8,7 @@ Includes functionality for managing relationships, such as adding,
 modifying, and removing edges, and querying related nodes and connections.
 """
 
-from typing import Callable
+from collections.abc import Callable
 
 from pandas import Series
 from pydantic import Field
@@ -62,7 +62,11 @@ class Node(Component, Relatable):
             List of node IDs related to this node.
         """
         all_nodes = set(
-            to_list([[i.head, i.tail] for i in self.edges], flatten=True, dropna=True)
+            to_list(
+                [[i.head, i.tail] for i in self.edges],
+                flatten=True,
+                dropna=True,
+            )
         )
         all_nodes.discard(self.ln_id)
         return list(all_nodes)
@@ -101,7 +105,9 @@ class Node(Component, Relatable):
             List of node IDs that precede this node.
         """
         return [
-            node_id for node_id, edges in self.node_relations["in"].items() if edges
+            node_id
+            for node_id, edges in self.node_relations["in"].items()
+            if edges
         ]
 
     @property
@@ -113,7 +119,9 @@ class Node(Component, Relatable):
             List of node IDs that succeed this node.
         """
         return [
-            node_id for node_id, edges in self.node_relations["out"].items() if edges
+            node_id
+            for node_id, edges in self.node_relations["out"].items()
+            if edges
         ]
 
     def relate(
@@ -141,7 +149,8 @@ class Node(Component, Relatable):
         """
         if direction not in ["in", "out"]:
             raise ValueError(
-                f"Invalid value for direction: {direction}, " "must be 'in' or 'out'"
+                f"Invalid value for direction: {direction}, "
+                "must be 'in' or 'out'"
             )
 
         edge = edge_class(

lionagi/core/generic/tree_node.py

@@ -28,7 +28,11 @@ class TreeNode(Node):
         if not self.parent:
             return list(self.related_nodes)
         else:
-            return [node for node in self.related_nodes if node != self.parent.ln_id]
+            return [
+                node
+                for node in self.related_nodes
+                if node != self.parent.ln_id
+            ]
 
     def relate_child(
         self,

lionagi/core/mail/mail_manager.py

@@ -126,7 +126,9 @@ class MailManager(Element, Executable):
             mail_id = mailbox.pending_outs.popleft()
             mail = mailbox.pile.pop(mail_id)
             if mail.recipient not in self.sources:
-                raise ValueError(f"Recipient source {mail.recipient} does not exist")
+                raise ValueError(
+                    f"Recipient source {mail.recipient} does not exist"
+                )
             if mail.sender not in self.mails[mail.recipient]:
                 self.mails[mail.recipient].update({mail.sender: deque()})
             self.mails[mail.recipient][mail.sender].append(mail)

lionagi/core/mail/package.py

@@ -44,4 +44,6 @@ class Package(Element):
             try:
                 return PackageCategory(value)
             except Exception as e:
-                raise ValueError(f"Invalid value for category: {value}.") from e
+                raise ValueError(
+                    f"Invalid value for category: {value}."
+                ) from e

lionagi/core/message/action_request.py

@@ -48,14 +48,21 @@ class ActionRequest(RoledMessage):
             sender (str, optional): The sender of the request.
             recipient (str, optional): The recipient of the request.
         """
-        function = function.__name__ if inspect.isfunction(function) else function
+        function = (
+            function.__name__ if inspect.isfunction(function) else function
+        )
         arguments = _prepare_arguments(arguments)
 
         super().__init__(
             role=MessageRole.ASSISTANT,
             sender=sender,
             recipient=recipient,
-            content={"action_request": {"function": function, "arguments": arguments}},
+            content={
+                "action_request": {
+                    "function": function,
+                    "arguments": arguments,
+                }
+            },
             **kwargs,
         )
         self.function = function

lionagi/core/message/action_response.py

@@ -25,8 +25,12 @@ class ActionResponse(RoledMessage):
         description="The id of the action request that this response corresponds to",
     )
 
-    function: str | None = Field(None, description="The name of the function called")
-    arguments: dict | None = Field(None, description="The keyword arguments provided")
+    function: str | None = Field(
+        None, description="The name of the function called"
+    )
+    arguments: dict | None = Field(
+        None, description="The keyword arguments provided"
+    )
     func_outputs: Any | None = Field(
         None, description="The output of the function call"
     )
@@ -108,7 +112,9 @@ class ActionResponse(RoledMessage):
         action_request = ActionRequest(
             function=self.function, arguments=json.loads(arguments)
         )
-        action_response_copy = ActionResponse(action_request=action_request, **kwargs)
+        action_response_copy = ActionResponse(
+            action_request=action_request, **kwargs
+        )
         action_response_copy.action_request = self.action_request
         action_response_copy.func_outputs = self.func_outputs
         action_response_copy.metadata["origin_ln_id"] = self.ln_id

lionagi/core/message/instruction.py

@@ -109,7 +109,11 @@ class Instruction(RoledMessage):
         self, context, requested_fields, images, image_detail, **kwargs
     ):
         if context:
-            context = {"context": context} if not isinstance(context, dict) else context
+            context = (
+                {"context": context}
+                if not isinstance(context, dict)
+                else context
+            )
             if (
                 additional_context := {
                     k: v for k, v in kwargs.items() if k not in SYSTEM_FIELDS
@@ -124,7 +128,9 @@ class Instruction(RoledMessage):
             )
 
         if images:
-            self.content["images"] = images if isinstance(images, list) else [images]
+            self.content["images"] = (
+                images if isinstance(images, list) else [images]
+            )
             self.content["image_detail"] = image_detail
 
     def clone(self, **kwargs):

lionagi/core/message/util.py

@@ -81,7 +81,9 @@ def create_message(
 
     if function:
         if not arguments:
-            raise ValueError("Error: please provide arguments for the function.")
+            raise ValueError(
+                "Error: please provide arguments for the function."
+            )
         return ActionRequest(
             function=function,
             arguments=arguments,
@@ -158,7 +160,9 @@ def _parse_action_request(response):
         content_ = message["content"]["tool_uses"]
 
     else:
-        json_block_pattern = re.compile(r"```json\n({.*?tool_uses.*?})\n```", re.DOTALL)
+        json_block_pattern = re.compile(
+            r"```json\n({.*?tool_uses.*?})\n```", re.DOTALL
+        )
 
         # Find the JSON block in the text
         match = json_block_pattern.search(str(message["content"]))
@@ -179,7 +183,9 @@ def _parse_action_request(response):
         outs = []
         for func_calling in content_:
             if "recipient_name" in func_calling:
-                func_calling["action"] = func_calling["recipient_name"].split(".")[1]
+                func_calling["action"] = func_calling["recipient_name"].split(
+                    "."
+                )[1]
                 func_calling["arguments"] = func_calling["parameters"]
             elif "function" in func_calling:
                 func_calling["action"] = func_calling["function"]
@@ -212,9 +218,13 @@ def _parse_action_request(response):
                     if "function" in func_calling:
                         func_calling["action"] = func_calling["function"]
                         if "parameters" in func_calling:
-                            func_calling["arguments"] = func_calling["parameters"]
+                            func_calling["arguments"] = func_calling[
+                                "parameters"
+                            ]
                         elif "arguments" in func_calling:
-                            func_calling["arguments"] = func_calling["arguments"]
+                            func_calling["arguments"] = func_calling[
+                                "arguments"
+                            ]
                     msg = ActionRequest(
                         function=func_calling["action"]
                         .replace("action_", "")

lionagi/core/report/base.py

@@ -71,12 +71,12 @@ class BaseForm(Component):
         examples=["input1, input2 -> output"],
     )
 
-    input_fields: List[str] = Field(
+    input_fields: list[str] = Field(
         default_factory=list,
         description="Fields required to carry out the objective of the form.",
     )
 
-    requested_fields: List[str] = Field(
+    requested_fields: list[str] = Field(
         default_factory=list,
         description="Fields requested to be filled by the user.",
     )
@@ -86,14 +86,14 @@ class BaseForm(Component):
         description="The work to be done by the form, including custom instructions.",
     )
 
-    validation_kwargs: Dict[str, Dict[str, Any]] = Field(
+    validation_kwargs: dict[str, dict[str, Any]] = Field(
         default_factory=dict,
         description="Additional validation constraints for the form fields.",
         examples=[{"field": {"config1": "a", "config2": "b"}}],
     )
 
     @property
-    def work_fields(self) -> Dict[str, Any]:
+    def work_fields(self) -> dict[str, Any]:
         """
         Get the fields relevant to the current task, including input and
         requested fields. Must be implemented by subclasses.
@@ -169,8 +169,8 @@ class BaseForm(Component):
         return True
 
     def _get_all_fields(
-        self, form: List["BaseForm"] = None, **kwargs
-    ) -> Dict[str, Any]:
+        self, form: list["BaseForm"] = None, **kwargs
+    ) -> dict[str, Any]:
         """
         Given a form or collections of forms, and additional fields, gather
         all fields together including self fields with valid value.
@@ -187,7 +187,12 @@ class BaseForm(Component):
         all_form_fields = (
             {}
             if not form
-            else {k: v for i in form for k, v in i.work_fields.items() if v is not None}
+            else {
+                k: v
+                for i in form
+                for k, v in i.work_fields.items()
+                if v is not None
+            }
         )
         all_fields.update({**all_form_fields, **kwargs})
         return all_fields

lionagi/core/report/form.py

@@ -84,7 +84,7 @@ class Form(BaseForm):
                 )
 
     @property
-    def work_fields(self) -> Dict[str, Any]:
+    def work_fields(self) -> dict[str, Any]:
         """
         Retrieves a dictionary of the fields relevant to the current task,
         excluding any SYSTEM_FIELDS and including only the input and requested
@@ -97,7 +97,8 @@ class Form(BaseForm):
         return {
             k: v
             for k, v in dict_.items()
-            if k not in SYSTEM_FIELDS and k in self.input_fields + self.requested_fields
+            if k not in SYSTEM_FIELDS
+            and k in self.input_fields + self.requested_fields
         }
 
     def fill(self, form: "Form" = None, strict: bool = True, **kwargs) -> None:
@@ -134,7 +135,9 @@ class Form(BaseForm):
             bool: True if the form is workable, otherwise raises ValueError.
         """
         if self.filled:
-            raise ValueError("Form is already filled, cannot be worked on again")
+            raise ValueError(
+                "Form is already filled, cannot be worked on again"
+            )
 
         for i in self.input_fields:
             if not getattr(self, i, None):
@@ -172,7 +175,7 @@ class Form(BaseForm):
         """
 
     @property
-    def _instruction_requested_fields(self) -> Dict[str, str]:
+    def _instruction_requested_fields(self) -> dict[str, str]:
         """
         Provides a dictionary mapping requested field names to their
         descriptions.

lionagi/core/report/report.py

@@ -25,7 +25,7 @@ class Report(BaseForm):
         examples=[["a, b -> c", "a -> e", "b -> f", "c -> g", "e, f, g -> h"]],
     )
 
-    form_template: Type[Form] = Field(
+    form_template: type[Form] = Field(
         Form, description="The template for the forms in the report."
     )
 
@@ -72,7 +72,12 @@ class Report(BaseForm):
                     all_fields[k] = v
         return all_fields
 
-    def fill(self, form: Form | list[Form] | dict[Form] = None, strict=True, **kwargs):
+    def fill(
+        self,
+        form: Form | list[Form] | dict[Form] = None,
+        strict=True,
+        **kwargs,
+    ):
         if self.filled:
             if strict:
                 raise ValueError("Form is filled, cannot be worked on again")
@@ -105,7 +110,9 @@ class Report(BaseForm):
             bool: True if the report is workable, otherwise raises ValueError.
         """
         if self.filled:
-            raise ValueError("Form is already filled, cannot be worked on again")
+            raise ValueError(
+                "Form is already filled, cannot be worked on again"
+            )
 
         for i in self.input_fields:
             if not getattr(self, i, None):

lionagi/core/report/util.py

@@ -3,7 +3,9 @@ def get_input_output_fields(str_: str) -> list[list[str]]:
         return [], []
 
     if "->" not in str_:
-        raise ValueError("Invalid assignment format. Expected 'inputs -> outputs'.")
+        raise ValueError(
+            "Invalid assignment format. Expected 'inputs -> outputs'."
+        )
 
     inputs, outputs = str_.split("->")
 

lionagi/core/rule/action.py

@@ -49,7 +49,10 @@ class ActionRequestRule(MappingRule):
         Raises:
             ActionError: If the action request is invalid.
         """
-        if isinstance(value, dict) and list(value.keys()) >= ["function", "arguments"]:
+        if isinstance(value, dict) and list(value.keys()) >= [
+            "function",
+            "arguments",
+        ]:
             return value
         raise ActionError(f"Invalid action request: {value}")
 

lionagi/core/rule/base.py

@@ -3,7 +3,12 @@ from typing import Any, Dict, List
 
 from pandas import Series
 
-from lionagi.core.collections.abc import Actionable, Component, Condition, FieldError
+from lionagi.core.collections.abc import (
+    Actionable,
+    Component,
+    Condition,
+    FieldError,
+)
 from lionagi.libs import SysUtil
 
 _rule_classes = {}
@@ -37,7 +42,9 @@ class Rule(Component, Condition, Actionable):
         if cls.__name__ not in _rule_classes:
             _rule_classes[cls.__name__] = cls
 
-    def add_log(self, field: str, form: Any, apply: bool = True, **kwargs) -> None:
+    def add_log(
+        self, field: str, form: Any, apply: bool = True, **kwargs
+    ) -> None:
         """
         Adds an entry to the applied or invoked log.
 
@@ -67,7 +74,7 @@ class Rule(Component, Condition, Actionable):
         value: Any,
         form: Any,
         *args,
-        annotation: List[str] = None,
+        annotation: list[str] = None,
         use_annotation: bool = True,
         **kwargs,
     ) -> bool:
@@ -93,7 +100,9 @@ class Rule(Component, Condition, Actionable):
 
         if use_annotation:
             annotation = annotation or form._get_field_annotation(field)
-            annotation = [annotation] if isinstance(annotation, str) else annotation
+            annotation = (
+                [annotation] if isinstance(annotation, str) else annotation
+            )
 
             for i in annotation:
                 if i in self.apply_type and i not in self.exclude_type:
@@ -132,7 +141,9 @@ class Rule(Component, Condition, Actionable):
             if self.fix:
                 try:
                     a = await self.perform_fix(value, **self.validation_kwargs)
-                    self.add_log(field, form, apply=False, **self.validation_kwargs)
+                    self.add_log(
+                        field, form, apply=False, **self.validation_kwargs
+                    )
                     return a
                 except Exception as e2:
                     raise FieldError(f"failed to fix field") from e2
@@ -184,7 +195,7 @@ class Rule(Component, Condition, Actionable):
         """
         pass
 
-    def _to_dict(self) -> Dict[str, Any]:
+    def _to_dict(self) -> dict[str, Any]:
         """
         Converts the rule's attributes to a dictionary.
 

lionagi/core/rule/rulebook.py

@@ -3,12 +3,12 @@ from lionfuncs import lcall
 from lionagi.core.rule.base import Rule
 
 """
-rule config schema 
+rule config schema
 
 {
     rule_name: {
         "fields: [],
-        "config": {}, 
+        "config": {},
         ...
     }
 }
@@ -30,12 +30,16 @@ class RuleBook:
     @property
     def _all_applied_log(self):
         """return all applied logs from all rules in the rulebook"""
-        return lcall(self.rules.values(), lambda x: x.applied_log, flatten=True)
+        return lcall(
+            self.rules.values(), lambda x: x.applied_log, flatten=True
+        )
 
     @property
     def _all_invoked_log(self):
         """return all invoked logs from all rules in the rulebook"""
-        return lcall(self.rules.values(), lambda x: x.invoked_log, flatten=True)
+        return lcall(
+            self.rules.values(), lambda x: x.invoked_log, flatten=True
+        )
 
     def __getitem__(self, key: str) -> Rule:
         return self.rules[key]

lionagi/core/rule/string.py

@@ -38,4 +38,6 @@ class StringRule(Rule):
         try:
             return to_str(value, **self.validation_kwargs)
         except Exception as e:
-            raise ValueError(f"Failed to convert {value} into a string value") from e
+            raise ValueError(
+                f"Failed to convert {value} into a string value"
+            ) from e

lionagi/core/session/branch.py

@@ -165,7 +165,9 @@ class Branch(Node, DirectiveMixin):
         )
 
         if isinstance(_msg, System):
-            _msg.recipient = self.ln_id  # the branch itself, system is to the branch
+            _msg.recipient = (
+                self.ln_id
+            )  # the branch itself, system is to the branch
             self._remove_system()
             self.system = _msg
 
@@ -255,8 +257,13 @@ class Branch(Node, DirectiveMixin):
             return True
         elif is_same_dtype(tools, Tool):
             for act_ in tools:
-                if act_.schema_["function"]["name"] in self.tool_manager.registry:
-                    self.tool_manager.registry.pop(act_.schema_["function"]["name"])
+                if (
+                    act_.schema_["function"]["name"]
+                    in self.tool_manager.registry
+                ):
+                    self.tool_manager.registry.pop(
+                        act_.schema_["function"]["name"]
+                    )
             if verbose:
                 print("tools successfully deleted")
             return True
@@ -330,7 +337,11 @@ class Branch(Node, DirectiveMixin):
         return isinstance(self.messages[-1], ActionResponse)
 
     def send(
-        self, recipient: str, category: str, package: Any, request_source: str = None
+        self,
+        recipient: str,
+        category: str,
+        package: Any,
+        request_source: str = None,
     ) -> None:
         """
         Sends a mail to a recipient.

lionagi/core/session/directive_mixin.py

@@ -1,4 +1,5 @@
 # lionagi/core/session/directive_mixin.py
+from pydantic import BaseModel
 
 from lionagi.core.unit import Unit
 
@@ -33,13 +34,14 @@ class DirectiveMixin:
         default=None,
         timeout: float = None,
         timing: bool = False,
-        return_branch=False,
         images=None,
         image_path=None,
         template=None,
         verbose=True,
         formatter=None,
         format_kwargs=None,
+        pydantic_model: type[BaseModel] = None,
+        return_pydantic_model: bool = False,
         **kwargs,
     ):
         """
@@ -120,7 +122,7 @@ class DirectiveMixin:
 
             images = ImageUtil.read_image_to_base64(image_path)
 
-        return await directive.chat(
+        output = await directive.chat(
             instruction=instruction,
             context=context,
             sender=sender,
@@ -139,10 +141,16 @@ class DirectiveMixin:
             timeout=timeout,
             timing=timing,
             clear_messages=clear_messages,
-            return_branch=return_branch,
             images=images,
+            return_pydantic_model=return_pydantic_model,
+            pydantic_model=pydantic_model,
+            return_branch=False,
             **kwargs,
         )
+        while isinstance(output, tuple | list):
+            if len(output) == 2 and output[0] == output[1]:
+                output = output[0]
+        return output
 
     async def direct(
         self,

lionagi/core/session/session.py

@@ -45,7 +45,9 @@ class Session:
     ):
         self.ln_id = SysUtil.create_id()
         self.timestamp = SysUtil.get_timestamp(sep=None)[:-6]
-        system = system or "You are a helpful assistant, let's think step by step"
+        system = (
+            system or "You are a helpful assistant, let's think step by step"
+        )
         self.system = System(system=system, sender=system_sender)
         self.system_sender = system_sender
         self.branches: Pile[Branch] = self._validate_branches(branches)
@@ -77,7 +79,9 @@ class Session:
         if isinstance(value, Pile):
             for branch in value:
                 if not isinstance(branch, Branch):
-                    raise ValueError("The branches pile contains non-Branch object")
+                    raise ValueError(
+                        "The branches pile contains non-Branch object"
+                    )
             return value
         else:
             try: