From 41d0ad11324580ac398f3076451b339f8cf2897b Mon Sep 17 00:00:00 2001 From: "Michael Pilosov, PhD" <40366263+mathematicalmichael@users.noreply.github.com> Date: Sat, 30 Aug 2025 19:03:37 -0600 Subject: [PATCH 1/6] 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. --- src/litai/tools.py | 18 ++++++++++++++++++ tests/test_tools.py | 44 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+) diff --git a/src/litai/tools.py b/src/litai/tools.py index 9598f19..c5af01a 100644 --- a/src/litai/tools.py +++ b/src/litai/tools.py @@ -117,6 +117,24 @@ 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: + # Default implementation: validate & return an instance + return model(*args, **kwargs) + + 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. diff --git a/tests/test_tools.py b/tests/test_tools.py index c642071..13318e2 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -15,10 +15,19 @@ 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 +235,38 @@ def get_weather(city: str) -> str: with pytest.raises(TypeError, match="Unsupported tool type: "): 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_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(), + }, + } From 9376d9f994509dc4fb101fa4edb113c822ff49a7 Mon Sep 17 00:00:00 2001 From: Michael Pilosov Date: Tue, 2 Sep 2025 18:24:50 +0000 Subject: [PATCH 2/6] add test for run method --- tests/test_tools.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/test_tools.py b/tests/test_tools.py index 13318e2..66d8b92 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -255,6 +255,13 @@ def test_tool_from_model_with_no_description(weather_tool_model): }, } + +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") + + def test_tool_from_model_with_description(weather_tool_model): lit_tool = LitTool.from_model(weather_tool_model) From b107fd612c0fcea18ffcba6a4d64224366b10d71 Mon Sep 17 00:00:00 2001 From: Michael Pilosov Date: Tue, 2 Sep 2025 19:13:05 +0000 Subject: [PATCH 3/6] mypy silence --- src/litai/tools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/litai/tools.py b/src/litai/tools.py index c5af01a..3d5fbdf 100644 --- a/src/litai/tools.py +++ b/src/litai/tools.py @@ -126,7 +126,7 @@ def setup(self) -> None: self.name = model.__name__ self.description = model.__doc__ or "" - def run(self, *args, **kwargs) -> Any: + def run(self, *args, **kwargs) -> Any: # type: ignore # Default implementation: validate & return an instance return model(*args, **kwargs) From 08155b80fa7486224c622a266d1dd4bc58f4a76f Mon Sep 17 00:00:00 2001 From: Michael Pilosov Date: Tue, 2 Sep 2025 19:14:55 +0000 Subject: [PATCH 4/6] ruff format --- src/litai/tools.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/litai/tools.py b/src/litai/tools.py index 3d5fbdf..86c23a8 100644 --- a/src/litai/tools.py +++ b/src/litai/tools.py @@ -120,19 +120,20 @@ def _extract_parameters(self) -> Dict[str, Any]: @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 an instance return model(*args, **kwargs) - + def _extract_parameters(self) -> Dict[str, Any]: return model.model_json_schema() - + return ModelTool() @classmethod From 25d95518feb94aa67fad5b9239530d4c8cbdc007 Mon Sep 17 00:00:00 2001 From: Michael Pilosov Date: Sat, 6 Sep 2025 02:46:35 +0000 Subject: [PATCH 5/6] lint --- tests/test_tools.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_tools.py b/tests/test_tools.py index 66d8b92..0f425f0 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -24,7 +24,9 @@ def weather_tool_model(): class WeatherRequest(BaseModel): """Get weather for location.""" + location: str + return WeatherRequest From 4d066d12a21bbe7974da1b7055285586e7071d8d Mon Sep 17 00:00:00 2001 From: Michael Pilosov Date: Fri, 12 Sep 2025 15:09:25 +0000 Subject: [PATCH 6/6] typehint change --- src/litai/llm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/litai/llm.py b/src/litai/llm.py index a70e940..3d451ea 100644 --- a/src/litai/llm.py +++ b/src/litai/llm.py @@ -31,6 +31,7 @@ if TYPE_CHECKING: from langchain_core.tools import StructuredTool + from pydantic import BaseModel CLOUDY_MODELS = { "openai/gpt-4o", @@ -358,7 +359,7 @@ def chat( # noqa: D417 @staticmethod def call_tool( response: Union[List[dict], dict, str], tools: Optional[Sequence[Union[LitTool, "StructuredTool"]]] = None - ) -> Optional[str]: + ) -> Optional[Union[str, "BaseModel", list["BaseModel"]]]: """Calls a tool with the given response.""" if tools is None: raise ValueError("No tools provided")