diff --git a/python/deepagents/README.md b/python/deepagents/README.md index 04ba90ef..cf8dcd4e 100644 --- a/python/deepagents/README.md +++ b/python/deepagents/README.md @@ -372,6 +372,31 @@ worker is unaffected, and the original function is restored when the worker stops. If you would rather be explicit, use `create_temporal_deep_agent` or pass `TemporalModel("provider:name")` yourself. +Each plugin instance registers four Activities: `deepagents.invoke_model`, +`deepagents.invoke_model_streaming`, `deepagents.invoke_tool` and +`deepagents.backend_op`. Use one `DeepAgentsPlugin` per Worker: one +`model_provider` receives every model name, so it can serve several models, and +`backends` takes every factory. The same instance on both the Client and the +Worker works. So do one instance on the Client and one on the Worker built with +the same `model_provider` and `backends` factories (the same objects) and the +same `streaming_batch_interval`; two default instances match. Any other second +instance fails at Worker start with a `ValueError` that names +`DeepAgentsPlugin`. With one instance on each side, the Worker instance does +not add an Activity whose name is already +registered, and still registered after its own step, by the same function, the +Client instance's copy of the same built-in Activity, or a wrapper of either +that sets `__wrapped__` (as `functools.wraps` does). +A dynamic Activity counts the same way, with the dynamic slot as its name. +An Activity registered under another name, such as a wrapper with an alias, +does not count. So the Worker keeps what the Client side +registered, including a wrapper a later Client plugin put around it if that +wrapper sets `__wrapped__`, and any +other Activity given to the Worker instance is still added. Any other repeated +Activity name fails at Worker start with `ValueError: More than one activity +named ...`, as it does without the plugin. That includes a caller Activity +with one of these names. A second, different +dynamic Activity fails with `TypeError: More than one dynamic activity`. + ## Composing with other plugins This plugin carries no tracing context of its own. For observability, compose it diff --git a/python/deepagents/src/temporalio/deepagents/_plugin.py b/python/deepagents/src/temporalio/deepagents/_plugin.py index fde00a00..dd4ab7b6 100644 --- a/python/deepagents/src/temporalio/deepagents/_plugin.py +++ b/python/deepagents/src/temporalio/deepagents/_plugin.py @@ -17,13 +17,23 @@ * pushes the model / tool dispatch defaults down to the seams that read them. The plugin auto-propagates from a ``Client`` to any ``Worker`` built from it, so -add it on exactly one side. ``configure_worker`` additionally de-duplicates -activities by name, so a user who mistakenly adds it on both sides gets a clean -no-op instead of a "More than one activity named ..." crash. +add it on exactly one side. One instance on the Client and one on the Worker +are accepted only when they are the same object or were built from the same +``model_provider``, ``backends`` factories and ``streaming_batch_interval``; any +other second instance raises ``ValueError``. In the accepted case, the Worker +instance does not add an activity whose name (or, for a dynamic activity, the +dynamic slot) is already registered, and still registered after its own step, +by the same callable, a wrapper of it (``__wrapped__``), or, for a second +instance, the Client instance's copy of the same built-in activity. The Worker +keeps what the Client side registered, instead of failing with "More than one +activity named ...". Any other repeated activity name, such as a caller +activity with a plugin activity's name, and any second dynamic activity, is +left to the Worker, which refuses it. """ from __future__ import annotations +import inspect import sys import warnings from collections.abc import AsyncGenerator, Mapping, Sequence @@ -126,6 +136,12 @@ def __init__( # (_resolve_tool_options), renamed to ``options`` at the callsite. _tools.set_tool_defaults(self._tool_activity_options) + # What the activities are built from; a second instance on the Worker + # must match it (see _differences). + self._model_provider = model_provider + self._backends = dict(backends or {}) + self._streaming_batch_interval = streaming_batch_interval + # The activities that carry the nondeterministic work. model_provider and # the batch interval are captured here so API keys never enter an input. self._activities = DeepAgentActivities( @@ -133,16 +149,17 @@ def __init__( backends=backends, streaming_batch_interval=streaming_batch_interval, ) - activities = [ + # Kept apart from the public ``activities`` setting, which callers may change. + self._builtin_activities = ( self._activities.invoke_model, self._activities.invoke_model_streaming, self._activities.invoke_tool, self._activities.backend_op, - ] + ) super().__init__( "langchain.DeepAgentsPlugin", - activities=activities, + activities=list(self._builtin_activities), data_converter=self._configure_data_converter, workflow_runner=self._make_workflow_runner(), workflow_failure_exception_types=[DeepAgentsWorkflowError], @@ -238,13 +255,161 @@ async def _run_context(self) -> AsyncGenerator[None, None]: # -- worker config ------------------------------------------------------- def configure_worker(self, config: Any) -> Any: - """Deduplicate activity registrations after the base configuration.""" + """Refuse a second instance; skip activities the Client's one registered.""" + client_instance = self._client_instance(config) + before = list(config.get("activities") or []) + if config.get("activities") is not None: + # A new list: the caller may pass a one-shot iterable, and a callable + # ``activities`` setting may append to the list it gets. + config["activities"] = list(before) config = super().configure_worker(config) - activities = config.get("activities") - if activities: - config["activities"] = _dedupe_activities(activities) + if client_instance is not None: + # A list: an ``activities`` callable may return a one-shot iterable. + config["activities"] = self._without_repeats( + before, list(config.get("activities") or []), client_instance + ) return config + def _client_instance(self, config: Any) -> DeepAgentsPlugin | None: + """The Client's instance, if this call is the Worker instance's. + + One instance on each side is accepted only when both are the same + object, which the Worker then calls twice, or were built from the same + sources. Any other second instance raises ``ValueError``. + """ + client = config.get("client") + client_plugins = client.config().get("plugins", []) if client else [] + on_client = [p for p in client_plugins if isinstance(p, DeepAgentsPlugin)] + on_worker = [ + p for p in config.get("plugins") or [] if isinstance(p, DeepAgentsPlugin) + ] + differences = ( + on_client[0]._differences(on_worker[0]) if on_client and on_worker else [] + ) + if len(on_client) > 1 or len(on_worker) > 1: + reason = "" + elif differences: + reason = f" with a different {' and '.join(differences)}" + elif on_client and on_worker and on_worker[0] is self: + return on_client[0] + else: + return None + raise ValueError( + f"DeepAgentsPlugin is registered {len(on_client) + len(on_worker)} " + f"times for this Worker ({len(on_client)} on the Client, " + f"{len(on_worker)} on the Worker){reason}. Use one DeepAgentsPlugin " + "instance for this Worker: one model_provider receives every model " + "name, so it can serve several models, and backends= takes every " + "factory. To run another instance, build its Worker from a Client " + "without DeepAgentsPlugin, on its own task queue; in one process, " + "model_activity_options, tool_activity_options and streaming_topic " + "still come from the instance built last." + ) + + def _differences(self, other: DeepAgentsPlugin) -> list[str]: + """The inputs of ``other``'s activities that differ from this instance's. + + ``model_provider`` and each backend factory must be the same object; + ``streaming_batch_interval`` must be equal. + """ + same_backends = self._backends.keys() == other._backends.keys() and all( + _same_callable(factory, other._backends[name]) + for name, factory in self._backends.items() + ) + same = { + "model_provider": _same_callable( + self._model_provider, other._model_provider + ), + "backends": same_backends, + "streaming_batch_interval": ( + self._streaming_batch_interval == other._streaming_batch_interval + ), + } + return [name for name, is_same in same.items() if not is_same] + + def _without_repeats( + self, + before: Sequence[Any], + after: Sequence[Any], + client_instance: DeepAgentsPlugin, + ) -> list[Any]: + """Drop the activities this call added that repeat a kept registration. + + An activity registered before this call that is still in ``after`` + keeps its activity name, or the dynamic slot if it is dynamic. An added + activity with that name or slot repeats it when the kept callable is it + or wraps it, or, for one of this instance's built-in activities, is or + wraps the Client instance's copy of it. Dropping such a repeat never + changes which code runs under the name; any other repeat reaches the + Worker's duplicate-name or duplicate-dynamic check. + """ + remaining = list(before) + added: list[bool] = [] + kept_keys: dict[tuple[str | None], list[Any]] = {} + for act in after: + index = _index_of(act, remaining) + added.append(index is None) + if index is not None: + del remaining[index] + key = _activity_key(act) + if key is not None: + kept_keys.setdefault(key, []).extend(_wrapped_chain(act)) + counterparts = list( + zip(self._builtin_activities, client_instance._builtin_activities) + ) + kept: list[Any] = [] + for act, is_added in zip(after, added): + key = _activity_key(act) if is_added else None + chain = kept_keys.get(key, []) if key is not None else [] + if _index_of(act, chain) is not None or any( + _same_callable(act, own) and _index_of(client_copy, chain) is not None + for own, client_copy in counterparts + ): + continue + kept.append(act) + return kept + + +def _same_callable(a: Any, b: Any) -> bool: + """Whether ``a`` and ``b`` are one callable, without calling ``__eq__``. + + Two bound methods of one function on one object count as one callable; + any other pair must be the same object. Callables that only compare equal + are different registrations. + """ + if a is b: + return True + return ( + inspect.ismethod(a) + and inspect.ismethod(b) + and a.__self__ is b.__self__ + and a.__func__ is b.__func__ + ) + + +def _index_of(act: Any, acts: Sequence[Any]) -> int | None: + """The index of the first entry of ``acts`` that is ``act``, or ``None``.""" + return next((i for i, other in enumerate(acts) if _same_callable(act, other)), None) + + +def _wrapped_chain(act: Any) -> list[Any]: + """``act`` and the callables it wraps, following ``__wrapped__``.""" + chain: list[Any] = [] + while act is not None and len(chain) < 100: + chain.append(act) + act = getattr(act, "__wrapped__", None) + return chain + + +def _activity_key(act: Any) -> tuple[str | None] | None: + """``(name,)`` for the activity ``act`` registers, or ``None`` if it has none. + + A dynamic activity has no name, so its key is ``(None,)``: the one dynamic + slot a Worker has. + """ + defn = activity_mod._Definition.from_callable(act) + return (defn.name,) if defn is not None else None + async def _temporal_aio_to_thread( default_aio_to_thread: Callable[..., Any], @@ -294,23 +459,3 @@ def _install_langsmith_temporal_override() -> None: langsmith.set_runtime_overrides(aio_to_thread=_temporal_aio_to_thread) except Exception: pass - - -def _dedupe_activities(activities: Sequence[Any]) -> list[Any]: - """Drop duplicate activity registrations by defn name, keeping the first. - - Guards the "plugin added on both Client and Worker" trap, where the plugin's - activities would otherwise be appended twice and the worker would reject the - duplicate names. - """ - seen: set[str] = set() - out: list[Any] = [] - for act in activities: - defn = activity_mod._Definition.from_callable(act) - name = defn.name if defn is not None else getattr(act, "__name__", repr(act)) - key = name or repr(act) - if key in seen: - continue - seen.add(key) - out.append(act) - return out diff --git a/python/deepagents/tests/test_activity_registration.py b/python/deepagents/tests/test_activity_registration.py new file mode 100644 index 00000000..156f070a --- /dev/null +++ b/python/deepagents/tests/test_activity_registration.py @@ -0,0 +1,366 @@ +"""The plugin leaves repeated activity names to the Worker, which refuses them. + +The one exception is one instance on the Client and one on the Worker, the same +object or built from the same sources: the Worker keeps what the Client side +registered, and the Worker instance does not add it again. Any other second +instance is refused by the plugin. +""" + +from __future__ import annotations + +import dataclasses +import functools +import itertools +import re +import uuid +from collections.abc import Callable, Sequence +from datetime import timedelta +from typing import Any, cast + +import pytest + +from temporalio import activity +from temporalio.client import Client +from temporalio.client import Plugin as ClientPlugin +from temporalio.common import RawValue +from temporalio.deepagents import DeepAgentsPlugin +from temporalio.plugin import SimplePlugin +from temporalio.worker import Worker, WorkerConfig + + +@activity.defn +async def fetch(key: str) -> str: + return key + + +@activity.defn(name="fetch") +async def fetch_v1(key: str) -> str: + return key + + +@activity.defn(name="deepagents.invoke_tool") +async def caller_invoke_tool(key: str) -> str: + return key + + +@activity.defn(dynamic=True) +async def fallback(_args: Sequence[RawValue]) -> str: + return activity.info().activity_type + + +@activity.defn(dynamic=True) +async def other_fallback(_args: Sequence[RawValue]) -> str: + return activity.info().activity_type + + +@dataclasses.dataclass +class _Endpoint: + """An async callable; any two instances compare equal.""" + + endpoint: str = dataclasses.field(compare=False) + + async def __call__(self, key: str) -> str: + return f"{self.endpoint}:{key}" + + +@dataclasses.dataclass +class _DynamicEndpoint: + """An async callable for a dynamic activity; any two instances compare equal.""" + + endpoint: str = dataclasses.field(compare=False) + + async def __call__(self, _args: Sequence[RawValue]) -> str: + return f"{self.endpoint}:{activity.info().activity_type}" + + +def _endpoint(endpoint: str) -> Callable[..., object]: + return activity.defn(name="endpoint")(_Endpoint(endpoint)) + + +def _dynamic_endpoint(endpoint: str) -> Callable[..., object]: + return activity.defn(dynamic=True)(_DynamicEndpoint(endpoint)) + + +def _client_with(client: Client, plugins: Sequence[ClientPlugin]) -> Client: + config = client.config() + config["plugins"] = list(plugins) + return Client(**config) + + +def _own(plugin: DeepAgentsPlugin) -> list[Callable[..., object]]: + return list(cast(Sequence[Callable[..., object]], plugin.activities)) + + +def _provider(model_name: str) -> Any: + raise AssertionError(f"model_provider called for {model_name}") + + +def _backend() -> Any: + raise AssertionError("backend factory called") + + +@pytest.mark.parametrize( + ("client_options", "worker_options", "message"), + [ + ([{}, {}], [], "2 times for this Worker (2 on the Client, 0 on the Worker)."), + ([], [{}, {}], "2 times for this Worker (0 on the Client, 2 on the Worker)."), + ([{}], [{}, {}], "3 times for this Worker (1 on the Client, 2 on the Worker)."), + ( + [{}], + [{"model_provider": _provider}], + "2 times for this Worker (1 on the Client, 1 on the Worker) with a different model_provider.", + ), + ( + [{"backends": {"fs": _backend}}], + [{"backends": {"fs": lambda: None}}], + "2 times for this Worker (1 on the Client, 1 on the Worker) with a different backends.", + ), + ( + [{}], + [{"streaming_batch_interval": timedelta(seconds=1)}], + "2 times for this Worker (1 on the Client, 1 on the Worker) with a different streaming_batch_interval.", + ), + ], + ids=[ + "two-on-client", + "two-on-worker", + "one-on-client-two-on-worker", + "different-model-provider", + "different-backend-factory", + "different-batch-interval", + ], +) +async def test_second_instance_is_refused( + client: Client, + client_options: list[dict[str, Any]], + worker_options: list[dict[str, Any]], + message: str, +) -> None: + with pytest.raises( + ValueError, match=re.escape(f"DeepAgentsPlugin is registered {message}") + ): + Worker( + _client_with(client, [DeepAgentsPlugin(**o) for o in client_options]), + task_queue=f"activity-registration-{uuid.uuid4()}", + activities=[fetch], + plugins=[DeepAgentsPlugin(**o) for o in worker_options], + ) + + +@pytest.mark.parametrize( + ("client_plugins", "worker_plugins", "activities", "worker_extra"), + [ + (1, 0, [fetch_v1, fetch], None), + (1, 0, [caller_invoke_tool], None), + (1, 1, [fetch], fetch_v1), + (1, 1, [_endpoint("client")], _endpoint("worker")), + ], + ids=[ + "caller-activities-share-a-name", + "caller-activity-uses-plugin-name", + "worker-instance-activity-repeats-caller-name", + "worker-instance-activity-equals-caller-activity", + ], +) +async def test_repeated_activity_names_are_refused( + client: Client, + client_plugins: int, + worker_plugins: int, + activities: list[Callable[..., object]], + worker_extra: Callable[..., object] | None, +) -> None: + on_worker = [DeepAgentsPlugin() for _ in range(worker_plugins)] + if worker_extra is not None: + # Another activity with a caller activity's name, given to the Worker + # instance; the _Endpoint pair is distinct but compares equal. + on_worker[0].activities = [*_own(on_worker[0]), worker_extra] + with pytest.raises(ValueError, match="More than one activity named"): + Worker( + _client_with(client, [DeepAgentsPlugin() for _ in range(client_plugins)]), + task_queue=f"activity-registration-{uuid.uuid4()}", + activities=activities, + plugins=on_worker, + ) + + +@pytest.mark.parametrize("handlers", ["functions", "equal-objects"]) +@pytest.mark.parametrize("same_instance", [True, False]) +async def test_two_dynamic_activities_are_refused( + client: Client, same_instance: bool, handlers: str +) -> None: + on_client = DeepAgentsPlugin() + on_worker = on_client if same_instance else DeepAgentsPlugin() + caller: Callable[..., object] = fallback + extra: Callable[..., object] = other_fallback + if handlers == "equal-objects": + caller, extra = _dynamic_endpoint("client"), _dynamic_endpoint("worker") + on_worker.activities = [*_own(on_worker), extra] + with pytest.raises(TypeError, match="More than one dynamic activity"): + Worker( + _client_with(client, [on_client]), + task_queue=f"activity-registration-{uuid.uuid4()}", + activities=[caller], + plugins=[on_worker], + ) + + +def _name(act: Callable[..., object]) -> str: + return activity._Definition.must_from_callable(act).name or "" # pyright: ignore[reportPrivateUsage] + + +def _wrap(act: Callable[..., Any], name: str | None = None) -> Callable[..., Any]: + # updated=(): act.__dict__ holds its activity definition. + @functools.wraps(act, updated=()) + async def wrapper(*args: Any, **kwargs: Any) -> Any: + return await act(*args, **kwargs) + + return activity.defn(name=name or _name(act))(wrapper) + + +def _wrap_all( + acts: Sequence[Callable[..., Any]] | None, +) -> Sequence[Callable[..., Any]]: + return [_wrap(act) for act in acts or []] + + +def _alias_all( + acts: Sequence[Callable[..., Any]] | None, +) -> Sequence[Callable[..., Any]]: + # Replaces each plugin activity with a wrapper named "alias.". + out: list[Callable[..., Any]] = [] + for act in acts or []: + if _name(act).startswith("deepagents."): + act = _wrap(act, f"alias.{_name(act)}") + out.append(act) + return out + + +_BUILTIN_NAMES = [ + "deepagents.backend_op", + "deepagents.invoke_model", + "deepagents.invoke_model_streaming", + "deepagents.invoke_tool", +] + + +class _CopyConfig(SimplePlugin): + """A Client plugin that hands the next plugin a copy of the config.""" + + def __init__(self) -> None: + super().__init__("copy-config") + + def configure_worker(self, config: WorkerConfig) -> WorkerConfig: + return WorkerConfig(**super().configure_worker(config)) + + +@pytest.mark.parametrize( + "setup", + [ + "same-instance", + "two-instances", + "shared-options", + "same-instance-generator", + "client-wrapper", + "same-instance-client-wrapper", + "worker-extra", + "same-instance-extra", + "same-instance-callable-extra", + "same-instance-dynamic-extra", + "same-instance-callable-dynamic-extra", + "same-instance-callable-iterator", + "same-instance-copying-client-plugin", + "same-instance-caller-copy", + "client-alias", + "same-instance-client-alias", + "worker-callable-replaces", + ], +) +async def test_one_instance_on_each_side_keeps_the_client_registration( + client: Client, setup: str +) -> None: + options: dict[str, Any] = {} + if setup == "shared-options": + options = { + "model_provider": _provider, + "backends": {"fs": _backend}, + "streaming_batch_interval": timedelta(seconds=1), + } + on_client = DeepAgentsPlugin(**options) + on_worker = ( + on_client if setup.startswith("same-instance") else DeepAgentsPlugin(**options) + ) + client_own = _own(on_client) + worker_own = _own(on_worker) + client_plugins: list[ClientPlugin] = [on_client] + activities: list[Callable[..., object]] = [fetch] + if setup.endswith("client-wrapper"): + client_plugins.append(SimplePlugin("wrap", activities=_wrap_all)) + if setup.endswith("client-alias"): + client_plugins.append(SimplePlugin("alias", activities=_alias_all)) + if setup.endswith("callable-replaces"): + + def replace_builtins( + existing: Sequence[Callable[..., Any]] | None, + ) -> Sequence[Callable[..., Any]]: + kept = [a for a in existing or [] if not _name(a).startswith("deepagents.")] + return [*kept, *worker_own] + + on_worker.activities = replace_builtins + if setup.endswith("callable-iterator"): + + def chain_own(existing: Any) -> Any: + # A one-shot iterable, outside the Sequence type the setting declares. + return itertools.chain(existing or [], worker_own) + + on_worker.activities = chain_own + if setup.endswith("copying-client-plugin"): + client_plugins.append(_CopyConfig()) + if setup.endswith("caller-copy"): + activities = [fetch, client_own[0]] + if setup.endswith("generator"): + # A one-shot iterable, outside the Sequence type the Worker declares. + one_shot: Any = (a for a in [fetch]) + activities = one_shot + if setup.endswith("-extra"): + extra = fallback if setup.endswith("dynamic-extra") else fetch + given: list[Callable[..., Any]] = [*_own(on_worker), extra] + if setup.endswith("callable-extra"): + + def add( + existing: Sequence[Callable[..., Any]] | None, + ) -> Sequence[Callable[..., Any]]: + return [*(existing or []), *given] + + on_worker.activities = add + else: + on_worker.activities = given + activities = [fetch] if extra is fallback else [] + + worker = Worker( + _client_with(client, client_plugins), + task_queue=f"activity-registration-{uuid.uuid4()}", + activities=activities, + plugins=[on_worker], + ) + + registered = list(worker.config(active_config=True).get("activities") or []) + aliases = [f"alias.{name}" for name in _BUILTIN_NAMES] + expected = [*_BUILTIN_NAMES, "fetch"] + if setup.endswith("dynamic-extra"): + # The dynamic activity has no name; it is registered once. + expected.append("") + if setup.endswith("client-alias"): + # The aliases do not stand in for the plugin's own names; both stay. + expected = [*aliases, *expected] + assert sorted(_name(act) for act in registered) == sorted(expected) + for act in registered: + if _name(act) in ("fetch", ""): + continue + if _name(act) in aliases: + assert getattr(act, "__wrapped__", None) in client_own + elif setup.endswith(("client-alias", "callable-replaces")): + assert act in worker_own + elif setup.endswith("client-wrapper"): + assert getattr(act, "__wrapped__", None) in client_own + else: + assert act in client_own