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

chuk_tool_processor/models/streaming_tool.py

@@ -0,0 +1,110 @@
+#!/usr/bin/env python
+# chuk_tool_processor/models/streaming_tool.py
+"""
+Base class for tools that support streaming results.
+
+This enables tools to yield incremental results during their execution,
+which is useful for long-running operations or real-time data processing.
+"""
+from __future__ import annotations
+
+import asyncio
+from abc import abstractmethod
+from typing import Any, AsyncIterator, List, TypeVar, Generic, ClassVar, Optional, Dict
+
+from pydantic import BaseModel, ConfigDict
+
+from chuk_tool_processor.models.validated_tool import ValidatedTool
+
+T = TypeVar('T')
+
+class StreamingTool(ValidatedTool):
+    """
+    Base class for tools that support streaming responses.
+    
+    Subclasses must implement _stream_execute which yields results one by one.
+    The executor should use stream_execute to access streaming results directly.
+    
+    Example:
+    ```python
+    class Counter(StreamingTool):
+        class Arguments(BaseModel):
+            count: int = 10
+            delay: float = 0.5
+            
+        class Result(BaseModel):
+            value: int
+            
+        async def _stream_execute(self, count: int, delay: float) -> AsyncIterator[Result]:
+            for i in range(count):
+                await asyncio.sleep(delay)
+                yield self.Result(value=i)
+    ```
+    
+    Streaming usage:
+    ```python
+    counter_tool = Counter()
+    async for result in counter_tool.stream_execute(count=5, delay=0.1):
+        print(f"Count: {result.value}")
+    ```
+    """
+    # Mark this as a ClassVar so Pydantic doesn't treat it as a field
+    supports_streaming: ClassVar[bool] = True
+    
+    # Use ConfigDict to configure model behavior
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    
+    async def stream_execute(self, **kwargs: Any) -> AsyncIterator[Any]:
+        """
+        Execute the tool and stream results incrementally.
+        
+        This public method validates arguments and then delegates to _stream_execute.
+        It should be used directly by the executor to support true streaming.
+        
+        Args:
+            **kwargs: Keyword arguments for the tool
+            
+        Yields:
+            Results as they are generated by the tool
+        """
+        # Validate arguments using the Arguments model
+        args = self.Arguments(**kwargs)
+        
+        # Stream results directly from _stream_execute
+        async for result in self._stream_execute(**args.model_dump()):
+            yield result
+    
+    async def execute(self, **kwargs: Any) -> Any:
+        """
+        Execute the tool and collect all results.
+        
+        For streaming tools, this collects all results from stream_execute
+        into a list for compatibility with the regular execution model.
+        
+        Args:
+            **kwargs: Keyword arguments for the tool
+            
+        Returns:
+            List of all streamed results
+        """
+        # Collect all streamed results into a list
+        results = []
+        async for chunk in self.stream_execute(**kwargs):
+            results.append(chunk)
+            
+        return results
+    
+    @abstractmethod
+    async def _stream_execute(self, **kwargs: Any) -> AsyncIterator[Any]:
+        """
+        Execute the tool and yield results incrementally.
+        
+        This must be implemented by streaming tool subclasses.
+        
+        Args:
+            **kwargs: Tool-specific arguments
+            
+        Yields:
+            Results as they are generated
+        """
+        yield NotImplemented

chuk_tool_processor/models/tool_call.py

@@ -1,7 +1,59 @@
 # chuk_tool_processor/models/tool_call.py
-from pydantic import BaseModel, Field
-from typing import Any, Dict
+"""
+Model representing a tool call with arguments.
+"""
+from __future__ import annotations
+
+import uuid
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel, Field, ConfigDict
 
 class ToolCall(BaseModel):
-    tool: str = Field(..., min_length=1, description="Name of the tool to call; must be non‐empty")
-    arguments: Dict[str, Any] = Field(default_factory=dict)
+    """
+    Represents a call to a tool with arguments.
+    
+    Attributes:
+        id: Unique identifier for the tool call
+        tool: Name of the tool to call
+        namespace: Namespace the tool belongs to
+        arguments: Arguments to pass to the tool
+    """
+    model_config = ConfigDict(extra='ignore')
+    
+    id: str = Field(
+        default_factory=lambda: str(uuid.uuid4()),
+        description="Unique identifier for the tool call"
+    )
+    tool: str = Field(
+        ..., 
+        min_length=1,
+        description="Name of the tool to call; must be non-empty"
+    )
+    namespace: str = Field(
+        default="default",
+        description="Namespace the tool belongs to"
+    )
+    arguments: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Arguments to pass to the tool"
+    )
+    
+    async def to_dict(self) -> Dict[str, Any]:
+        """Convert to a dictionary for serialization."""
+        return {
+            "id": self.id,
+            "tool": self.tool,
+            "namespace": self.namespace,
+            "arguments": self.arguments
+        }
+    
+    @classmethod
+    async def from_dict(cls, data: Dict[str, Any]) -> ToolCall:
+        """Create a ToolCall from a dictionary."""
+        return cls(**data)
+    
+    def __str__(self) -> str:
+        """String representation of the tool call."""
+        args_str = ", ".join(f"{k}={v!r}" for k, v in self.arguments.items())
+        return f"ToolCall({self.tool}, {args_str})"

chuk_tool_processor/models/tool_result.py

@@ -1,23 +1,44 @@
 # chuk_tool_processor/models/tool_result.py
+"""
+Model representing the result of a tool execution.
+"""
+from __future__ import annotations
+
 import os
-from pydantic import BaseModel, Field, ConfigDict
-from typing import Any, Optional
+import uuid
 from datetime import datetime, timezone
+from typing import Any, Dict, Optional, List
+
+from pydantic import BaseModel, Field, ConfigDict
 
 class ToolResult(BaseModel):
     """
     Represents the result of executing a tool.
-    Includes timing, host, and process metadata for diagnostics.
+    
+    Includes timing, host, and process metadata for diagnostics and tracing.
+    
+    Attributes:
+        id: Unique identifier for the result
+        tool: Name of the tool that was executed
+        result: Return value from the tool execution
+        error: Error message if execution failed
+        start_time: UTC timestamp when execution started
+        end_time: UTC timestamp when execution finished
+        machine: Hostname where the tool ran
+        pid: Process ID of the worker
+        cached: Flag indicating if the result was retrieved from cache
+        attempts: Number of execution attempts made
+        stream_id: Optional identifier for streaming results
+        is_partial: Whether this is a partial streaming result
     """
-    # Configure Pydantic to ignore any extra fields
     model_config = ConfigDict(extra='ignore')
 
-    # Flag indicating whether this result was retrieved from cache
-    cached: bool = Field(
-        default=False,
-        description="True if this result was retrieved from cache"
+    id: str = Field(
+        default_factory=lambda: str(uuid.uuid4()),
+        description="Unique identifier for this result"
     )
-
+    
+    # Core fields
     tool: str = Field(
         ...,
         min_length=1,
@@ -31,6 +52,8 @@ class ToolResult(BaseModel):
         None,
         description="Error message if execution failed"
     )
+    
+    # Execution metadata
     start_time: datetime = Field(
         default_factory=lambda: datetime.now(timezone.utc),
         description="UTC timestamp when execution started"
@@ -47,3 +70,86 @@ class ToolResult(BaseModel):
         default_factory=lambda: os.getpid(),
         description="Process ID of the worker"
     )
+    
+    # Extended features
+    cached: bool = Field(
+        default=False,
+        description="True if this result was retrieved from cache"
+    )
+    attempts: int = Field(
+        default=1,
+        description="Number of execution attempts made"
+    )
+    
+    # Streaming support
+    stream_id: Optional[str] = Field(
+        default=None,
+        description="Identifier for this stream of results (for streaming tools)"
+    )
+    is_partial: bool = Field(
+        default=False,
+        description="True if this is a partial result in a stream"
+    )
+    
+    @property
+    def is_success(self) -> bool:
+        """Check if the execution was successful (no error)."""
+        return self.error is None
+    
+    @property
+    def duration(self) -> float:
+        """Calculate the execution duration in seconds."""
+        if not self.start_time or not self.end_time:
+            return 0.0
+        return (self.end_time - self.start_time).total_seconds()
+    
+    async def to_dict(self) -> Dict[str, Any]:
+        """Convert to a dictionary for serialization."""
+        return {
+            "id": self.id,
+            "tool": self.tool,
+            "result": self.result,
+            "error": self.error,
+            "success": self.is_success,
+            "duration": self.duration,
+            "start_time": self.start_time.isoformat(),
+            "end_time": self.end_time.isoformat(),
+            "machine": self.machine,
+            "pid": self.pid,
+            "cached": self.cached,
+            "attempts": self.attempts,
+            "stream_id": self.stream_id,
+            "is_partial": self.is_partial
+        }
+    
+    @classmethod
+    def create_stream_chunk(
+        cls, 
+        tool: str, 
+        result: Any, 
+        stream_id: Optional[str] = None
+    ) -> ToolResult:
+        """Create a partial streaming result."""
+        stream_id = stream_id or str(uuid.uuid4())
+        return cls(
+            tool=tool,
+            result=result,
+            stream_id=stream_id,
+            is_partial=True
+        )
+    
+    @classmethod
+    async def from_dict(cls, data: Dict[str, Any]) -> ToolResult:
+        """Create a ToolResult from a dictionary."""
+        # Handle datetime fields
+        if isinstance(data.get("start_time"), str):
+            data["start_time"] = datetime.fromisoformat(data["start_time"])
+        if isinstance(data.get("end_time"), str):
+            data["end_time"] = datetime.fromisoformat(data["end_time"])
+            
+        return cls(**data)
+    
+    def __str__(self) -> str:
+        """String representation of the tool result."""
+        status = "success" if self.is_success else f"error: {self.error}"
+        return f"ToolResult({self.tool}, {status}, duration={self.duration:.3f}s)"

chuk_tool_processor/models/validated_tool.py

@@ -1,6 +1,6 @@
 # chuk_tool_processor/models/validated_tool.py
 """
-Self-contained base-class for *declarative* tools.
+Self-contained base-class for *declarative* async-native tools.
 
 Subclass it like so:
 
@@ -12,7 +12,7 @@ Subclass it like so:
         class Result(BaseModel):
             sum: int
 
-        def _execute(self, *, x: int, y: int) -> Result:
+        async def _execute(self, *, x: int, y: int) -> Result:
             return self.Result(sum=x + y)
 """
 from __future__ import annotations
@@ -93,7 +93,7 @@ class _ExportMixin:
 # The public validated base-class
 # --------------------------------------------------------------------------- #
 class ValidatedTool(_ExportMixin, BaseModel):
-    """Pydantic-validated base for new tools."""
+    """Pydantic-validated base for new async-native tools."""
 
     # ------------------------------------------------------------------ #
     # Inner models – override in subclasses
@@ -107,11 +107,11 @@ class ValidatedTool(_ExportMixin, BaseModel):
     # ------------------------------------------------------------------ #
     # Public entry-point called by the processor
     # ------------------------------------------------------------------ #
-    def execute(self: T_Validated, **kwargs: Any) -> BaseModel:
+    async def execute(self: T_Validated, **kwargs: Any) -> BaseModel:
         """Validate *kwargs*, run `_execute`, validate the result."""
         try:
             args = self.Arguments(**kwargs)  # type: ignore[arg-type]
-            res = self._execute(**args.model_dump())  # type: ignore[arg-type]
+            res = await self._execute(**args.model_dump())  # type: ignore[arg-type]
 
             return (
                 res
@@ -124,16 +124,16 @@ class ValidatedTool(_ExportMixin, BaseModel):
     # ------------------------------------------------------------------ #
     # Sub-classes must implement this
     # ------------------------------------------------------------------ #
-    def _execute(self, **_kwargs: Any):  # noqa: D401 – expected override
-        raise NotImplementedError("Tool must implement _execute()")
+    async def _execute(self, **_kwargs: Any):  # noqa: D401 – expected override
+        raise NotImplementedError("Tool must implement async _execute()")
 
 
 # --------------------------------------------------------------------------- #
-# Decorator to retrofit validation onto classic “imperative” tools
+# Decorator to retrofit validation onto classic "imperative" tools
 # --------------------------------------------------------------------------- #
 def with_validation(cls):  # noqa: D401 – factory
     """
-    Decorator that wraps an existing ``execute`` method with:
+    Decorator that wraps an existing async ``execute`` method with:
 
     * argument validation (based on type hints)
     * result validation (based on return annotation)
@@ -143,13 +143,15 @@ def with_validation(cls):  # noqa: D401 – factory
         validate_result,
     )
 
-    original: Callable[..., Any] = cls.execute  # type: ignore[attr-defined]
+    original = cls.execute  # type: ignore[attr-defined]
+    if not inspect.iscoroutinefunction(original):
+        raise TypeError(f"Tool {cls.__name__} must have an async execute method")
 
-    def _wrapper(self, **kwargs):  # type: ignore[override]
+    async def _async_wrapper(self, **kwargs):  # type: ignore[override]
         tool_name = cls.__name__
         validated = validate_arguments(tool_name, original, kwargs)
-        result = original(self, **validated)
+        result = await original(self, **validated)
         return validate_result(tool_name, original, result)
 
-    cls.execute = _wrapper  # type: ignore[assignment]
+    cls.execute = _async_wrapper  # type: ignore[assignment]
     return cls

chuk_tool_processor/plugins/discovery.py

@@ -1,121 +1,164 @@
 # chuk_tool_processor/plugins/discovery.py
-"""Plugin discovery & registry utilities for chuk_tool_processor"""
+"""Async-friendly plugin discovery & registry utilities for chuk_tool_processor."""
+
 from __future__ import annotations
 
 import importlib
 import inspect
 import logging
 import pkgutil
+from types import ModuleType
 from typing import Any, Dict, List, Optional, Set, Type
 
+from chuk_tool_processor.plugins.parsers.base import ParserPlugin
+from chuk_tool_processor.models.execution_strategy import ExecutionStrategy
+
+__all__ = [
+    "plugin_registry",
+    "PluginRegistry",
+    "PluginDiscovery",
+    "discover_default_plugins",
+    "discover_plugins",
+    "plugin",
+]
+
 logger = logging.getLogger(__name__)
 
 
+# -----------------------------------------------------------------------------
+# In-memory registry
+# -----------------------------------------------------------------------------
 class PluginRegistry:
-    """In‑memory registry keyed by *category → name*."""
+    """Thread-safe (GIL) in-memory registry keyed by *category → name*."""
 
-    def __init__(self) -> None:  # no side‑effects in import time
+    def __init__(self) -> None:
+        # category → {name → object}
         self._plugins: Dict[str, Dict[str, Any]] = {}
 
-    # ------------------------------------------------------------------
+    # --------------------------------------------------------------------- #
+    # Public API
+    # --------------------------------------------------------------------- #
     def register_plugin(self, category: str, name: str, plugin: Any) -> None:
         self._plugins.setdefault(category, {})[name] = plugin
         logger.debug("Registered plugin %s.%s", category, name)
 
-    def get_plugin(self, category: str, name: str) -> Optional[Any]:
+    def get_plugin(self, category: str, name: str) -> Optional[Any]:  # noqa: D401
         return self._plugins.get(category, {}).get(name)
 
     def list_plugins(self, category: str | None = None) -> Dict[str, List[str]]:
-        if category:
-            return {category: list(self._plugins.get(category, {}))}
-        return {cat: list(names) for cat, names in self._plugins.items()}
+        if category is not None:
+            return {category: sorted(self._plugins.get(category, {}))}
+        return {cat: sorted(names) for cat, names in self._plugins.items()}
 
 
+# -----------------------------------------------------------------------------
+# Discovery
+# -----------------------------------------------------------------------------
 class PluginDiscovery:
-    """Recursively scans packages for plugin classes and registers them."""
+    """
+    Recursively scans *package_paths* for plugin classes and registers them.
 
-    def __init__(self, registry: PluginRegistry) -> None:
-        self.registry = registry
-        self._seen: Set[str] = set()
+    * Parser plugins – concrete subclasses of :class:`ParserPlugin`
+      with an **async** ``try_parse`` coroutine.
 
-        # optional parser subsystem
-        try:
-            from chuk_tool_processor.parsers.base import ParserPlugin as _PP  # noqa: WPS433
-        except ModuleNotFoundError:
-            _PP = None
-        self.ParserPlugin = _PP
+    * Execution strategies – concrete subclasses of
+      :class:`ExecutionStrategy`.
 
-        # ExecutionStrategy always present inside core models
-        from chuk_tool_processor.models.execution_strategy import ExecutionStrategy  # noqa: WPS433
+    * Explicitly-decorated plugins – classes tagged with ``@plugin(...)``.
+    """
 
-        self.ExecutionStrategy = ExecutionStrategy
+    # ------------------------------------------------------------------ #
+    def __init__(self, registry: PluginRegistry) -> None:
+        self._registry = registry
+        self._seen_modules: Set[str] = set()
 
-    # ------------------------------------------------------------------
+    # ------------------------------------------------------------------ #
     def discover_plugins(self, package_paths: List[str]) -> None:
-        for pkg in package_paths:
-            self._walk(pkg)
+        """Import every package in *package_paths* and walk its subtree."""
+        for pkg_path in package_paths:
+            self._walk(pkg_path)
 
-    # ------------------------------------------------------------------
+    # ------------------------------------------------------------------ #
+    # Internal helpers
+    # ------------------------------------------------------------------ #
     def _walk(self, pkg_path: str) -> None:
         try:
-            pkg = importlib.import_module(pkg_path)
+            root_pkg = importlib.import_module(pkg_path)
         except ImportError as exc:  # pragma: no cover
             logger.warning("Cannot import package %s: %s", pkg_path, exc)
             return
 
-        for _, mod_name, is_pkg in pkgutil.iter_modules(pkg.__path__, pkg.__name__ + "."):
-            if mod_name in self._seen:
+        self._inspect_module(root_pkg)
+
+        for _, mod_name, is_pkg in pkgutil.iter_modules(root_pkg.__path__, root_pkg.__name__ + "."):
+            if mod_name in self._seen_modules:
+                continue
+            self._seen_modules.add(mod_name)
+
+            try:
+                mod = importlib.import_module(mod_name)
+            except ImportError as exc:  # pragma: no cover
+                logger.debug("Cannot import module %s: %s", mod_name, exc)
                 continue
-            self._seen.add(mod_name)
-            self._inspect_module(mod_name)
+
+            self._inspect_module(mod)
+
             if is_pkg:
                 self._walk(mod_name)
 
-    # ------------------------------------------------------------------
-    def _inspect_module(self, mod_name: str) -> None:
-        try:
-            module = importlib.import_module(mod_name)
-        except ImportError as exc:  # pragma: no cover
-            logger.warning("Cannot import module %s: %s", mod_name, exc)
-            return
-
+    # ------------------------------------------------------------------ #
+    def _inspect_module(self, module: ModuleType) -> None:
         for attr in module.__dict__.values():
             if inspect.isclass(attr):
                 self._maybe_register(attr)
 
-    # ------------------------------------------------------------------
+    # ------------------------------------------------------------------ #
     def _maybe_register(self, cls: Type) -> None:
-        """Register *cls* in all relevant plugin categories."""
-
-        # ---------------- parser plugins ------------------------------
-        looks_like_parser = callable(getattr(cls, "try_parse", None))
-        if looks_like_parser and not inspect.isabstract(cls):
-            # skip ABC base itself if available
-            if self.ParserPlugin and cls is self.ParserPlugin:
-                pass
-            else:
-                self.registry.register_plugin("parser", cls.__name__, cls())
-
-        # --------------- execution strategies -------------------------
-        if (
-            issubclass(cls, self.ExecutionStrategy)
-            and cls is not self.ExecutionStrategy
-            and not inspect.isabstract(cls)
-        ):
-            self.registry.register_plugin("execution_strategy", cls.__name__, cls)
-
-        # --------------- explicit @plugin decorator -------------------
-        meta = getattr(cls, "_plugin_meta", None)
-        if meta and not inspect.isabstract(cls):
-            self.registry.register_plugin(meta.get("category", "unknown"), meta.get("name", cls.__name__), cls())
-
-
-# ----------------------------------------------------------------------
-# public decorator helper
-# ----------------------------------------------------------------------
+        """Register *cls* in all matching plugin categories."""
+        if inspect.isabstract(cls):
+            return
 
+        # ------------------- Parser plugins -------------------------
+        if issubclass(cls, ParserPlugin) and cls is not ParserPlugin:
+            if not inspect.iscoroutinefunction(getattr(cls, "try_parse", None)):
+                logger.warning("Skipping parser plugin %s: try_parse is not async", cls.__qualname__)
+            else:
+                try:
+                    self._registry.register_plugin("parser", cls.__name__, cls())
+                except Exception as exc:  # pragma: no cover
+                    logger.warning("Cannot instantiate parser plugin %s: %s", cls.__qualname__, exc)
+
+        # ---------------- Execution strategies ---------------------
+        if issubclass(cls, ExecutionStrategy) and cls is not ExecutionStrategy:
+            self._registry.register_plugin("execution_strategy", cls.__name__, cls)
+
+        # ------------- Explicit @plugin decorator ------------------
+        meta: Optional[dict] = getattr(cls, "_plugin_meta", None)
+        if meta:
+            category = meta.get("category", "unknown")
+            name = meta.get("name", cls.__name__)
+            try:
+                plugin_obj: Any = cls() if callable(getattr(cls, "__init__", None)) else cls
+                self._registry.register_plugin(category, name, plugin_obj)
+            except Exception as exc:  # pragma: no cover
+                logger.warning("Cannot instantiate decorated plugin %s: %s", cls.__qualname__, exc)
+
+
+# -----------------------------------------------------------------------------
+# Decorator helper
+# -----------------------------------------------------------------------------
 def plugin(category: str, name: str | None = None):
-    """Decorator to mark a class as a plugin for explicit registration."""
+    """
+    Decorator that marks a concrete class as a plugin for *category*.
+
+    Example
+    -------
+    ```python
+    @plugin("transport", name="sse")
+    class MySSETransport:
+        ...
+    ```
+    """
 
     def decorator(cls):
         cls._plugin_meta = {"category": category, "name": name or cls.__name__}
@@ -124,15 +167,17 @@ def plugin(category: str, name: str | None = None):
     return decorator
 
 
-# ----------------------------------------------------------------------
+# -----------------------------------------------------------------------------
 # Singletons & convenience wrappers
-# ----------------------------------------------------------------------
+# -----------------------------------------------------------------------------
 plugin_registry = PluginRegistry()
 
 
 def discover_default_plugins() -> None:
+    """Discover plugins shipped inside *chuk_tool_processor.plugins*."""
     PluginDiscovery(plugin_registry).discover_plugins(["chuk_tool_processor.plugins"])
 
 
 def discover_plugins(package_paths: List[str]) -> None:
+    """Discover plugins from arbitrary external *package_paths*."""
     PluginDiscovery(plugin_registry).discover_plugins(package_paths)