🔧 Refactor helper functions names

shroominic · Dec 20, 2023 · ce5a2a5 · ce5a2a5
1 parent 1e8a59b
commit ce5a2a5
Showing 1 changed file with 10 additions and 18 deletions.
diff --git a/src/funcchain/utils/helpers.py b/src/funcchain/utils/helpers.py
@@ -97,9 +97,10 @@ def _remove_a_key(d: dict, remove_key: str) -> None:
                 _remove_a_key(d[key], remove_key)
 
 
-def pydantic_to_functions(pydantic_object: Type[BaseModel]) -> dict[str, Any]:
-    schema = pydantic_object.schema()
-    docstring = parse(pydantic_object.__doc__ or "")
+def pydantic_to_functions(pydantic_type: Type[BaseModel]) -> dict[str, Any]:
+    schema = pydantic_type.model_json_schema()
+
+    docstring = parse(pydantic_type.__doc__ or "")
     parameters = {k: v for k, v in schema.items() if k not in ("title", "description")}
 
     for param in docstring.params:
@@ -116,29 +117,20 @@ def pydantic_to_functions(pydantic_object: Type[BaseModel]) -> dict[str, Any]:
             schema["description"] = docstring.short_description
         else:
             schema["description"] = (
-                f"Correctly extracted `{pydantic_object.__name__.lower()}` with all "
+                f"Correctly extracted `{pydantic_type.__name__.lower()}` with all "
                 f"the required parameters with correct types"
             )
 
     _remove_a_key(parameters, "title")
     _remove_a_key(parameters, "additionalProperties")
 
-    # print(
-    #     "pydantic_to_functions",
-    #     {
-    #         "name": pydantic_object.__name__.lower(),
-    #         "description": schema["description"],
-    #         "parameters": parameters,
-    #     },
-    # )
-
     return {
         "function_call": {
-            "name": pydantic_object.__name__.lower(),
+            "name": pydantic_type.__name__.lower(),
         },
         "functions": [
             {
-                "name": pydantic_object.__name__.lower(),
+                "name": pydantic_type.__name__.lower(),
                 "description": schema["description"],
                 "parameters": parameters,
             },
@@ -147,11 +139,11 @@ def pydantic_to_functions(pydantic_object: Type[BaseModel]) -> dict[str, Any]:
 
 
 def multi_pydantic_to_functions(
-    pydantic_objects: list[Type[BaseModel]],
+    pydantic_types: list[Type[BaseModel]],
 ) -> dict[str, Any]:
     functions: list[dict[str, Any]] = [
-        pydantic_to_functions(pydantic_object)["functions"][0]
-        for pydantic_object in pydantic_objects
+        pydantic_to_functions(pydantic_type)["functions"][0]
+        for pydantic_type in pydantic_types
     ]
 
     return {