lionagi 0.3.3__py3-none-any.whl → 0.3.4__py3-none-any.whl

lionagi/core/unit/unit.py

@@ -12,7 +12,7 @@ from lionagi.libs.ln_func_call import rcall
 
 from .unit_form import UnitForm
 from .unit_mixin import DirectiveMixin
-from .util import retry_kwargs
+from .util import break_down_annotation, retry_kwargs
 
 
 class Unit(Directive, DirectiveMixin):
@@ -75,7 +75,7 @@ class Unit(Directive, DirectiveMixin):
         return_branch=False,
         formatter=None,
         format_kwargs={},
-        pydantic_model: type[BaseModel] = None,
+        pydantic_model: type[BaseModel] | BaseModel = None,
         return_pydantic_model: bool = False,
         **kwargs,
     ):
@@ -114,7 +114,13 @@ class Unit(Directive, DirectiveMixin):
                 raise ValueError(
                     "Cannot use both requested_fields and pydantic_model."
                 )
-            requested_fields = pydantic_model.model_json_schema()["properties"]
+            requested_fields = break_down_annotation(pydantic_model)
+            context = {
+                "info": context,
+                "return_guidance": pydantic_model.model_json_schema()[
+                    "properties"
+                ],
+            }
 
         output, branch = await rcall(
             self._chat,
@@ -148,8 +154,13 @@ class Unit(Directive, DirectiveMixin):
 
         if return_pydantic_model:
             try:
-                a_ = to_dict(output, recursive=True, max_recursive_depth=3)
-                output = pydantic_model(**a_)
+                a_ = to_dict(
+                    output,
+                    recursive=True,
+                    max_recursive_depth=5,
+                    fuzzy_parse=True,
+                )
+                output = pydantic_model.model_validate(a_)
                 return output, branch if return_branch else output
             except Exception as e:
                 logging.error(f"Error converting to pydantic model: {e}")

lionagi/core/unit/util.py

@@ -1,3 +1,8 @@
+import inspect
+from typing import get_args
+
+from pydantic import BaseModel
+
 retry_kwargs = {
     "retries": 0,  # kwargs for rcall, number of retries if failed
     "delay": 0,  # number of seconds to delay before retrying
@@ -74,3 +79,25 @@ async def _direct(
         )
 
     return await directive.direct(form=form, return_form=True, **kwargs)
+
+
+def break_down_annotation(model: type[BaseModel]):
+
+    def _ispydantic_model(x):
+        return inspect.isclass(x) and issubclass(x, BaseModel)
+
+    if not _ispydantic_model(model):
+        return model
+
+    out = {}
+    for k, v in model.__annotations__.items():
+        if _ispydantic_model(v):
+            out[k] = break_down_annotation(v)
+        elif "list" in str(v) and get_args(v):
+            v = get_args(v)[0]
+            if _ispydantic_model(v):
+                v = break_down_annotation(v)
+            out[k] = [v]
+        else:
+            out[k] = v
+    return out

lionagi/version.py

@@ -1 +1 @@
-__version__ = "0.3.3"
+__version__ = "0.3.4"

{lionagi-0.3.3.dist-info → lionagi-0.3.4.dist-info}/METADATA

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

{lionagi-0.3.3.dist-info → lionagi-0.3.4.dist-info}/RECORD

@@ -105,10 +105,10 @@ lionagi/core/unit/template/plan.py,sha256=i4FmKEB8eRsRCsTanvfoX-2RZ8SaM1qvLBluuY
 lionagi/core/unit/template/predict.py,sha256=-EIZQo0ZjGKy3MiM0AtqmbnJpbXcrwtSCON5n3jcyVo,3160
 lionagi/core/unit/template/score.py,sha256=ReUaIIr-NLjunSy4NNXQpIsH28NNceGBAUuPCRptzMc,3809
 lionagi/core/unit/template/select.py,sha256=VSpkphJl9bHSE8i0X6MMJD8LB5QwOj1UORHm8VDIRKE,3047
-lionagi/core/unit/unit.py,sha256=1kFz_rLM1f8OKn4pbwL2WiVyqTzyEPOmHL9OdIiwzHI,15252
+lionagi/core/unit/unit.py,sha256=5XJ8k401cqVNLOJpNcHul3IG6WflC--AlSg2IprlSkk,15600
 lionagi/core/unit/unit_form.py,sha256=zK_ij3Tod5FwMVdIIhdVoEFvD3br-YM9RPe7WsOIW2s,10980
 lionagi/core/unit/unit_mixin.py,sha256=c8GvHzgc65iJKQBKv71ET3afLPsIz5-Ce-4Eo6_bZiw,38823
-lionagi/core/unit/util.py,sha256=yEIf4ksWLB4X0IZNmwvIphvXiRbtWtGqbWCPnIt1u-s,2048
+lionagi/core/unit/util.py,sha256=WN2Jop-LUwQNYJNubFPhOZrisQ6SQq-XMhD_KhzLkgE,2707
 lionagi/core/validator/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 lionagi/core/validator/validator.py,sha256=901wwmqL92XNi25ajv57bNKIKZhmu-KprNLArMUiGqg,12255
 lionagi/core/work/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -229,8 +229,8 @@ lionagi/lions/researcher/data_source/finhub_.py,sha256=W63daXgIwHJQ6TDMR2ALQIDk1
 lionagi/lions/researcher/data_source/google_.py,sha256=401SKHQaSpxiOUoXl7stadl4qeF7SIX72lUNK7bKesg,6797
 lionagi/lions/researcher/data_source/wiki_.py,sha256=UPoa2dk_y5sELu7_rkdme2auDpUmc_Dn0Avgjwr2X2g,3145
 lionagi/lions/researcher/data_source/yfinance_.py,sha256=snAf897J69MyAc6fcFjF0irrMjbAh81EZ3RvaFT3hxE,977
-lionagi/version.py,sha256=8KcCYTXH99C2-gCLuPILJvtT9YftRWJsartIx6TQ2ZY,22
-lionagi-0.3.3.dist-info/LICENSE,sha256=VXFWsdoN5AAknBCgFqQNgPWYx7OPp-PFEP961zGdOjc,11288
-lionagi-0.3.3.dist-info/METADATA,sha256=aPmmyz7FbSYaeDChIRXx9r3C6-2ZVs-S8YzEZtnpyZg,3149
-lionagi-0.3.3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-lionagi-0.3.3.dist-info/RECORD,,
+lionagi/version.py,sha256=oYLGMpySamd16KLiaBTfRyrAS7_oyp-TOEHmzmeumwg,22
+lionagi-0.3.4.dist-info/LICENSE,sha256=VXFWsdoN5AAknBCgFqQNgPWYx7OPp-PFEP961zGdOjc,11288
+lionagi-0.3.4.dist-info/METADATA,sha256=VM0JO4wP1qKYkLZ7Al0lLiOi0ZaWreAjxbJmkyhdaOo,3149
+lionagi-0.3.4.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+lionagi-0.3.4.dist-info/RECORD,,