LitTool.from_model method to create LitTool from Pydantic (#57)

* from_model method to create LitTool from Pydantic

Added a class method 'from_model' to create a LitTool from a Pydantic model, including setup and run methods for validation.

* add test for run method

* mypy silence

* ruff format

* lint

* typehint change

* revert to Optional[str]

* Modify run method to return serialized instance

Updated the run method to return a serialized instance instead of a model instance.

* Update assertion in test_tool_run_from_model

Author: mathematicalmichael

src/litai/tools.py

@@ -117,6 +117,25 @@ def _extract_parameters(self) -> Dict[str, Any]:
 
         return LangchainTool()
 
+    @classmethod
+    def from_model(cls, model: type[BaseModel]) -> "LitTool":
+        """Create a LitTool that exposes a Pydantic model as a structured schema."""
+
+        class ModelTool(LitTool):
+            def setup(self) -> None:
+                super().setup()
+                self.name = model.__name__
+                self.description = model.__doc__ or ""
+
+            def run(self, *args, **kwargs) -> Any:  # type: ignore
+                # Default implementation: validate & return a serialized instance
+                return model(*args, **kwargs).model_dump()
+
+            def _extract_parameters(self) -> Dict[str, Any]:
+                return model.model_json_schema()
+
+        return ModelTool()
+
     @classmethod
     def convert_tools(cls, tools: Optional[Sequence[Union["LitTool", "StructuredTool"]]]) -> List["LitTool"]:
         """Convert a list of tools into LitTool instances.

tests/test_tools.py

@@ -15,10 +15,21 @@
 
 import pytest
 from langchain_core.tools import tool as langchain_tool
+from pydantic import BaseModel
 
 from litai import LitTool, tool
 
 
+@pytest.fixture
+def weather_tool_model():
+    class WeatherRequest(BaseModel):
+        """Get weather for location."""
+
+        location: str
+
+    return WeatherRequest
+
+
 @pytest.fixture
 def basic_tool_class():
     class TestTool(LitTool):
@@ -226,3 +237,45 @@ def get_weather(city: str) -> str:
 
     with pytest.raises(TypeError, match="Unsupported tool type: <class 'function'>"):
         LitTool.convert_tools([get_weather])
+
+
+def test_tool_from_model_with_no_description(weather_tool_model):
+    weather_tool_model.__doc__ = None
+
+    lit_tool = LitTool.from_model(weather_tool_model)
+
+    assert isinstance(lit_tool, LitTool)
+    assert lit_tool.name == "WeatherRequest"
+    assert lit_tool.description == ""
+
+    assert lit_tool.as_tool() == {
+        "type": "function",
+        "function": {
+            "name": "WeatherRequest",
+            "description": "",
+            "parameters": weather_tool_model.model_json_schema(),
+        },
+    }
+
+
+def test_tool_run_from_model(weather_tool_model):
+    lit_tool = LitTool.from_model(weather_tool_model)
+
+    assert lit_tool.run(location="NYC") == weather_tool_model(location="NYC").model_dump()
+
+
+def test_tool_from_model_with_description(weather_tool_model):
+    lit_tool = LitTool.from_model(weather_tool_model)
+
+    assert isinstance(lit_tool, LitTool)
+    assert lit_tool.name == "WeatherRequest"
+    assert lit_tool.description == "Get weather for location."
+
+    assert lit_tool.as_tool() == {
+        "type": "function",
+        "function": {
+            "name": "WeatherRequest",
+            "description": "Get weather for location.",
+            "parameters": weather_tool_model.model_json_schema(),
+        },
+    }