Fix support for applying tool_method to pydantic base model classes. Migrate a bunch of optional import logic to be lazy. Clean up dependencies and typing for optional packages. (#237)

Author: monoxgas

.hooks/generate_docs.py

@@ -3,8 +3,8 @@
 import typing as t
 from pathlib import Path
 
-from markdown import Markdown  # type: ignore[import-untyped]
-from markdownify import MarkdownConverter  # type: ignore[import-untyped]
+from markdown import Markdown  # type: ignore [import-untyped]
+from markdownify import MarkdownConverter  # type: ignore [import-untyped]
 from markupsafe import Markup
 from mkdocstrings_handlers.python._internal.config import PythonConfig
 from mkdocstrings_handlers.python._internal.handler import (
@@ -14,7 +14,7 @@
 # ruff: noqa: T201
 
 
-class CustomMarkdownConverter(MarkdownConverter):  # type: ignore[misc]
+class CustomMarkdownConverter(MarkdownConverter):  # type: ignore [misc]
     # Strip extra whitespace from code blocks
     def convert_pre(self, el: t.Any, text: str, parent_tags: t.Any) -> t.Any:
         return super().convert_pre(el, text.strip(), parent_tags)

docs/api/data.mdx

@@ -58,7 +58,7 @@ chat\_id column.
 
 <Accordion title="Source code in rigging/data.py" icon="code">
 ```python
-def chats_to_df(chats: Chat | t.Sequence[Chat]) -> pd.DataFrame:
+def chats_to_df(chats: Chat | t.Sequence[Chat]) -> "pd.DataFrame":
     """
     Convert a Chat or list of Chat objects into a pandas DataFrame.
 
@@ -73,6 +73,13 @@ def chats_to_df(chats: Chat | t.Sequence[Chat]) -> pd.DataFrame:
         A pandas DataFrame containing the chat data.
 
     """
+    try:
+        import pandas as pd
+    except ImportError as e:
+        raise ImportError(
+            "Pandas is not available. Please install `pandas` or use `rigging[data]`.",
+        ) from e
+
     chats = [chats] if isinstance(chats, Chat) else chats
 
     flattened = flatten_chats(chats)
@@ -176,10 +183,10 @@ async def chats_to_elastic(
         The indexed count from the bulk operation
     """
     try:
-        import elasticsearch.helpers
+        import elasticsearch.helpers  # type: ignore [import-not-found, unused-ignore]
     except ImportError as e:
         raise ImportError(
-            "Elasticsearch is not available. Please install `elasticsearch` or use `rigging[extra]`.",
+            "Elasticsearch is not available. Please install `elasticsearch` or use `rigging[data]`.",
         ) from e
 
     es_data = chats_to_elastic_data(chats, index, op_type=op_type)
@@ -190,7 +197,7 @@ async def chats_to_elastic(
             await client.indices.put_mapping(index=index, properties=ElasticMapping["properties"])
 
     results = await elasticsearch.helpers.async_bulk(client, es_data, **kwargs)
-    return results[0]  # Return modified count
+    return results[0]  # type: ignore [no-any-return, unused-ignore]
 ```
 
 
@@ -286,7 +293,7 @@ generated by the `chats_to_df` function.
 
 <Accordion title="Source code in rigging/data.py" icon="code">
 ```python
-def df_to_chats(df: pd.DataFrame) -> list[Chat]:
+def df_to_chats(df: "pd.DataFrame") -> list[Chat]:
     """
     Convert a pandas DataFrame into a list of Chat objects.
 
@@ -301,6 +308,7 @@ def df_to_chats(df: pd.DataFrame) -> list[Chat]:
         A list of Chat objects.
 
     """
+
     chats = []
     for chat_id, chat_group in df.groupby("chat_id"):
         chat_data = chat_group.iloc[0]
@@ -564,7 +572,7 @@ Determine if an S3 bucket exists.
 
 <Accordion title="Source code in rigging/data.py" icon="code">
 ```python
-async def s3_bucket_exists(client: S3Client, bucket: str) -> bool:
+async def s3_bucket_exists(client: "S3Client", bucket: str) -> bool:
     """
     Determine if an S3 bucket exists.
 
@@ -618,7 +626,7 @@ Determine if an S3 object exists.
 
 <Accordion title="Source code in rigging/data.py" icon="code">
 ```python
-async def s3_object_exists(client: S3Client, bucket: str, key: str) -> bool:
+async def s3_object_exists(client: "S3Client", bucket: str, key: str) -> bool:
     """
     Determine if an S3 object exists.
 

docs/api/logging.mdx

@@ -32,7 +32,7 @@ configure\_logging
 
 ```python
 configure_logging(
-    log_level: LogLevelLiteral,
+    log_level: LogLevelLiteral = "info",
     log_file: Path | None = None,
     log_file_level: LogLevelLiteral = "debug",
 ) -> None
@@ -43,7 +43,9 @@ Configures common loguru handlers.
 **Parameters:**
 
 * **`log_level`**
-  (`LogLevelLiteral`)
+  (`LogLevelLiteral`, default:
+  `'info'`
+  )
   –The desired log level.
 * **`log_file`**
   (`Path | None`, default:
@@ -60,7 +62,7 @@ Configures common loguru handlers.
 <Accordion title="Source code in rigging/logging.py" icon="code">
 ```python
 def configure_logging(
-    log_level: LogLevelLiteral,
+    log_level: LogLevelLiteral = "info",
     log_file: pathlib.Path | None = None,
     log_file_level: LogLevelLiteral = "debug",
 ) -> None:

docs/api/tools.mdx

@@ -41,10 +41,9 @@ Base class for representing a tool to a generator.
 ### catch
 
 ```python
-catch: bool | set[type[Exception]] = {
-    JSONDecodeError,
-    ValidationError,
-}
+catch: bool | set[type[Exception]] = set(
+    DEFAULT_CATCH_EXCEPTIONS
+)
 ```
 
 Whether to catch exceptions and return them as messages.
@@ -207,6 +206,8 @@ async def handle_tool_call(  # noqa: PLR0912
             kwargs = json.loads(tool_call.function.arguments)
             if self._type_adapter is not None:
                 kwargs = self._type_adapter.validate_python(kwargs)
+            kwargs = kwargs or {}
+
             dn.log_inputs(**kwargs)
 
             # Call the function
@@ -367,7 +368,51 @@ def with_(
 ToolMethod
 ----------
 
-A Tool wrapping a class method.
+```python
+ToolMethod(
+    fget: Callable[..., Any],
+    name: str,
+    description: str,
+    parameters_schema: dict[str, Any],
+    catch: bool | Iterable[type[Exception]] | None,
+    truncate: int | None,
+    signature: Signature,
+    type_adapter: TypeAdapter[Any],
+)
+```
+
+A descriptor that acts as a factory for creating bound Tool instances.
+
+It inherits from `property` to be ignored by pydantic's `ModelMetaclass`
+during field inspection. This prevents validation errors which would
+otherwise treat the descriptor as a field and stop tool\_method decorators
+from being applied in BaseModel classes.
+
+<Accordion title="Source code in rigging/tools/base.py" icon="code">
+```python
+def __init__(
+    self,
+    fget: t.Callable[..., t.Any],
+    name: str,
+    description: str,
+    parameters_schema: dict[str, t.Any],
+    catch: bool | t.Iterable[type[Exception]] | None,
+    truncate: int | None,
+    signature: inspect.Signature,
+    type_adapter: TypeAdapter[t.Any],
+):
+    super().__init__(fget)
+    self.tool_name = name
+    self.tool_description = description
+    self.tool_parameters_schema = parameters_schema
+    self.tool_catch = catch
+    self.tool_truncate = truncate
+    self._tool_signature = signature
+    self._tool_type_adapter = type_adapter
+```
+
+
+</Accordion>
 
 tool
 ----
@@ -621,41 +666,38 @@ def tool_method(
         ~~~
     """
 
-    def make_tool(func: t.Callable[..., t.Any]) -> ToolMethod[P, R]:
-        # TODO: Improve consistency of detection here before enabling this warning
-        # if not _is_unbound_method(func):
-        #     warnings.warn(
-        #         "Passing a regular function to @tool_method improperly handles the 'self' argument, use @tool instead.",
-        #         SyntaxWarning,
-        #         stacklevel=3,
-        #     )
+    def make_tool(f: t.Callable[t.Concatenate[t.Any, P], R]) -> ToolMethod[P, R]:
+        # This logic is specialized from `Tool.from_callable` to correctly
+        # handle the `self` parameter in method signatures.
 
-        # Strip the `self` argument from the function signature so
-        # our schema generation doesn't include it under the hood.
+        signature = inspect.signature(f)
+        params_without_self = [p for p_name, p in signature.parameters.items() if p_name != "self"]
+        schema_signature = signature.replace(parameters=params_without_self)
 
-        @functools.wraps(func)
-        def wrapper(self: t.Any, *args: P.args, **kwargs: P.kwargs) -> R:
-            return func(self, *args, **kwargs)  # type: ignore [no-any-return]
+        @functools.wraps(f)
+        def empty_func(*_: t.Any, **kwargs: t.Any) -> t.Any:
+            return kwargs
 
-        wrapper.__signature__ = inspect.signature(func).replace(  # type: ignore [attr-defined]
-            parameters=tuple(
-                param
-                for param in inspect.signature(func).parameters.values()
-                if param.name != "self"
-            ),
-        )
+        empty_func.__signature__ = schema_signature  # type: ignore [attr-defined]
+        type_adapter: TypeAdapter[t.Any] = TypeAdapter(empty_func)
+        schema = deref_json(type_adapter.json_schema(), is_json_schema=True)
 
-        return ToolMethod.from_callable(
-            wrapper,  # type: ignore [arg-type]
-            name=name,
-            description=description,
+        tool_name = name or f.__name__
+        tool_description = inspect.cleandoc(description or f.__doc__ or "")
+
+        return ToolMethod(
+            fget=f,
+            name=tool_name,
+            description=tool_description,
+            parameters_schema=schema,
             catch=catch,
             truncate=truncate,
+            signature=schema_signature,
+            type_adapter=type_adapter,
         )
 
     if func is not None:
         return make_tool(func)
-
     return make_tool
 ```
 
@@ -677,7 +719,7 @@ A client for communicating with MCP servers.
 
 <Accordion title="Source code in rigging/tools/mcp.py" icon="code">
 ```python
-def __init__(self, transport: Transport, connection: StdioConnection | SSEConnection) -> None:
+def __init__(self, transport: Transport, connection: "StdioConnection | SSEConnection") -> None:
     self.transport = transport
     self.connection = connection
     self.tools = []
@@ -793,6 +835,9 @@ def as_mcp(
             )
         ~~~
     """
+    from mcp.server.fastmcp import FastMCP
+    from mcp.server.fastmcp.tools import Tool as FastMCPTool
+
     rigging_tools: list[Tool[..., t.Any]] = []
     for tool in flatten_list(list(tools)):
         interior_tools = [

docs/api/watchers.mdx

@@ -343,7 +343,7 @@ Create a watcher to write each chat to an Amazon S3 bucket.
 <Accordion title="Source code in rigging/watchers.py" icon="code">
 ```python
 def write_chats_to_s3(
-    client: S3Client,
+    client: "S3Client",
     bucket: str,
     key: str,
     *,