refactor tool

Author: vilsonrodrigues

src/msgflux/nn/modules/__init__.py

@@ -6,15 +6,15 @@
 from msgflux.nn.modules.module import Module
 from msgflux.nn.modules.retriever import Retriever
 from msgflux.nn.modules.speaker import Speaker
-from msgflux.nn.modules.tool import Tool, ToolLibrary
-
-ToolBase = Tool
+from msgflux.nn.modules.tool import LocalTool, MCPTool, Tool, ToolLibrary
 from msgflux.nn.modules.transcriber import Transcriber
 
 __all__ = [
     "Agent",
     "Embedder",
     "LM",
+    "LocalTool",
+    "MCPTool",
     "MediaMaker",
     "Module",
     "ModuleDict",
@@ -23,7 +23,6 @@
     "Sequential",
     "Speaker",
     "Tool",
-    "ToolBase",
     "ToolLibrary",
     "Transcriber",
 ]

src/msgflux/nn/modules/lm.py

@@ -1,4 +1,4 @@
-from typing import Any, Union
+from typing import Union
 
 from msgflux.models.gateway import ModelGateway
 from msgflux.models.response import ModelResponse, ModelStreamResponse
@@ -44,4 +44,4 @@ def _set_model(self, model: Union[ChatCompletionModel, ModelGateway]):
             raise TypeError(
                 f"`model` must be a `chat_completion` model, given `{type(model)}`"
             )
-        self.register_buffer("model", model)
+        self.register_buffer("model", model)

src/msgflux/nn/modules/module.py

@@ -848,8 +848,9 @@ async def _aexecute_input_guardrail(self, model_execution_params: Dict[str, Any]
         elif inspect.iscoroutinefunction(input_guardrail):
             guardrail_response = await input_guardrail(**guardrail_params)
         else:
-            # Fallback to sync call
-            guardrail_response = input_guardrail(**guardrail_params)
+            # Fallback to sync call in executor to avoid blocking event loop
+            loop = asyncio.get_event_loop()
+            guardrail_response = await loop.run_in_executor(None, lambda: input_guardrail(**guardrail_params))
 
         if isinstance(guardrail_response, ModelResponse):
             guardrail_response = self._extract_raw_response(guardrail_response)
@@ -884,8 +885,9 @@ async def _aexecute_output_guardrail(self, model_response: Dict[str, Any]):
         elif inspect.iscoroutinefunction(output_guardrail):
             guardrail_response = await output_guardrail(**guardrail_params)
         else:
-            # Fallback to sync call
-            guardrail_response = output_guardrail(**guardrail_params)
+            # Fallback to sync call in executor to avoid blocking event loop
+            loop = asyncio.get_event_loop()
+            guardrail_response = await loop.run_in_executor(None, lambda: output_guardrail(**guardrail_params))
 
         if isinstance(guardrail_response, ModelResponse):
             guardrail_response = self._extract_raw_response(guardrail_response)

src/msgflux/nn/modules/tool.py

@@ -1,14 +1,12 @@
+import asyncio
 import inspect
 from dataclasses import asdict, dataclass, field
 from functools import partial
 from typing import Any, Callable, Dict, Iterator, List, Mapping, Optional, Union, Tuple
 
 import msgspec
-from opentelemetry import trace
-from opentelemetry.trace import Status, StatusCode
 
 from msgflux.dotdict import dotdict
-from msgflux.envs import envs
 from msgflux.logger import logger
 from msgflux.nn import functional as F
 from msgflux.nn.modules.container import ModuleDict
@@ -149,6 +147,41 @@ async def aforward(self, **kwargs) -> Any:
         # Extract and return result
         return extract_tool_result_text(result)
 
+class LocalTool(Tool):
+    """Local tool implementation."""
+    def __init__(
+        self,
+        name: str,
+        description: str,
+        annotations: Dict[str, Any],
+        tool_config: Dict[str, Any],
+        impl: Callable,
+    ):
+        super().__init__()
+        self.set_name(name)
+        self.set_description(description)
+        self.set_annotations(annotations)
+        self.register_buffer("tool_config", tool_config)
+        self.impl = impl  # Not a buffer for now
+
+    @tool_retry
+    @set_tool_attributes(execution_type="local")
+    def forward(self, **kwargs):
+        if inspect.iscoroutinefunction(self.impl):
+            return F.wait_for(self.impl, **kwargs)
+        return self.impl(**kwargs)
+
+    @tool_retry
+    @aset_tool_attributes(execution_type="local")
+    async def aforward(self, *args, **kwargs):
+        if hasattr(self.impl, "acall"):
+            return await self.impl.acall(*args, **kwargs)
+        elif inspect.iscoroutinefunction(self.impl):
+            return await self.impl(*args, **kwargs)
+        # Fall back to sync call in executor to avoid blocking event loop
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(None, lambda: self.impl(*args, **kwargs))
+
 def _convert_module_to_nn_tool(impl: Callable) -> Tool: # noqa: C901
     """Convert a callable in nn.Tool."""
     tool_config = impl.__dict__.get("tool_config", dotdict())
@@ -168,7 +201,7 @@ def _convert_module_to_nn_tool(impl: Callable) -> Tool: # noqa: C901
             or
             getattr(impl, "__doc__", None)
             or
-            getattr(impl.__call__, "__doc__", None)            
+            getattr(impl.__call__, "__doc__", None)
         )
         if doc is None:
             raise NotImplementedError(
@@ -183,7 +216,7 @@ def _convert_module_to_nn_tool(impl: Callable) -> Tool: # noqa: C901
             or
             getattr(impl, "__annotations__", None)
             or
-            getattr(impl.__call__, "__annotations__", None)            
+            getattr(impl.__call__, "__annotations__", None)
         )
         if annotations is None:
             if fn_has_parameters(impl.__call__):
@@ -217,7 +250,7 @@ def _convert_module_to_nn_tool(impl: Callable) -> Tool: # noqa: C901
             )
 
         annotations = impl.__annotations__
-        
+
         if annotations is None:
             if fn_has_parameters(impl):
                 raise NotImplementedError(
@@ -241,33 +274,13 @@ def _convert_module_to_nn_tool(impl: Callable) -> Tool: # noqa: C901
     if tool_config.get("background"):
         doc = "This tool will run in the background. \n" + doc
 
-    class LocalTool(Tool):
-        """Local tool implementation."""
-        def __init__(self):
-            super().__init__()
-            self.set_name(name)
-            self.set_description(doc)
-            self.set_annotations(annotations)
-            self.register_buffer("tool_config", tool_config)
-            self.impl = impl  # Not a buffer for now
-
-        @tool_retry
-        @set_tool_attributes(execution_type="local")
-        def forward(self, **kwargs):
-            if inspect.iscoroutinefunction(self.impl):
-                return F.wait_for(self.impl, **kwargs)
-            return self.impl(**kwargs)
-
-        @tool_retry
-        @aset_tool_attributes(execution_type="local")
-        async def aforward(self, *args, **kwargs):
-            if hasattr(self.impl, "acall"):
-                return await self.impl.acall(*args, **kwargs)
-            elif inspect.iscoroutinefunction(self.impl):
-                return await self.impl(*args, **kwargs)
-            # Fall back to sync call
-            return self.impl(*args, **kwargs)
-    return LocalTool()
+    return LocalTool(
+        name=name,
+        description=doc,
+        annotations=annotations,
+        tool_config=tool_config,
+        impl=impl,
+    )
 
 
 class ToolLibrary(Module):