dreadnode 1.16.0__tar.gz → 1.17.0__tar.gz

{dreadnode-1.16.0 → dreadnode-1.17.0}/.gitignore

@@ -156,6 +156,7 @@ venv.bak/
 
 # mkdocs documentation
 /site
+debug.html
 
 # mypy
 .mypy_cache/

{dreadnode-1.16.0 → dreadnode-1.17.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dreadnode
-Version: 1.16.0
+Version: 1.17.0
 Summary: Dreadnode SDK
 Project-URL: Homepage, https://github.com/dreadnode/sdk
 Project-URL: Repository, https://github.com/dreadnode/sdk
@@ -208,14 +208,22 @@ License:                                  Apache License
         See the License for the specific language governing permissions and
         limitations under the License.
 License-File: LICENSE
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Requires-Python: <3.14,>=3.10
+Requires-Dist: aiofiles<25.0.0,>=24.1.0
 Requires-Dist: coolname<3.0.0,>=2.2.0
 Requires-Dist: cyclopts>=4.2.0
-Requires-Dist: fsspec[s3]<=2025.3.0,>=2023.1.0
+Requires-Dist: fsspec[s3]<=2025.12.0,>=2023.1.0
 Requires-Dist: httpx<1.0.0,>=0.28.0
 Requires-Dist: logfire<=3.20.0,>=3.5.3
 Requires-Dist: loguru>=0.7.3
-Requires-Dist: numpy<=2.2.6
+Requires-Dist: numpy<=2.3.5
 Requires-Dist: optuna<5.0.0,>=4.5.0
 Requires-Dist: pandas<3.0.0,>=2.2.3
 Requires-Dist: pydantic<3.0.0,>=2.9.2
@@ -278,7 +286,7 @@ Dreadnode Strikes SDK
     <img alt="PyPI - Python Version" src="https://img.shields.io/pypi/pyversions/dreadnode">
     <img alt="PyPI - Version" src="https://img.shields.io/pypi/v/dreadnode">
     <img alt="GitHub License" src="https://img.shields.io/github/license/dreadnode/sdk">
-    <img alt="Tests" src="https://img.shields.io/github/actions/workflow/status/dreadnode/sdk/tests.yaml">
+    <img alt="Tests" src="https://img.shields.io/github/actions/workflow/status/dreadnode/sdk/test.yaml">
     <img alt="Pre-Commit" src="https://img.shields.io/github/actions/workflow/status/dreadnode/sdk/pre-commit.yaml">
     <img alt="Renovate" src="https://img.shields.io/github/actions/workflow/status/dreadnode/sdk/renovate.yaml">
 </h4>

{dreadnode-1.16.0 → dreadnode-1.17.0}/README.md

@@ -16,7 +16,7 @@ Dreadnode Strikes SDK
     <img alt="PyPI - Python Version" src="https://img.shields.io/pypi/pyversions/dreadnode">
     <img alt="PyPI - Version" src="https://img.shields.io/pypi/v/dreadnode">
     <img alt="GitHub License" src="https://img.shields.io/github/license/dreadnode/sdk">
-    <img alt="Tests" src="https://img.shields.io/github/actions/workflow/status/dreadnode/sdk/tests.yaml">
+    <img alt="Tests" src="https://img.shields.io/github/actions/workflow/status/dreadnode/sdk/test.yaml">
     <img alt="Pre-Commit" src="https://img.shields.io/github/actions/workflow/status/dreadnode/sdk/pre-commit.yaml">
     <img alt="Renovate" src="https://img.shields.io/github/actions/workflow/status/dreadnode/sdk/renovate.yaml">
 </h4>

{dreadnode-1.16.0 → dreadnode-1.17.0}/dreadnode/agent/agent.py

@@ -2,7 +2,7 @@ import inspect
 import json
 import re
 import typing as t
-from contextlib import aclosing, asynccontextmanager
+from contextlib import AsyncExitStack, aclosing, asynccontextmanager
 from copy import deepcopy
 from textwrap import dedent
 
@@ -875,8 +875,16 @@ class Agent(Model):
         commit: CommitBehavior = "always",
     ) -> t.AsyncIterator[t.AsyncGenerator[AgentEvent, None]]:
         thread = thread or self.thread
-        async with aclosing(self._stream_traced(thread, user_input, commit=commit)) as stream:
-            yield stream
+
+        async with AsyncExitStack() as stack:
+            # Ensure all tools are properly entered if they
+            # are context managers before we start using them
+            for tool_container in self.tools:
+                if hasattr(tool_container, "__aenter__") and hasattr(tool_container, "__aexit__"):
+                    await stack.enter_async_context(tool_container)
+
+            async with aclosing(self._stream_traced(thread, user_input, commit=commit)) as stream:
+                yield stream
 
     async def run(
         self,

{dreadnode-1.16.0 → dreadnode-1.17.0}/dreadnode/agent/tools/base.py

@@ -1,6 +1,8 @@
+import asyncio
+import functools
 import typing as t
 
-from pydantic import ConfigDict
+from pydantic import ConfigDict, PrivateAttr
 from rigging import tools
 from rigging.tools.base import ToolMethod as RiggingToolMethod
 
@@ -171,18 +173,82 @@ class Toolset(Model):
     - Pydantic's declarative syntax for defining state (fields).
     - Automatic application of the `@configurable` decorator.
     - A `get_tools` method for discovering methods decorated with `@dreadnode.tool_method`.
+    - Support for async context management, with automatic re-entrancy handling.
     """
 
+    model_config = ConfigDict(arbitrary_types_allowed=True, use_attribute_docstrings=True)
+
     variant: str | None = None
     """The variant for filtering tools available in this toolset."""
 
-    model_config = ConfigDict(arbitrary_types_allowed=True, use_attribute_docstrings=True)
+    # Context manager magic
+    _entry_ref_count: int = PrivateAttr(default=0)
+    _context_handle: object = PrivateAttr(default=None)
+    _entry_lock: asyncio.Lock = PrivateAttr(default_factory=asyncio.Lock)
 
     @property
     def name(self) -> str:
         """The name of the toolset, derived from the class name."""
         return self.__class__.__name__
 
+    def __init_subclass__(cls, **kwargs: t.Any) -> None:
+        super().__init_subclass__(**kwargs)
+
+        # This essentially ensures that if the Toolset is any kind of context manager,
+        # it will be re-entrant, and only actually enter/exit once. This means we can
+        # safely build auto-entry/exit logic into our Agent class without worrying about
+        # breaking the code if the user happens to enter a toolset manually before using
+        # it in an agent.
+
+        original_aenter = cls.__dict__.get("__aenter__")
+        original_enter = cls.__dict__.get("__enter__")
+        original_aexit = cls.__dict__.get("__aexit__")
+        original_exit = cls.__dict__.get("__exit__")
+
+        has_enter = callable(original_aenter) or callable(original_enter)
+        has_exit = callable(original_aexit) or callable(original_exit)
+
+        if has_enter and not has_exit:
+            raise TypeError(
+                f"{cls.__name__} defining __aenter__ or __enter__ must also define __aexit__ or __exit__"
+            )
+        if has_exit and not has_enter:
+            raise TypeError(
+                f"{cls.__name__} defining __aexit__ or __exit__ must also define __aenter__ or __enter__"
+            )
+        if original_aenter and original_enter:
+            raise TypeError(f"{cls.__name__} cannot define both __aenter__ and __enter__")
+        if original_aexit and original_exit:
+            raise TypeError(f"{cls.__name__} cannot define both __aexit__ and __exit__")
+
+        @functools.wraps(original_aenter or original_enter)  # type: ignore[arg-type]
+        async def aenter_wrapper(self: "Toolset", *args: t.Any, **kwargs: t.Any) -> t.Any:
+            async with self._entry_lock:
+                if self._entry_ref_count == 0:
+                    handle = None
+                    if original_aenter:
+                        handle = await original_aenter(self, *args, **kwargs)
+                    elif original_enter:
+                        handle = original_enter(self, *args, **kwargs)
+                    self._context_handle = handle if handle is not None else self
+                self._entry_ref_count += 1
+                return self._context_handle
+
+        cls.__aenter__ = aenter_wrapper  # type: ignore[attr-defined]
+
+        @functools.wraps(original_aexit or original_exit)  # type: ignore[arg-type]
+        async def aexit_wrapper(self: "Toolset", *args: t.Any, **kwargs: t.Any) -> t.Any:
+            async with self._entry_lock:
+                self._entry_ref_count -= 1
+                if self._entry_ref_count == 0:
+                    if original_aexit:
+                        await original_aexit(self, *args, **kwargs)
+                    elif original_exit:
+                        original_exit(self, *args, **kwargs)
+                    self._context_handle = None
+
+        cls.__aexit__ = aexit_wrapper  # type: ignore[attr-defined]
+
     def get_tools(self, *, variant: str | None = None) -> list[AnyTool]:
         variant = variant or self.variant