chuk-tool-processor 0.1.6__py3-none-any.whl → 0.2__py3-none-any.whl

chuk_tool_processor/registry/decorators.py

@@ -1,24 +1,42 @@
 # chuk_tool_processor/registry/decorators.py
 """
-Decorators for registering tools with the registry.
+Decorators for registering tools with the registry asynchronously.
 """
 
-from functools import wraps
-from typing import Any, Callable, Dict, Optional, Type, TypeVar
+import asyncio
+import functools
+import inspect
+import sys
+import weakref
+import atexit
+import warnings
+from typing import Any, Callable, Dict, Optional, Type, TypeVar, cast, Set, List, Awaitable
 
 from chuk_tool_processor.registry.provider import ToolRegistryProvider
 
 T = TypeVar('T')
 
+# Global tracking of classes to be registered
+# Store coroutines rather than awaitables to avoid warnings
+_PENDING_REGISTRATIONS: List[Callable[[], Awaitable]] = []
+_REGISTERED_CLASSES = weakref.WeakSet()
+
+# Keep track of whether we're shutting down
+_SHUTTING_DOWN = False
+
 
 def register_tool(name: Optional[str] = None, namespace: str = "default", **metadata):
     """
     Decorator for registering tools with the global registry.
     
+    This decorator will queue the registration to happen asynchronously.
+    You must call `await ensure_registrations()` in your application startup
+    to complete all registrations.
+    
     Example:
         @register_tool(name="my_tool", namespace="math", description="Performs math operations")
         class MyTool:
-            def execute(self, x: int, y: int) -> int:
+            async def execute(self, x: int, y: int) -> int:
                 return x + y
     
     Args:
@@ -30,13 +48,118 @@ def register_tool(name: Optional[str] = None, namespace: str = "default", **meta
         A decorator function that registers the class with the registry.
     """
     def decorator(cls: Type[T]) -> Type[T]:
-        registry = ToolRegistryProvider.get_registry()
-        registry.register_tool(cls, name=name, namespace=namespace, metadata=metadata)
+        # Skip if already registered
+        if cls in _REGISTERED_CLASSES:
+            return cls
+            
+        # Skip if shutting down
+        if _SHUTTING_DOWN:
+            return cls
+
+        # Ensure execute method is async
+        if hasattr(cls, 'execute') and not inspect.iscoroutinefunction(cls.execute):
+            raise TypeError(f"Tool {cls.__name__} must have an async execute method")
+            
+        # Create registration function (not coroutine)
+        async def do_register():
+            registry = await ToolRegistryProvider.get_registry()
+            await registry.register_tool(
+                cls, 
+                name=name, 
+                namespace=namespace, 
+                metadata=metadata
+            )
         
-        @wraps(cls)
-        def wrapper(*args: Any, **kwargs: Dict[str, Any]) -> T:
-            return cls(*args, **kwargs)
+        # Store the function, not the coroutine
+        _PENDING_REGISTRATIONS.append(do_register)
+        _REGISTERED_CLASSES.add(cls)
+        
+        # Add class attribute so we can identify decorated classes
+        cls._tool_registration_info = {
+            'name': name or cls.__name__,
+            'namespace': namespace,
+            'metadata': metadata
+        }
+        
+        # Don't modify the original class
+        return cls
+    
+    return decorator
+
+
+async def ensure_registrations() -> None:
+    """
+    Process all pending tool registrations.
+    
+    This must be called during application startup to register
+    all tools decorated with @register_tool.
+    
+    Returns:
+        None
+    """
+    global _PENDING_REGISTRATIONS
+    
+    if not _PENDING_REGISTRATIONS:
+        return
         
-        return wrapper
+    # Create tasks from the stored functions
+    tasks = []
+    for registration_fn in _PENDING_REGISTRATIONS:
+        # Now we await the function to get the coroutine, then create a task
+        tasks.append(asyncio.create_task(registration_fn()))
+    
+    # Clear the pending list
+    _PENDING_REGISTRATIONS.clear()
     
-    return decorator
+    # Wait for all registrations to complete
+    if tasks:
+        await asyncio.gather(*tasks)
+
+
+def discover_decorated_tools() -> List[Type]:
+    """
+    Discover all tool classes decorated with @register_tool.
+    
+    This can be used to inspect what tools have been registered
+    without awaiting ensure_registrations().
+    
+    Returns:
+        List of tool classes that have been decorated
+    """
+    tools = []
+    
+    # Search all loaded modules
+    for module_name, module in list(sys.modules.items()):
+        if not module_name.startswith('chuk_tool_processor'):
+            continue
+            
+        for attr_name in dir(module):
+            try:
+                attr = getattr(module, attr_name)
+                if hasattr(attr, '_tool_registration_info'):
+                    tools.append(attr)
+            except (AttributeError, ImportError):
+                pass
+                
+    return tools
+
+
+# Register atexit handler to prevent warnings at shutdown
+def _handle_shutdown():
+    """
+    Handle shutdown by marking shutdown flag and clearing pending registrations.
+    This prevents warnings about unawaited coroutines.
+    """
+    global _SHUTTING_DOWN, _PENDING_REGISTRATIONS
+    
+    # Set the shutdown flag
+    _SHUTTING_DOWN = True
+    
+    # Clear without creating any coroutines
+    _PENDING_REGISTRATIONS = []
+
+# Register the shutdown handler
+atexit.register(_handle_shutdown)
+
+# Filter the coroutine never awaited warning
+warnings.filterwarnings("ignore", message="coroutine.*was never awaited")

chuk_tool_processor/registry/interface.py

@@ -1,19 +1,23 @@
 # chuk_tool_processor/registry/interface.py
 """
-Defines the interface for tool registries.
+Defines the interface for asynchronous tool registries.
 """
-from typing import Protocol, Any, Dict, List, Optional, Tuple
+from __future__ import annotations
 
-# imports
+from typing import Protocol, Any, Dict, List, Optional, Tuple, TypeVar, runtime_checkable
+
+# imports
 from chuk_tool_processor.registry.metadata import ToolMetadata
 
+T = TypeVar('T')
 
+@runtime_checkable
 class ToolRegistryInterface(Protocol):
     """
-    Protocol for a tool registry. Implementations should allow registering tools
+    Protocol for an async tool registry. Implementations should allow registering tools
     and retrieving them by name and namespace.
     """
-    def register_tool(
+    async def register_tool(
         self, 
         tool: Any, 
         name: Optional[str] = None,
@@ -21,7 +25,7 @@ class ToolRegistryInterface(Protocol):
         metadata: Optional[Dict[str, Any]] = None
     ) -> None:
         """
-        Register a tool implementation.
+        Register a tool implementation asynchronously.
 
         Args:
             tool: The tool class or instance with an `execute` method.
@@ -31,9 +35,9 @@ class ToolRegistryInterface(Protocol):
         """
         ...
 
-    def get_tool(self, name: str, namespace: str = "default") -> Optional[Any]:
+    async def get_tool(self, name: str, namespace: str = "default") -> Optional[Any]:
         """
-        Retrieve a registered tool by name and namespace.
+        Retrieve a registered tool by name and namespace asynchronously.
         
         Args:
             name: The name of the tool.
@@ -44,9 +48,25 @@ class ToolRegistryInterface(Protocol):
         """
         ...
 
-    def get_metadata(self, name: str, namespace: str = "default") -> Optional[ToolMetadata]:
+    async def get_tool_strict(self, name: str, namespace: str = "default") -> Any:
+        """
+        Retrieve a registered tool by name and namespace, raising if not found.
+        
+        Args:
+            name: The name of the tool.
+            namespace: The namespace of the tool (default: "default").
+            
+        Returns:
+            The tool implementation.
+            
+        Raises:
+            ToolNotFoundError: If the tool is not found in the registry.
+        """
+        ...
+
+    async def get_metadata(self, name: str, namespace: str = "default") -> Optional[ToolMetadata]:
         """
-        Retrieve metadata for a registered tool.
+        Retrieve metadata for a registered tool asynchronously.
         
         Args:
             name: The name of the tool.
@@ -57,9 +77,9 @@ class ToolRegistryInterface(Protocol):
         """
         ...
 
-    def list_tools(self, namespace: Optional[str] = None) -> List[Tuple[str, str]]:
+    async def list_tools(self, namespace: Optional[str] = None) -> List[Tuple[str, str]]:
         """
-        List all registered tool names, optionally filtered by namespace.
+        List all registered tool names asynchronously, optionally filtered by namespace.
         
         Args:
             namespace: Optional namespace filter.
@@ -69,11 +89,25 @@ class ToolRegistryInterface(Protocol):
         """
         ...
 
-    def list_namespaces(self) -> List[str]:
+    async def list_namespaces(self) -> List[str]:
         """
-        List all registered namespaces.
+        List all registered namespaces asynchronously.
         
         Returns:
             List of namespace names.
         """
+        ...
+        
+    async def list_metadata(self, namespace: Optional[str] = None) -> List[ToolMetadata]:
+        """
+        Return all ToolMetadata objects asynchronously.
+
+        Args:
+            namespace: Optional filter by namespace.
+                • None (default) – metadata from all namespaces
+                • "some_ns" – only that namespace
+
+        Returns:
+            List of ToolMetadata objects.
+        """
         ...

chuk_tool_processor/registry/metadata.py

@@ -1,9 +1,12 @@
 # chuk_tool_processor/registry/metadata.py
 """
-Tool metadata models for the registry.
+Tool metadata models for the registry with async-native support.
 """
-from typing import Any, Dict, Optional, Set
-from pydantic import BaseModel, Field
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Set, List, Union
+from datetime import datetime
+from pydantic import BaseModel, Field, model_validator
 
 
 class ToolMetadata(BaseModel):
@@ -20,17 +23,60 @@ class ToolMetadata(BaseModel):
         result_schema: Optional schema for the tool's result.
         requires_auth: Whether the tool requires authentication.
         tags: Set of tags associated with the tool.
+        created_at: When the tool was first registered.
+        updated_at: When the tool was last updated.
+        source: Optional source information (e.g., "function", "class", "langchain").
+        source_name: Optional source identifier.
+        concurrency_limit: Optional maximum concurrent executions.
+        timeout: Optional default timeout in seconds.
+        rate_limit: Optional rate limiting configuration.
     """
     name: str = Field(..., description="Tool name")
     namespace: str = Field("default", description="Namespace the tool belongs to")
     description: Optional[str] = Field(None, description="Tool description")
     version: str = Field("1.0.0", description="Tool implementation version")
-    is_async: bool = Field(False, description="Whether the tool's execute method is asynchronous")
+    is_async: bool = Field(True, description="Whether the tool's execute method is asynchronous")
     argument_schema: Optional[Dict[str, Any]] = Field(None, description="Schema for the tool's arguments")
     result_schema: Optional[Dict[str, Any]] = Field(None, description="Schema for the tool's result")
     requires_auth: bool = Field(False, description="Whether the tool requires authentication")
     tags: Set[str] = Field(default_factory=set, description="Tags associated with the tool")
-
+    created_at: datetime = Field(default_factory=datetime.utcnow, description="When the tool was first registered")
+    updated_at: datetime = Field(default_factory=datetime.utcnow, description="When the tool was last updated")
+    source: Optional[str] = Field(None, description="Source of the tool (e.g., 'function', 'class', 'langchain')")
+    source_name: Optional[str] = Field(None, description="Source identifier (e.g., function name, class name)")
+    concurrency_limit: Optional[int] = Field(None, description="Maximum concurrent executions (None = unlimited)")
+    timeout: Optional[float] = Field(None, description="Default timeout in seconds (None = no timeout)")
+    rate_limit: Optional[Dict[str, Any]] = Field(None, description="Rate limiting configuration")
+    
+    # Additional fields for async-native architecture
+    supports_streaming: bool = Field(False, description="Whether the tool supports streaming responses")
+    execution_options: Dict[str, Any] = Field(default_factory=dict, description="Additional execution options")
+    dependencies: List[str] = Field(default_factory=list, description="Dependencies on other tools")
+    
+    @model_validator(mode='after')
+    def ensure_async(self) -> 'ToolMetadata':
+        """Ensure all tools are marked as async in the async-native architecture."""
+        self.is_async = True
+        return self
+    
+    def with_updated_timestamp(self) -> 'ToolMetadata':
+        """Create a copy with updated timestamp."""
+        return self.model_copy(update={"updated_at": datetime.utcnow()})
+    
     def __str__(self) -> str:
         """String representation of the tool metadata."""
-        return f"{self.namespace}.{self.name} (v{self.version})"
+        return f"{self.namespace}.{self.name} (v{self.version})"
+
+
+class RateLimitConfig(BaseModel):
+    """Rate limiting configuration for tools."""
+    requests: int = Field(..., description="Maximum number of requests")
+    period: float = Field(..., description="Time period in seconds")
+    scope: str = Field("global", description="Scope of rate limiting: 'global', 'user', 'ip'")
+
+
+class StreamingToolMetadata(ToolMetadata):
+    """Extended metadata for tools that support streaming responses."""
+    supports_streaming: bool = Field(True, description="Whether the tool supports streaming responses")
+    chunk_size: Optional[int] = Field(None, description="Suggested chunk size for streaming")
+    content_type: Optional[str] = Field(None, description="Content type for streaming responses")

chuk_tool_processor/registry/provider.py

@@ -1,5 +1,6 @@
+# chuk_tool_processor/registry/provider.py
 """
-Global access to *the* tool registry instance.
+Global access to the async tool registry instance.
 
 There are two public faces:
 
@@ -8,92 +9,130 @@ There are two public faces:
       and memoises it in the module-level variable ``_REGISTRY``.  
     • `set_registry()` lets callers replace or reset that singleton.
 
-2.  **`ToolRegistryProvider` shim**  
-    Earlier versions exposed a static wrapper.  Tests rely on being able to
-    monkey-patch the *module-level* factory and to clear the cached instance
-    by setting `ToolRegistryProvider._registry = None`.  We therefore keep a
-    **separate class-level cache** (`_registry`) and call the *current*
-    module-level `get_registry()` **only when the cache is empty**.
+2.  **`ToolRegistryProvider` class**  
+    Provides static methods for async-safe access to the registry.
 
 The contract verified by the test-suite is:
 
 * The module-level factory is invoked **exactly once** per fresh cache.
-* `ToolRegistryProvider.set_registry(obj)` overrides subsequent retrievals.
-* `ToolRegistryProvider.set_registry(None)` resets the cache so the next
-  `get_registry()` call invokes (and honours any monkey-patched) factory.
+* `await ToolRegistryProvider.set_registry(obj)` overrides subsequent retrievals.
+* `await ToolRegistryProvider.set_registry(None)` resets the cache so the next
+  `await get_registry()` call invokes (and honours any monkey-patched) factory.
 """
 from __future__ import annotations
 
-from typing import Optional
+import asyncio
+import importlib
+import sys
+from typing import Optional, Callable, Awaitable, Dict, Any
 
+# registry
 from .interface import ToolRegistryInterface
-from .providers.memory import InMemoryToolRegistry
 
 # --------------------------------------------------------------------------- #
 # Module-level singleton used by the helper functions
 # --------------------------------------------------------------------------- #
 _REGISTRY: Optional[ToolRegistryInterface] = None
+_REGISTRY_LOCK = asyncio.Lock()
 # --------------------------------------------------------------------------- #
 
 
-def _default_registry() -> ToolRegistryInterface:
-    """Create the default in-memory registry."""
+async def _default_registry() -> ToolRegistryInterface:
+    """Create the default in-memory registry asynchronously."""
+    # Import here to avoid circular import
+    from .providers.memory import InMemoryToolRegistry
     return InMemoryToolRegistry()
 
 
-def get_registry() -> ToolRegistryInterface:
+async def get_registry() -> ToolRegistryInterface:
     """
-    Return the process-wide registry, creating it on first use.
-
-    This function *may* be monkey-patched in tests; call it via
-    ``globals()["get_registry"]()`` if you need the latest binding.
+    Return the process-wide registry asynchronously, creating it on first use.
+    
+    This function is thread-safe and will only create the registry once,
+    even with concurrent calls.
     """
     global _REGISTRY
     if _REGISTRY is None:
-        _REGISTRY = _default_registry()
+        async with _REGISTRY_LOCK:
+            # Double-check pattern: check again after acquiring the lock
+            if _REGISTRY is None:
+                _REGISTRY = await _default_registry()
     return _REGISTRY
 
 
-def set_registry(registry: ToolRegistryInterface | None) -> None:
+async def set_registry(registry: ToolRegistryInterface | None) -> None:
     """
-    Replace or clear the global registry.
+    Replace or clear the global registry asynchronously.
 
     Passing ``None`` resets the singleton so that the next `get_registry()`
     call recreates it (useful in tests).
     """
     global _REGISTRY
-    _REGISTRY = registry
+    async with _REGISTRY_LOCK:
+        _REGISTRY = registry
 
 
 # --------------------------------------------------------------------------- #
-# Back-compat shim used by legacy import paths and the test-suite
+# Provider class for consistent access to the registry
 # --------------------------------------------------------------------------- #
-class ToolRegistryProvider:                       # noqa: D401
-    """Legacy static wrapper retaining historical semantics."""
+class ToolRegistryProvider:
+    """Async static wrapper for registry access."""
 
-    # The test-suite directly mutates this attribute, so we keep it.
+    # Thread-safe singleton management
     _registry: Optional[ToolRegistryInterface] = None
+    _lock = asyncio.Lock()
 
     # ------------------------ public API ------------------------ #
     @staticmethod
-    def get_registry() -> ToolRegistryInterface:
+    async def get_registry() -> ToolRegistryInterface:
         """
-        Return the cached instance or, if absent, call the *current*
-        module-level `get_registry()` exactly once to populate it.
+        Return the cached instance or initialize a new one asynchronously.
+        
+        This method ensures thread-safety when initializing the registry.
         """
         if ToolRegistryProvider._registry is None:
-            # Honour any runtime monkey-patching of the factory.
-            ToolRegistryProvider._registry = globals()["get_registry"]()
+            async with ToolRegistryProvider._lock:
+                # Check again after acquiring the lock
+                if ToolRegistryProvider._registry is None:
+                    # Dynamically import to get the latest definition
+                    module = sys.modules[__name__]
+                    get_registry_func: Callable[[], Awaitable[ToolRegistryInterface]] = getattr(
+                        module, "get_registry"
+                    )
+                    # Call it to get the registry
+                    ToolRegistryProvider._registry = await get_registry_func()
+        
         return ToolRegistryProvider._registry
 
     @staticmethod
-    def set_registry(registry: ToolRegistryInterface | None) -> None:
+    async def set_registry(registry: ToolRegistryInterface | None) -> None:
         """
-        Override the cached registry.
+        Override the cached registry asynchronously.
 
         *   If ``registry`` is an object, all subsequent `get_registry()`
             calls return it without touching the factory.
         *   If ``registry`` is ``None``, the cache is cleared so the next
-            `get_registry()` call invokes the (possibly patched) factory.
+            `get_registry()` call invokes the factory.
+        """
+        async with ToolRegistryProvider._lock:
+            ToolRegistryProvider._registry = registry
+            
+    @staticmethod
+    async def reset() -> None:
+        """
+        Reset both the module-level and class-level registry caches.
+        
+        This is primarily used in tests to ensure a clean state.
+        """
+        async with ToolRegistryProvider._lock:
+            ToolRegistryProvider._registry = None
+            await set_registry(None)
+            
+    @staticmethod
+    async def get_global_registry() -> ToolRegistryInterface:
+        """
+        Get the module-level registry directly.
+        
+        This bypasses the class-level cache and always returns the module-level registry.
         """
-        ToolRegistryProvider._registry = registry
+        return await get_registry()

chuk_tool_processor/registry/providers/__init__.py

@@ -1,20 +1,29 @@
+# chuk_tool_processor/registry/providers/__init__.py
 """
-Registry provider implementations and factory functions.
+Async registry provider implementations and factory functions.
 """
 
 import os
-from typing import Optional
+import asyncio
+from typing import Optional, Dict, Any
 
 from chuk_tool_processor.registry.interface import ToolRegistryInterface
-from chuk_tool_processor.registry.providers.memory import InMemoryToolRegistry
 
+# Cache for initialized registries
+_REGISTRY_CACHE: Dict[str, ToolRegistryInterface] = {}
+_REGISTRY_LOCKS: Dict[str, asyncio.Lock] = {}
 
-def get_registry(
+
+async def get_registry(
     provider_type: Optional[str] = None, 
     **kwargs
 ) -> ToolRegistryInterface:
     """
-    Factory function to get a registry implementation.
+    Factory function to get a registry implementation asynchronously.
+    
+    This function caches registry instances by provider_type to avoid
+    creating multiple instances unnecessarily. The cache is protected
+    by locks to ensure thread safety.
     
     Args:
         provider_type: Type of registry provider to use. Options:
@@ -34,8 +43,38 @@ def get_registry(
     if provider_type is None:
         provider_type = os.environ.get("CHUK_TOOL_REGISTRY_PROVIDER", "memory")
     
-    # Create the appropriate provider
-    if provider_type == "memory":
-        return InMemoryToolRegistry()
-    else:
-        raise ValueError(f"Unknown registry provider type: {provider_type}")
+    # Check cache first
+    cache_key = f"{provider_type}:{hash(frozenset(kwargs.items()))}"
+    if cache_key in _REGISTRY_CACHE:
+        return _REGISTRY_CACHE[cache_key]
+    
+    # Create lock if needed
+    if cache_key not in _REGISTRY_LOCKS:
+        _REGISTRY_LOCKS[cache_key] = asyncio.Lock()
+    
+    # Acquire lock to ensure only one registry is created
+    async with _REGISTRY_LOCKS[cache_key]:
+        # Double-check pattern: check cache again after acquiring lock
+        if cache_key in _REGISTRY_CACHE:
+            return _REGISTRY_CACHE[cache_key]
+            
+        # Create the appropriate provider
+        if provider_type == "memory":
+            # Import here to avoid circular imports
+            from chuk_tool_processor.registry.providers.memory import InMemoryToolRegistry
+            registry = InMemoryToolRegistry()
+        else:
+            raise ValueError(f"Unknown registry provider type: {provider_type}")
+        
+        # Cache the registry
+        _REGISTRY_CACHE[cache_key] = registry
+        return registry
+
+
+async def clear_registry_cache() -> None:
+    """
+    Clear the registry cache.
+    
+    This is useful in tests or when configuration changes.
+    """
+    _REGISTRY_CACHE.clear()