typehint change

Author: mathematicalmichael

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")