Add Pydantic model support for tool parameters

- Implement _deserialize_pydantic_args() method to handle Pydantic BaseModel parameters in tool calls
- Update tool execution to automatically deserialize Pydantic model arguments
- Add JSON schema generation for Pydantic model parameters in tool definitions
- Bump version to 5.4.0 to reflect the new Pydantic model support
- Add test case to verify tool definition and execution with Pydantic models

Author: knowsuchagency

promptic.py

@@ -17,7 +17,7 @@
 from litellm import completion as litellm_completion
 from pydantic import BaseModel
 
-__version__ = "5.3.1"
+__version__ = "5.4.0"
 
 SystemPrompt = Optional[Union[str, List[str], List[Dict[str, str]]]]
 
@@ -288,6 +288,8 @@ def _generate_tool_definition(self, fn: Callable) -> dict:
                 param_info["type"] = "number"
             elif param_type == bool:
                 param_info["type"] = "boolean"
+            elif inspect.isclass(param_type) and issubclass(param_type, BaseModel):
+                param_info = param_type.model_json_schema()
 
             parameters["properties"][name] = param_info
 
@@ -378,6 +380,28 @@ def llm(self, func: Callable = None, **kwargs):
 
         return new_instance._decorator(func) if func else new_instance._decorator
 
+    def _deserialize_pydantic_args(self, fn: Callable, function_args: dict) -> dict:
+        """Deserialize any Pydantic model parameters in the function arguments.
+
+        Args:
+            fn: The function whose parameters to check
+            function_args: The arguments to deserialize
+
+        Returns:
+            The function arguments with any Pydantic models deserialized
+        """
+        sig = inspect.signature(fn)
+        for param_name, param in sig.parameters.items():
+            param_type = param.annotation if param.annotation != inspect._empty else Any
+            if (
+                inspect.isclass(param_type)
+                and issubclass(param_type, BaseModel)
+                and param_name in function_args
+                and isinstance(function_args[param_name], dict)
+            ):
+                function_args[param_name] = param_type(**function_args[param_name])
+        return function_args
+
     def _decorator(self, func: Callable):
         return_type = func.__annotations__.get("return")
 
@@ -590,9 +614,14 @@ def wrapper(*args, **kwargs):
                                     function_response = f"[DRY RUN] Would have called {function_name = } {function_args = }"
                                 else:
                                     try:
-                                        function_response = self.tools[function_name](
-                                            **function_args
+                                        self.logger.debug(
+                                            f"Calling tool {function_name}({function_args}) using {self.model = }"
+                                        )
+                                        fn = self.tools[function_name]
+                                        function_args = self._deserialize_pydantic_args(
+                                            fn, function_args
                                         )
+                                        function_response = fn(**function_args)
                                     except Exception as e:
                                         self.logger.error(
                                             f"Error calling tool {function_name}({function_args}): {e}"
@@ -625,6 +654,7 @@ def wrapper(*args, **kwargs):
         wrapper.tool = self.tool
         wrapper.clear = self.clear
         wrapper.message = self.message
+        wrapper.instance = self
 
         # Automatically expose all other attributes from self
         for attr_name, attr_value in self.__dict__.items():
@@ -683,9 +713,16 @@ def _stream_response(self, response, call=None):
                                                 f"[DRY RUN] Would have called {tool_info['name']} with {function_args}"
                                             )
                                         else:
-                                            self.tools[tool_info["name"]](
-                                                **function_args
+                                            self.logger.debug(
+                                                f"Calling tool {tool_info['name']}({function_args}) using {self.model = }"
+                                            )
+                                            fn = self.tools[tool_info["name"]]
+                                            function_args = (
+                                                self._deserialize_pydantic_args(
+                                                    fn, function_args
+                                                )
                                             )
+                                            fn(**function_args)
                                         # Clear after successful execution
                                         del current_tool_calls[current_index]
                                 except json.JSONDecodeError:

tests/test_promptic.py

@@ -1668,3 +1668,46 @@ def child_function(command):
     # It should respond without using the tool
     result = child_function("What time is it?")
     assert "12:00" not in result
+
+
+@pytest.mark.parametrize("model", REGULAR_MODELS)
+@pytest.mark.parametrize(
+    "create_completion_fn", [openai_completion_fn, litellm_completion]
+)
+def test_tool_definition_with_pydantic_param(model, create_completion_fn):
+    """Test that tool definitions properly handle Pydantic BaseModel parameters"""
+    if create_completion_fn == openai_completion_fn and not model.startswith("gpt"):
+        pytest.skip("Non-GPT models are not supported with OpenAI client")
+    if "gemini" in model:
+        pytest.skip("The Gemini model is flaky with this test.")
+
+    class UserProfile(BaseModel):
+        name: str
+        age: int
+        email: str
+
+    llm = Promptic(
+        model=model,
+        temperature=0,
+        timeout=5,
+        create_completion_fn=create_completion_fn,
+    )
+
+    @llm
+    def assistant(command):
+        """{command}"""
+
+    name = None
+
+    @assistant.tool
+    def update_user(profile: UserProfile) -> str:
+        """Update user profile"""
+        nonlocal name
+        name = profile.name
+        return f"Updated profile for {profile.name}"
+
+    assistant(
+        "Update the user's profile with the name 'John Doe', age 30, and email '[email protected]'"
+    )
+
+    assert name == "John Doe"