lionagi 0.3.0__tar.gz → 0.3.1__tar.gz

{lionagi-0.3.0 → lionagi-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lionagi
-Version: 0.3.0
+Version: 0.3.1
 Summary: Towards automated general intelligence.
 Author: HaiyangLi
 Author-email: quantocean.li@gmail.com

{lionagi-0.3.0 → lionagi-0.3.1}/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-0.3.0 → lionagi-0.3.1}/lionagi/core/unit/unit.py

@@ -1,5 +1,9 @@
+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
 from lionagi.core.report.form import Form
@@ -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,

lionagi-0.3.1/lionagi/version.py

@@ -0,0 +1 @@
+__version__ = "0.3.1"

{lionagi-0.3.0 → lionagi-0.3.1}/pyproject.toml

@@ -7,7 +7,7 @@ classifiers=[
 
 [tool.poetry]
 name = "lionagi"
-version = "0.3.0"
+version = "0.3.1"
 description = "Towards automated general intelligence."
 authors = ["HaiyangLi <quantocean.li@gmail.com>"]
 readme = "README.md"

lionagi-0.3.0/lionagi/version.py

@@ -1 +0,0 @@
-__version__ = "0.3.0"