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

lionagi/core/unit/parallel_unit.py

@@ -59,7 +59,9 @@ class ParallelUnit(Directive):
         else:
             self.imodel = session.imodel
         self.form_template = template or self.default_template
-        self.validator = Validator(rulebook=rulebook) if rulebook else Validator()
+        self.validator = (
+            Validator(rulebook=rulebook) if rulebook else Validator()
+        )
 
     async def pchat(self, *args, **kwargs):
         """
@@ -174,13 +176,17 @@ class ParallelUnit(Directive):
 
         async def _inner_3(i):
             """different instructions but same context"""
-            tasks = [_inner_2(i, ins_=ins_) for ins_ in convert.to_list(instruction)]
+            tasks = [
+                _inner_2(i, ins_=ins_) for ins_ in convert.to_list(instruction)
+            ]
             ress = await AsyncUtil.execute_tasks(*tasks)
             return convert.to_list(ress)
 
         async def _inner_3_b(i):
             """different context but same instruction"""
-            tasks = [_inner_2(i, cxt_=cxt_) for cxt_ in convert.to_list(context)]
+            tasks = [
+                _inner_2(i, cxt_=cxt_) for cxt_ in convert.to_list(context)
+            ]
             ress = await AsyncUtil.execute_tasks(*tasks)
             return convert.to_list(ress)
 

lionagi/core/unit/template/action.py

@@ -58,7 +58,7 @@ class ActionTemplate(BaseUnitForm):
 
         self.task = f"""
 Perform reasoning and prepare actions with GIVEN TOOLS ONLY.
-1. additional instruction: {instruction or "N/A"}. 
+1. additional instruction: {instruction or "N/A"}.
 2. additional context: {context or "N/A"}.
 """
         if confidence_score:

lionagi/core/unit/template/predict.py

@@ -34,7 +34,9 @@ class PredictTemplate(BaseUnitForm):
 
     template_name: str = "predict_template"
 
-    num_sentences: int = Field(2, description="the number of sentences to predict")
+    num_sentences: int = Field(
+        2, description="the number of sentences to predict"
+    )
 
     prediction: None | str | list = Field(
         None,

lionagi/core/unit/template/select.py

@@ -39,7 +39,9 @@ class SelectTemplate(BaseUnitForm):
     selection: Enum | str | list | None = Field(
         None, description="selection from given choices"
     )
-    choices: list = Field(default_factory=list, description="the given choices")
+    choices: list = Field(
+        default_factory=list, description="the given choices"
+    )
 
     assignment: str = "task -> selection"
 
@@ -78,9 +80,9 @@ class SelectTemplate(BaseUnitForm):
 
         self.choices = choices
         self.task = f"""
-select 1 item from the provided choices {choices}.        
+select 1 item from the provided choices {choices}.
 1. additional objective: {instruction or "N/A"}.
-2. additional information: {context or "N/A"}.     
+2. additional information: {context or "N/A"}.
 """
         if reason:
             self.append_to_request("reason")

lionagi/core/unit/unit.py

@@ -1,4 +1,8 @@
-from typing import Callable
+import logging
+from collections.abc import Callable
+
+from lionfuncs import to_dict
+from pydantic import BaseModel
 
 from lionagi.core.collections import iModel
 from lionagi.core.collections.abc import Directive
@@ -71,6 +75,8 @@ class Unit(Directive, DirectiveMixin):
         return_branch=False,
         formatter=None,
         format_kwargs={},
+        pydantic_model: type[BaseModel] = None,
+        return_pydantic_model: bool = False,
         **kwargs,
     ):
         """
@@ -100,7 +106,17 @@ class Unit(Directive, DirectiveMixin):
             Any: The processed response.
         """
         kwargs = {**retry_kwargs, **kwargs}
-        return await rcall(
+
+        if pydantic_model:
+            if form:
+                raise ValueError("Cannot use both form and pydantic_model.")
+            if requested_fields:
+                raise ValueError(
+                    "Cannot use both requested_fields and pydantic_model."
+                )
+            requested_fields = pydantic_model.model_json_schema()["properties"]
+
+        output, branch = await rcall(
             self._chat,
             instruction=instruction,
             context=context,
@@ -118,11 +134,27 @@ class Unit(Directive, DirectiveMixin):
             imodel=imodel,
             clear_messages=clear_messages,
             use_annotation=use_annotation,
-            return_branch=return_branch,
+            return_branch=True,
             formatter=formatter,
             format_kwargs=format_kwargs,
             **kwargs,
         )
+        if isinstance(output, tuple | list) and len(output) == 1:
+            output = output[0]
+
+        if isinstance(output, tuple | list) and len(output) == 2:
+            if output[0] == output[1]:
+                output = output[0]
+
+        if return_pydantic_model:
+            try:
+                a_ = to_dict(output, recursive=True, max_recursive_depth=3)
+                output = pydantic_model(**a_)
+                return output, branch if return_branch else output
+            except Exception as e:
+                logging.error(f"Error converting to pydantic model: {e}")
+
+        return output, branch if return_branch else output
 
     async def direct(
         self,
@@ -232,7 +264,9 @@ class Unit(Directive, DirectiveMixin):
         )
 
         if verbose:
-            print("--------------------------------------------------------------")
+            print(
+                "--------------------------------------------------------------"
+            )
             print(f"Directive successfully completed!")
 
         return out

lionagi/core/unit/unit_form.py

@@ -38,7 +38,9 @@ class UnitForm(BaseUnitForm):
             "number, you should provide a number like 1, 23, or 1.1 if float is "
             "allowed."
         ),
-        examples=["{action_1: {function: 'add', arguments: {num1: 1, num2: 2}}}"],
+        examples=[
+            "{action_1: {function: 'add', arguments: {num1: 1, num2: 2}}}"
+        ],
     )
 
     action_required: bool | None = Field(
@@ -188,13 +190,13 @@ class UnitForm(BaseUnitForm):
 
         if allow_action:
             self.append_to_request("actions, action_required, reason")
-            self.task += "- Reason and prepare actions with GIVEN TOOLS ONLY.\n"
+            self.task += (
+                "- Reason and prepare actions with GIVEN TOOLS ONLY.\n"
+            )
 
         if allow_extension:
             self.append_to_request("extension_required")
-            self.task += (
-                f"- Allow auto-extension up to another {max_extension} rounds.\n"
-            )
+            self.task += f"- Allow auto-extension up to another {max_extension} rounds.\n"
 
         if tool_schema:
             self.append_to_input("tool_schema")
@@ -209,21 +211,15 @@ class UnitForm(BaseUnitForm):
             max_extension = max_extension or plan_num_step
             allow_extension = True
             self.append_to_request("plan, extension_required")
-            self.task += (
-                f"- Generate a {plan_num_step}-step plan based on the context.\n"
-            )
+            self.task += f"- Generate a {plan_num_step}-step plan based on the context.\n"
 
         if predict:
             self.append_to_request("prediction")
-            self.task += (
-                f"- Predict the next {predict_num_sentences or 1} sentence(s).\n"
-            )
+            self.task += f"- Predict the next {predict_num_sentences or 1} sentence(s).\n"
 
         if select:
             self.append_to_request("selection")
-            self.task += (
-                f"- Select 1 item from the provided choices: {select_choices}.\n"
-            )
+            self.task += f"- Select 1 item from the provided choices: {select_choices}.\n"
 
         if confidence:
             self.append_to_request("confidence_score")
@@ -238,7 +234,9 @@ class UnitForm(BaseUnitForm):
                 "upper_bound": score_range[1],
                 "lower_bound": score_range[0],
                 "num_type": int if score_num_digits == 0 else float,
-                "precision": score_num_digits if score_num_digits != 0 else None,
+                "precision": (
+                    score_num_digits if score_num_digits != 0 else None
+                ),
             }
 
             self.task += (

lionagi/core/unit/unit_mixin.py

@@ -22,16 +22,16 @@ class DirectiveMixin(ABC):
 
     def _create_chat_config(
         self,
-        system: Optional[str] = None,
-        instruction: Optional[str] = None,
-        context: Optional[str] = None,
-        images: Optional[str] = None,
-        sender: Optional[str] = None,
-        recipient: Optional[str] = None,
-        requested_fields: Optional[list] = None,
+        system: str | None = None,
+        instruction: str | None = None,
+        context: str | None = None,
+        images: str | None = None,
+        sender: str | None = None,
+        recipient: str | None = None,
+        requested_fields: list | None = None,
         form: Form = None,
         tools: bool = False,
-        branch: Optional[Any] = None,
+        branch: Any | None = None,
         **kwargs,
     ) -> Any:
         """
@@ -87,7 +87,7 @@ class DirectiveMixin(ABC):
         return config
 
     async def _call_chatcompletion(
-        self, imodel: Optional[Any] = None, branch: Optional[Any] = None, **kwargs
+        self, imodel: Any | None = None, branch: Any | None = None, **kwargs
     ) -> Any:
         """
         Calls the chat completion model.
@@ -102,7 +102,9 @@ class DirectiveMixin(ABC):
         """
         imodel = imodel or self.imodel
         branch = branch or self.branch
-        return await imodel.call_chat_completion(branch.to_chat_messages(), **kwargs)
+        return await imodel.call_chat_completion(
+            branch.to_chat_messages(), **kwargs
+        )
 
     async def _process_chatcompletion(
         self,
@@ -110,8 +112,8 @@ class DirectiveMixin(ABC):
         completion: dict,
         sender: str,
         invoke_tool: bool = True,
-        branch: Optional[Any] = None,
-        action_request: Optional[Any] = None,
+        branch: Any | None = None,
+        action_request: Any | None = None,
         costs=None,
     ) -> Any:
         """
@@ -157,7 +159,9 @@ class DirectiveMixin(ABC):
                 m = completion.get("model", None)
                 if m:
                     ttl = (a * price[0] + b * price[1]) / 1000000
-                    branch.messages[-1]._meta_insert(["extra", "usage", "expense"], ttl)
+                    branch.messages[-1]._meta_insert(
+                        ["extra", "usage", "expense"], ttl
+                    )
                 return msg
 
             if _choices and not isinstance(_choices, list):
@@ -180,10 +184,10 @@ class DirectiveMixin(ABC):
 
     async def _process_action_request(
         self,
-        _msg: Optional[dict] = None,
-        branch: Optional[Any] = None,
+        _msg: dict | None = None,
+        branch: Any | None = None,
         invoke_tool: bool = True,
-        action_request: Optional[Any] = None,
+        action_request: Any | None = None,
     ) -> Any:
         """
         Processes an action request from the assistant response.
@@ -204,9 +208,13 @@ class DirectiveMixin(ABC):
         if action_request:
             for i in action_request:
                 if i.function in branch.tool_manager.registry:
-                    i.recipient = branch.tool_manager.registry[i.function].ln_id
+                    i.recipient = branch.tool_manager.registry[
+                        i.function
+                    ].ln_id
                 else:
-                    raise ActionError(f"Tool {i.function} not found in registry")
+                    raise ActionError(
+                        f"Tool {i.function} not found in registry"
+                    )
                 branch.add_message(action_request=i, recipient=i.recipient)
 
         if invoke_tool:
@@ -305,7 +313,7 @@ class DirectiveMixin(ABC):
         requested_fields: dict = None,
         form: Form = None,
         tools: Any = False,
-        images: Optional[str] = None,
+        images: str | None = None,
         invoke_tool: bool = True,
         return_form: bool = True,
         strict: bool = False,
@@ -396,7 +404,7 @@ class DirectiveMixin(ABC):
         return_form=True,
         strict=False,
         imodel=None,
-        images: Optional[str] = None,
+        images: str | None = None,
         clear_messages=False,
         use_annotation=True,
         timeout: float = None,
@@ -493,7 +501,7 @@ class DirectiveMixin(ABC):
         predict_num_sentences=None,
         clear_messages=False,
         return_branch=False,
-        images: Optional[str] = None,
+        images: str | None = None,
         verbose=None,
         **kwargs,
     ):
@@ -584,7 +592,7 @@ class DirectiveMixin(ABC):
         predict_num_sentences=None,
         clear_messages=False,
         return_branch=False,
-        images: Optional[str] = None,
+        images: str | None = None,
         verbose=True,
         formatter=None,
         format_kwargs=None,
@@ -736,10 +744,14 @@ class DirectiveMixin(ABC):
             )
 
             if verbose:
-                print(f"------------------- Extension completed -------------------\n")
+                print(
+                    f"------------------- Extension completed -------------------\n"
+                )
 
             extension_forms.extend(new_form)
-            last_form = new_form[-1] if isinstance(new_form, list) else new_form
+            last_form = (
+                new_form[-1] if isinstance(new_form, list) else new_form
+            )
             ctr += len(form)
 
         if extension_forms:
@@ -977,7 +989,9 @@ class DirectiveMixin(ABC):
                 context=context,
             )
 
-        return await self._chat(form=form, return_form=True, branch=branch, **kwargs)
+        return await self._chat(
+            form=form, return_form=True, branch=branch, **kwargs
+        )
 
     async def _predict(
         self,
@@ -1019,7 +1033,9 @@ class DirectiveMixin(ABC):
                 reason=reason,
             )
 
-        return await self._chat(form=form, return_form=True, branch=branch, **kwargs)
+        return await self._chat(
+            form=form, return_form=True, branch=branch, **kwargs
+        )
 
     async def _score(
         self,
@@ -1066,7 +1082,9 @@ class DirectiveMixin(ABC):
                 context=context,
             )
 
-        return await self._chat(form=form, return_form=True, branch=branch, **kwargs)
+        return await self._chat(
+            form=form, return_form=True, branch=branch, **kwargs
+        )
 
     async def _plan(
         self,

lionagi/core/unit/util.py

@@ -21,7 +21,7 @@ choices_fields = ["index", "message", "logprobs", "finish_reason"]
 
 usage_fields = ["prompt_tokens", "completion_tokens", "total_tokens"]
 
-from typing import Callable
+from collections.abc import Callable
 
 from lionagi.core.action.tool import Tool
 from lionagi.core.action.tool_manager import func_to_tool
@@ -40,12 +40,16 @@ def process_tools(tool_obj, branch):
 def _process_tool(tool_obj, branch):
     if (
         isinstance(tool_obj, Tool)
-        and tool_obj.schema_["function"]["name"] not in branch.tool_manager.registry
+        and tool_obj.schema_["function"]["name"]
+        not in branch.tool_manager.registry
     ):
         branch.register_tools(tool_obj)
     if isinstance(tool_obj, Callable):
         tool = func_to_tool(tool_obj)[0]
-        if tool.schema_["function"]["name"] not in branch.tool_manager.registry:
+        if (
+            tool.schema_["function"]["name"]
+            not in branch.tool_manager.registry
+        ):
             branch.register_tools(tool)
 
 

lionagi/core/validator/validator.py

@@ -1,5 +1,6 @@
 import asyncio
-from typing import Any, Callable, Dict, List, Union
+from collections.abc import Callable
+from typing import Any, Dict, List, Union
 
 from lionfuncs import lcall
 
@@ -40,10 +41,10 @@ class Validator:
         self,
         *,
         rulebook: RuleBook = None,
-        rules: Dict[str, Rule] = None,
-        order: List[str] = None,
-        init_config: Dict[str, Dict] = None,
-        active_rules: Dict[str, Rule] = None,
+        rules: dict[str, Rule] = None,
+        order: list[str] = None,
+        init_config: dict[str, dict] = None,
+        active_rules: dict[str, Rule] = None,
         formatter: Callable = None,
         format_kwargs: dict = {},
     ):
@@ -63,12 +64,14 @@ class Validator:
         self.rulebook = rulebook or RuleBook(
             rules or _DEFAULT_RULES, order or _DEFAULT_RULEORDER, init_config
         )
-        self.active_rules: Dict[str, Rule] = active_rules or self._initiate_rules()
+        self.active_rules: dict[str, Rule] = (
+            active_rules or self._initiate_rules()
+        )
         self.validation_log = []
         self.formatter = formatter
         self.format_kwargs = format_kwargs
 
-    def _initiate_rules(self) -> Dict[str, Rule]:
+    def _initiate_rules(self) -> dict[str, Rule]:
         """
         Initialize rules from the rulebook.
 
@@ -157,7 +160,7 @@ class Validator:
             raise FieldError(error_message)
 
     async def validate_report(
-        self, report: Report, forms: List[Form], strict: bool = True
+        self, report: Report, forms: list[Form], strict: bool = True
     ) -> Report:
         """
         Validate a report based on active rules.
@@ -176,7 +179,7 @@ class Validator:
     async def validate_response(
         self,
         form: Form,
-        response: Union[dict, str],
+        response: dict | str,
         strict: bool = True,
         use_annotation: bool = True,
     ) -> Form:
@@ -201,21 +204,29 @@ class Validator:
             else:
                 if self.formatter:
                     if asyncio.iscoroutinefunction(self.formatter):
-                        response = await self.formatter(response, **self.format_kwargs)
+                        response = await self.formatter(
+                            response, **self.format_kwargs
+                        )
                         print("formatter used")
                     else:
-                        response = self.formatter(response, **self.format_kwargs)
+                        response = self.formatter(
+                            response, **self.format_kwargs
+                        )
                         print("formatter used")
 
         if not isinstance(response, dict):
-            raise ValueError(f"The form response format is invalid for filling.")
+            raise ValueError(
+                f"The form response format is invalid for filling."
+            )
 
         dict_ = {}
         for k, v in response.items():
             if k in form.requested_fields:
                 kwargs = form.validation_kwargs.get(k, {})
                 _annotation = form._field_annotations[k]
-                if (keys := form._get_field_attr(k, "choices", None)) is not None:
+                if (
+                    keys := form._get_field_attr(k, "choices", None)
+                ) is not None:
                     v = await self.validate_field(
                         field=k,
                         value=v,
@@ -227,7 +238,9 @@ class Validator:
                         **kwargs,
                     )
 
-                elif (_keys := form._get_field_attr(k, "keys", None)) is not None:
+                elif (
+                    _keys := form._get_field_attr(k, "keys", None)
+                ) is not None:
 
                     v = await self.validate_field(
                         field=k,
@@ -346,7 +359,7 @@ class Validator:
         }
         self.validation_log.append(log_entry)
 
-    def get_validation_summary(self) -> Dict[str, Any]:
+    def get_validation_summary(self) -> dict[str, Any]:
         """
         Get a summary of validation results.
 

lionagi/core/work/work_edge.py

@@ -1,5 +1,5 @@
 import inspect
-from typing import Callable
+from collections.abc import Callable
 
 from pydantic import Field, field_validator
 
@@ -54,7 +54,9 @@ class WorkEdge(Edge, Progressable):
             getattr(func, "_worklink_decorator_params")
             return func
         except:
-            raise ValueError("convert_function must be a worklink decorated function")
+            raise ValueError(
+                "convert_function must be a worklink decorated function"
+            )
 
     @property
     def name(self):
@@ -94,5 +96,7 @@ class WorkEdge(Edge, Progressable):
             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)
+        next_work = await self.convert_function(
+            self=self.associated_worker, **kwargs
+        )
         return next_work

lionagi/core/work/work_task.py

@@ -1,5 +1,5 @@
 import inspect
-from typing import Callable
+from collections.abc import Callable
 
 from pydantic import Field, field_validator
 
@@ -31,9 +31,13 @@ class WorkTask(Component):
 
     work_history: list[Work] = Field([], description="List of works processed")
 
-    max_steps: int | None = Field(10, description="Maximum number of works allowed")
+    max_steps: int | None = Field(
+        10, description="Maximum number of works allowed"
+    )
 
-    current_work: Work | None = Field(None, description="The current work in progress")
+    current_work: Work | None = Field(
+        None, description="The current work in progress"
+    )
 
     post_processing: Callable | None = Field(
         None,
@@ -55,7 +59,9 @@ class WorkTask(Component):
             ValueError: If value is not a positive integer.
         """
         if value <= 0:
-            raise ValueError("Invalid value: max_steps must be a positive integer.")
+            raise ValueError(
+                "Invalid value: max_steps must be a positive integer."
+            )
         return value
 
     @field_validator("post_processing", mode="before")
@@ -72,7 +78,7 @@ class WorkTask(Component):
         Raises:
             ValueError: If value is not an asynchronous function.
         """
-        if value is not None and not inspect.iscoroutinefunction((value)):
+        if value is not None and not inspect.iscoroutinefunction(value):
             raise ValueError("post_processing must be a async function")
         return value
 

lionagi/core/work/worker.py

@@ -62,7 +62,12 @@ class Worker(ABC):
         """
 
         return (
-            any([await i.is_progressable() for i in self.work_functions.values()])
+            any(
+                [
+                    await i.is_progressable()
+                    for i in self.work_functions.values()
+                ]
+            )
             and not self.stopped
         )
 
@@ -261,7 +266,9 @@ def work(
             **kwargs,
         ):
             if not inspect.iscoroutinefunction(func):
-                raise TypeError(f"{func.__name__} must be an asynchronous function")
+                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._work_wrapper(
@@ -309,7 +316,9 @@ def worklink(from_: str, to_: str, auto_schedule: bool = True):
             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")
+                raise TypeError(
+                    f"{func.__name__} must be an asynchronous function"
+                )
 
             work_funcs = self._get_decorated_functions(
                 decorator_attr="_work_decorator_params"
@@ -363,7 +372,9 @@ def worklink(from_: str, to_: str, auto_schedule: bool = True):
                     next_params[1], dict
                 ):
                     if wrapper.auto_schedule:
-                        return await to_work_func(*next_params[0], **next_params[1])
+                        return await to_work_func(
+                            *next_params[0], **next_params[1]
+                        )
                 else:
                     raise TypeError(f"Invalid return type {func.__name__}")
             else:
@@ -372,7 +383,11 @@ def worklink(from_: str, to_: str, auto_schedule: bool = True):
             return next_params
 
         wrapper.auto_schedule = auto_schedule
-        wrapper._worklink_decorator_params = {"func": func, "from_": from_, "to_": to_}
+        wrapper._worklink_decorator_params = {
+            "func": func,
+            "from_": from_,
+            "to_": to_,
+        }
 
         return wrapper
 

lionagi/core/work/worker_engine.py

@@ -172,8 +172,12 @@ class WorkerEngine:
         )
         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])
+                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