diff --git a/temporalio/contrib/_langchain/__init__.py b/temporalio/contrib/_langchain/__init__.py new file mode 100644 index 000000000..0b9247c45 --- /dev/null +++ b/temporalio/contrib/_langchain/__init__.py @@ -0,0 +1,10 @@ +"""Internal shared machinery for the LangChain-family plugins. + +This package is shared infrastructure for ``temporalio.contrib.langgraph``, +``temporalio.contrib.langsmith``, and ``temporalio.contrib.deepagents``. It is +NOT a public API: names, modules, and behavior may change without notice. + +Import discipline: this ``__init__`` performs no imports, and submodules defer +every third-party import into function bodies, so the package is importable +with none of the LangChain-family distributions installed. +""" diff --git a/temporalio/contrib/_langchain/_activity_helpers.py b/temporalio/contrib/_langchain/_activity_helpers.py new file mode 100644 index 000000000..15551f819 --- /dev/null +++ b/temporalio/contrib/_langchain/_activity_helpers.py @@ -0,0 +1,90 @@ +"""Activity-side helpers shared by the LangChain-family plugins.""" + +from __future__ import annotations + +import asyncio +from datetime import timedelta +from functools import wraps +from typing import Any, Callable + +from temporalio import activity +from temporalio.exceptions import ApplicationError + + +def auto_heartbeater(fn: Callable) -> Callable: + """Heartbeat at half the configured ``heartbeat_timeout`` while ``fn`` runs. + + Long LLM calls (thinking mode, long context, streaming accumulation) can run + well past a scheduler's patience; without a heartbeat Temporal would cancel + them and surface a ``HeartbeatTimeoutError`` instead of the real problem. + """ + + @wraps(fn) + async def wrapped(*args: Any, **kwargs: Any) -> Any: + heartbeat_timeout = activity.info().heartbeat_timeout + beat_task: asyncio.Task | None = None + if heartbeat_timeout: + interval = heartbeat_timeout.total_seconds() / 2 + + async def beat() -> None: + while True: + activity.heartbeat() + await asyncio.sleep(interval) + + beat_task = asyncio.create_task(beat()) + try: + return await fn(*args, **kwargs) + finally: + if beat_task is not None: + beat_task.cancel() + # Let the cancellation land before returning so no pending task + # outlives the activity (a bare ``cancel()`` leaves the task to + # be destroyed while pending if the loop shuts down first). + # ``asyncio.wait`` never re-raises the task's CancelledError. + await asyncio.wait([beat_task]) + + return wrapped + + +def translate_api_error(exc: Exception) -> ApplicationError | None: + """Map an LLM SDK HTTP error onto Temporal's retry contract. + + Works by duck typing so neither ``openai`` nor ``anthropic`` needs to be + imported here: both expose ``status_code`` and ``response.headers``. Returns + ``None`` when ``exc`` is not a recognizable HTTP status error, so the caller + can fall through to its generic handling. + """ + status = getattr(exc, "status_code", None) + if status is None: + return None + headers: dict[str, Any] = {} + response = getattr(exc, "response", None) + if response is not None: + headers = dict(getattr(response, "headers", {}) or {}) + # Case-insensitive header access. + lower = {str(k).lower(): v for k, v in headers.items()} + + retryable = status in (408, 409, 429) or 500 <= status < 600 + should_retry = lower.get("x-should-retry") + if should_retry == "false": + retryable = False + elif should_retry == "true": + retryable = True + + delay_ms = lower.get("retry-after-ms") + retry_after = lower.get("retry-after") + next_delay: timedelta | None = None + try: + if delay_ms is not None: + next_delay = timedelta(milliseconds=int(delay_ms)) + elif retry_after is not None: + next_delay = timedelta(seconds=int(retry_after)) + except (TypeError, ValueError): + next_delay = None + + return ApplicationError( + str(exc), + type=type(exc).__name__, + non_retryable=not retryable, + next_retry_delay=next_delay, + ) diff --git a/temporalio/contrib/_langchain/_aio_to_thread.py b/temporalio/contrib/_langchain/_aio_to_thread.py new file mode 100644 index 000000000..fc55c53dd --- /dev/null +++ b/temporalio/contrib/_langchain/_aio_to_thread.py @@ -0,0 +1,59 @@ +"""LangSmith ``aio_to_thread`` override shared by the LangChain-family plugins.""" + +from __future__ import annotations + +from typing import Any, Callable + +import temporalio.workflow + +_installed = False + + +async def _temporal_aio_to_thread( + default_aio_to_thread: Callable[..., Any], + ctx: Any, + func: Callable[..., Any], + /, + *args: Any, + **kwargs: Any, +) -> Any: + """Run LangSmith's ``aio_to_thread`` synchronously inside Temporal workflows. + + The ``@traceable`` decorator on async functions uses ``aio_to_thread()`` → + ``loop.run_in_executor()`` for run setup/teardown. The Temporal workflow + event loop does not support ``run_in_executor``. This override runs those + functions synchronously on the workflow thread when inside a workflow, + and delegates to the default implementation outside workflows. + + Registered via ``langsmith.set_runtime_overrides(aio_to_thread=...)``. + """ + if not temporalio.workflow.in_workflow(): + return await default_aio_to_thread(ctx, func, *args, **kwargs) + with temporalio.workflow.unsafe.sandbox_unrestricted(): + return ctx.run(func, *args, **kwargs) + + +def install_aio_to_thread_override() -> None: + """Install the ``aio_to_thread`` override via LangSmith's official API. + + Safe to call multiple times and from multiple plugins; the override is + installed once per process. It is deliberately never uninstalled: + LangSmith exposes a single process-wide override slot (each + ``set_runtime_overrides`` call replaces it wholesale), so resetting it on + one worker's shutdown would strip a composed plugin's still-needed + override. Leaving it installed is safe — the override defers to + LangSmith's default thread hop whenever ``workflow.in_workflow()`` is + false, so it is inert outside workflows. + + Raises whatever the lazy ``langsmith`` import or + ``set_runtime_overrides`` call raises (e.g. ``ImportError`` when + LangSmith is absent); the installed flag stays unset on failure so a + later call can retry. + """ + global _installed # noqa: PLW0603 + if _installed: + return + import langsmith + + langsmith.set_runtime_overrides(aio_to_thread=_temporal_aio_to_thread) + _installed = True diff --git a/temporalio/contrib/_langchain/_converter.py b/temporalio/contrib/_langchain/_converter.py new file mode 100644 index 000000000..e5fc39756 --- /dev/null +++ b/temporalio/contrib/_langchain/_converter.py @@ -0,0 +1,61 @@ +"""Payload converter shared by the LangChain-family plugins. + +This module eagerly imports ``temporalio.contrib.pydantic``; only plugins +that actually wire the converter import it, so the core package itself stays +importable without pydantic installed. +""" + +from __future__ import annotations + +import dataclasses + +from temporalio.contrib.pydantic import PydanticPayloadConverter, ToJsonOptions +from temporalio.converter import DataConverter + + +class LangChainPayloadConverter(PydanticPayloadConverter): + """Pydantic payload converter pinned to ``exclude_unset=True``. + + LangChain request/response types are deeply nested with many + ``Optional[...] = None`` fields. Shipping every unset default inflates + payloads several-fold and some peers reject the explicit nulls on + round-trip, so unset fields are excluded by convention. + """ + + def __init__(self) -> None: + """Construct the converter with ``exclude_unset`` serialization.""" + super().__init__(ToJsonOptions(exclude_unset=True)) + + +data_converter = DataConverter(payload_converter_class=LangChainPayloadConverter) +"""The family default data converter (LangChain messages are shipped as their +``dumpd`` JSON form, so the Pydantic converter only ever sees plain +containers).""" + + +def build_data_converter( + user_converter: DataConverter | None, + *, + plugin_name: str, +) -> DataConverter: + """Compose the family converter with whatever the caller already set. + + * ``None`` — install the family default. + * the SDK default converter — swap in the LangChain-aware Pydantic + converter via :func:`dataclasses.replace`. + * a custom converter — refuse rather than silently clobber it; the caller + must fold ``LangChainPayloadConverter`` into their own converter. + ``plugin_name`` names the refusing plugin in the error. + """ + if user_converter is None: + return data_converter + if user_converter is DataConverter.default: + return dataclasses.replace( + user_converter, payload_converter_class=LangChainPayloadConverter + ) + raise ValueError( + f"{plugin_name} cannot compose with a custom data_converter " + "automatically. Set payload_converter_class=LangChainPayloadConverter " + "on your own DataConverter (so LangChain messages serialize with " + "exclude_unset=True), or omit data_converter to use the plugin default." + ) diff --git a/temporalio/contrib/_langchain/_messages.py b/temporalio/contrib/_langchain/_messages.py new file mode 100644 index 000000000..946f9383e --- /dev/null +++ b/temporalio/contrib/_langchain/_messages.py @@ -0,0 +1,44 @@ +"""LangChain object serialization shared by the LangChain-family plugins.""" + +from __future__ import annotations + +from typing import Any + + +def dump_object(obj: Any) -> Any: + """Serialize a single LangChain ``Serializable`` (message, tool call, …).""" + from langchain_core.load import dumpd + + return dumpd(obj) + + +def load_object(data: Any) -> Any: + """Rehydrate a value produced by :func:`dump_object`, preserving subtype.""" + from langchain_core.load import load + + return load(data) + + +def dump_messages(messages: Any) -> list[Any]: + """Serialize a sequence of LangChain messages to their ``dumpd`` form.""" + from langchain_core.load import dumpd + + return [dumpd(m) for m in messages] + + +def load_messages(dumped: list[Any]) -> list[Any]: + """Rehydrate messages serialized by :func:`dump_messages`.""" + from langchain_core.load import load + + return [load(d) for d in dumped] + + +def tool_to_schema(tool: Any) -> dict[str, Any]: + """Advertise a tool to the model as a full OpenAI tool schema. + + Carries name + description + argument JSON schema so the model can build + valid arguments, not just select the tool by name. + """ + from langchain_core.utils.function_calling import convert_to_openai_tool + + return convert_to_openai_tool(tool) diff --git a/temporalio/contrib/_langchain/_passthrough.py b/temporalio/contrib/_langchain/_passthrough.py new file mode 100644 index 000000000..8979683b8 --- /dev/null +++ b/temporalio/contrib/_langchain/_passthrough.py @@ -0,0 +1,22 @@ +"""Sandbox passthrough-list merging shared by the LangChain-family plugins. + +Only the MECHANISM lives here; each plugin owns its module list (the lists +are behavior for released plugins and must not drift by sharing). +""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence + + +def merge_passthrough_modules( + defaults: Sequence[str], user: Iterable[str] | None +) -> tuple[str, ...]: + """Merge caller-supplied passthrough modules with a plugin's defaults. + + Order-preserving: defaults first, then user additions, first occurrence + wins on duplicates. + """ + merged = [*defaults, *(user or ())] + # dict.fromkeys preserves order while de-duplicating. + return tuple(dict.fromkeys(merged)) diff --git a/temporalio/contrib/_langchain/_runnable_config.py b/temporalio/contrib/_langchain/_runnable_config.py new file mode 100644 index 000000000..8ee9d3a71 --- /dev/null +++ b/temporalio/contrib/_langchain/_runnable_config.py @@ -0,0 +1,103 @@ +"""``RunnableConfig`` strip/rebuild shared by the LangChain-family plugins.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, cast + +if TYPE_CHECKING: + from langchain_core.runnables import RunnableConfig + + +def is_jsonish(value: Any) -> bool: + """Return True for values representable as plain JSON. + + Containers are checked recursively. Tuples are accepted — they encode + as JSON arrays and rehydrate as lists, so a round-trip preserves + content but may change the container type. + """ + if value is None or isinstance(value, (str, int, float, bool)): + return True + if isinstance(value, (list, tuple)): + return all(is_jsonish(v) for v in value) + if isinstance(value, dict): + return all(isinstance(k, str) and is_jsonish(v) for k, v in value.items()) + return False + + +def strip_runnable_config( + config: Mapping[str, Any] | None, + *, + configurable_keys: Sequence[str] | None = None, + configurable_filter: Callable[[str, Any], bool] | None = None, + metadata_filter: Callable[[Any], bool] | None = None, +) -> dict[str, Any]: + """Return a serializable subset of a ``RunnableConfig``. + + The full object holds non-serializable things (callbacks, checkpointer / + store / cache handles, pregel send/read callables) that cannot cross an + activity boundary, so only primitive fields and a caller-selected subset + of ``configurable`` are kept. The output shape follows the langgraph + plugin's released behavior: ``tags`` and ``metadata`` are always present + (empty when absent), ``run_name`` / ``run_id`` are kept when truthy, + ``recursion_limit`` when not None, and ``configurable`` only when the + selection is non-empty. + + Exactly one of ``configurable_keys`` (keys kept in whitelist order) or + ``configurable_filter`` (entries kept in source order) must be provided. + ``metadata_filter`` optionally drops metadata VALUES (default: keep all, + matching langgraph). + """ + if (configurable_keys is None) == (configurable_filter is None): + raise ValueError( + "exactly one of configurable_keys or configurable_filter is required" + ) + orig: Mapping[str, Any] = config or {} + configurable: Mapping[str, Any] = orig.get("configurable") or {} + + metadata = dict(orig.get("metadata") or {}) + if metadata_filter is not None: + metadata = {k: v for k, v in metadata.items() if metadata_filter(v)} + result: dict[str, Any] = { + "tags": list(orig.get("tags") or []), + "metadata": metadata, + } + if run_name := orig.get("run_name"): + result["run_name"] = run_name + if run_id := orig.get("run_id"): + result["run_id"] = run_id + if (recursion_limit := orig.get("recursion_limit")) is not None: + result["recursion_limit"] = recursion_limit + + if configurable_keys is not None: + stripped_configurable: dict[str, Any] = { + key: configurable[key] for key in configurable_keys if key in configurable + } + else: + stripped_configurable = { + k: v + for k, v in configurable.items() + # The guard above makes configurable_filter non-None on this + # branch; the re-check narrows for type checkers. + if configurable_filter is not None and configurable_filter(k, v) + } + if stripped_configurable: + result["configurable"] = stripped_configurable + return result + + +def rebuild_runnable_config(data: dict[str, Any]) -> "RunnableConfig": + """Reconstruct a minimal ``RunnableConfig`` from :func:`strip_runnable_config`.""" + config: dict[str, Any] = {"metadata": dict(data.get("metadata", {}))} + if data.get("tags"): + config["tags"] = list(data["tags"]) + for key in ("run_id", "run_name"): + if data.get(key) is not None: + config[key] = data[key] + if data.get("recursion_limit") is not None: + config["recursion_limit"] = data["recursion_limit"] + if data.get("configurable"): + config["configurable"] = dict(data["configurable"]) + # Double cast: a TypedDict and dict[str, Any] "insufficiently overlap" + # for basedpyright's reportInvalidCast; object is the sanctioned bridge. + return cast("RunnableConfig", cast(object, config)) diff --git a/temporalio/contrib/_langchain/_task_cache.py b/temporalio/contrib/_langchain/_task_cache.py new file mode 100644 index 000000000..66bd43c4c --- /dev/null +++ b/temporalio/contrib/_langchain/_task_cache.py @@ -0,0 +1,91 @@ +"""Continue-as-new result caching shared by the LangChain-family plugins.""" + +from __future__ import annotations + +from contextvars import ContextVar +from hashlib import sha256 +from json import dumps +from typing import Any + + +class TaskResultCache: + """A workflow-scoped result cache carried across continue-as-new. + + Each plugin owns its own instance, so caches never share state across + plugins. The backing store is a ``ContextVar``, so update / signal / + query handler tasks spawned by the workflow inherit the same cache + automatically. The cache itself is a plain dict that can travel through + ``workflow.continue_as_new()``. + + ``set_cache`` stores the mapping AS GIVEN (identity semantics — the + langgraph plugin exposes the live dict so mutations ride into the next + run); callers that want copy-on-set semantics wrap at their own layer. + """ + + def __init__(self, context_var_name: str) -> None: + """Create the cache with a uniquely named backing ``ContextVar``.""" + self._var: ContextVar[dict[str, Any] | None] = ContextVar( + context_var_name, default=None + ) + + def set_cache(self, cache: dict[str, Any] | None) -> None: + """Set the result cache for the current context (stored as given).""" + self._var.set(cache) + + def get_cache(self) -> dict[str, Any] | None: + """Get the result cache for the current context.""" + return self._var.get() + + def lookup(self, key: str) -> tuple[bool, Any]: + """Return ``(True, value)`` if cached, ``(False, None)`` otherwise.""" + cache = self._var.get() + if cache is not None and key in cache: + return True, cache[key] + return False, None + + def put(self, key: str, value: Any) -> None: + """Store a value when a cache is active for this context.""" + cache = self._var.get() + if cache is not None: + cache[key] = value + + +def task_id(func: Any) -> str: + """Return the fully-qualified module.qualname for a function. + + Raises ValueError for functions that cannot be identified unambiguously + (lambdas, closures, __main__ functions). + """ + module = getattr(func, "__module__", None) + qualname = getattr(func, "__qualname__", None) or getattr(func, "__name__", None) + + if module is None or qualname is None: + raise ValueError( + f"Cannot identify task {func}: missing __module__ or __qualname__. " + "Tasks must be defined at module level." + ) + if module == "__main__": + raise ValueError( + f"Cannot identify task {qualname}: defined in __main__. " + "Tasks must be importable from a named module." + ) + if "" in qualname: + raise ValueError( + f"Cannot identify task {qualname}: closures/local functions are not supported. " + "Tasks must be defined at module level." + ) + return f"{module}.{qualname}" + + +def cache_key( + task_id: str, + args: tuple[Any, ...], + kwargs: dict[str, Any], + context: Any = None, +) -> str: + """Build a cache key from the full task identifier, arguments, and runtime context.""" + try: + key_str = dumps([task_id, args, kwargs, context], sort_keys=True, default=str) + except (TypeError, ValueError): + key_str = repr([task_id, args, kwargs, context]) + return sha256(key_str.encode()).hexdigest()[:32] diff --git a/temporalio/contrib/deepagents/_activity.py b/temporalio/contrib/deepagents/_activity.py index 3c1501a42..4846d51c2 100644 --- a/temporalio/contrib/deepagents/_activity.py +++ b/temporalio/contrib/deepagents/_activity.py @@ -27,10 +27,10 @@ import dataclasses import importlib from datetime import timedelta -from functools import wraps from typing import Any, Callable from temporalio import activity +from temporalio.contrib._langchain import _activity_helpers from temporalio.contrib.deepagents import _serde from temporalio.exceptions import ApplicationError @@ -124,90 +124,6 @@ class BackendOpOutput: # --------------------------------------------------------------------------- -def _auto_heartbeater(fn: Callable) -> Callable: - """Heartbeat at half the configured ``heartbeat_timeout`` while ``fn`` runs. - - Long LLM calls (thinking mode, long context, streaming accumulation) can run - well past a scheduler's patience; without a heartbeat Temporal would cancel - them and surface a ``HeartbeatTimeoutError`` instead of the real problem. - """ - - @wraps(fn) - async def wrapped(*args: Any, **kwargs: Any) -> Any: - heartbeat_timeout = activity.info().heartbeat_timeout - beat_task: asyncio.Task | None = None - if heartbeat_timeout: - interval = heartbeat_timeout.total_seconds() / 2 - - async def beat() -> None: - while True: - activity.heartbeat() - await asyncio.sleep(interval) - - beat_task = asyncio.create_task(beat()) - try: - return await fn(*args, **kwargs) - finally: - if beat_task is not None: - beat_task.cancel() - # Let the cancellation land before returning so no pending task - # outlives the activity (a bare ``cancel()`` leaves the task to - # be destroyed while pending if the loop shuts down first). - # ``asyncio.wait`` never re-raises the task's CancelledError. - await asyncio.wait([beat_task]) - - return wrapped - - -def _translate_api_error(exc: Exception) -> ApplicationError | None: - """Map an LLM SDK HTTP error onto Temporal's retry contract. - - Works by duck typing so neither ``openai`` nor ``anthropic`` needs to be - imported here: both expose ``status_code`` and ``response.headers``. Returns - ``None`` when ``exc`` is not a recognizable HTTP status error, so the caller - can fall through to its generic handling. - """ - status = getattr(exc, "status_code", None) - if status is None: - return None - headers: dict[str, Any] = {} - response = getattr(exc, "response", None) - if response is not None: - headers = dict(getattr(response, "headers", {}) or {}) - # Case-insensitive header access. - lower = {str(k).lower(): v for k, v in headers.items()} - - retryable = status in (408, 409, 429) or 500 <= status < 600 - should_retry = lower.get("x-should-retry") - if should_retry == "false": - retryable = False - elif should_retry == "true": - retryable = True - - delay_ms = lower.get("retry-after-ms") - retry_after = lower.get("retry-after") - next_delay: timedelta | None = None - try: - if delay_ms is not None: - next_delay = timedelta(milliseconds=int(delay_ms)) - elif retry_after is not None: - next_delay = timedelta(seconds=int(retry_after)) - except (TypeError, ValueError): - next_delay = None - - return ApplicationError( - str(exc), - type=type(exc).__name__, - non_retryable=not retryable, - next_retry_delay=next_delay, - ) - - -# --------------------------------------------------------------------------- -# The activities -# --------------------------------------------------------------------------- - - def _default_model_provider(model_name: str) -> Any: """Build a chat model from a name string with LLM-SDK retries disabled. @@ -249,7 +165,7 @@ def _build_bound_model(self, input: ModelActivityInput) -> Any: return model @activity.defn(name=INVOKE_MODEL) - @_auto_heartbeater + @_activity_helpers.auto_heartbeater async def invoke_model(self, input: ModelActivityInput) -> ModelActivityOutput: """Run exactly one LLM call and return the resulting ``AIMessage``.""" messages = _serde.load_messages(input.messages) @@ -258,7 +174,7 @@ async def invoke_model(self, input: ModelActivityInput) -> ModelActivityOutput: try: message = await model.ainvoke(messages, config=config) except Exception as exc: - translated = _translate_api_error(exc) + translated = _activity_helpers.translate_api_error(exc) if translated is not None: activity.logger.warning( "Model call failed with an HTTP status error", exc_info=True @@ -268,7 +184,7 @@ async def invoke_model(self, input: ModelActivityInput) -> ModelActivityOutput: return ModelActivityOutput(message=_serde.dump_object(message)) @activity.defn(name=INVOKE_MODEL_STREAMING) - @_auto_heartbeater + @_activity_helpers.auto_heartbeater async def invoke_model_streaming( self, input: ModelActivityInput ) -> ModelActivityOutput: @@ -300,7 +216,7 @@ async def invoke_model_streaming( topic.publish(_serde.dump_object(chunk)) final = chunk if final is None else final + chunk except Exception as exc: - translated = _translate_api_error(exc) + translated = _activity_helpers.translate_api_error(exc) if translated is not None: activity.logger.warning( "Streaming model call failed with an HTTP status error", @@ -311,7 +227,7 @@ async def invoke_model_streaming( return ModelActivityOutput(message=_serde.dump_object(final)) @activity.defn(name=INVOKE_TOOL) - @_auto_heartbeater + @_activity_helpers.auto_heartbeater async def invoke_tool(self, input: ToolActivityInput) -> ToolActivityOutput: """Execute one registered tool and return its ``ToolMessage``.""" from temporalio.contrib.deepagents._tools import get_registered_tool @@ -335,7 +251,7 @@ async def invoke_tool(self, input: ToolActivityInput) -> ToolActivityOutput: return ToolActivityOutput(message=_serde.dump_object(message)) @activity.defn(name=BACKEND_OP) - @_auto_heartbeater + @_activity_helpers.auto_heartbeater async def backend_op(self, input: BackendOpInput) -> BackendOpOutput: """Run one operation against a registered (real-I/O) backend.""" from temporalio.contrib.deepagents._tools import registered_backends diff --git a/temporalio/contrib/deepagents/_plugin.py b/temporalio/contrib/deepagents/_plugin.py index 3b95fcdaf..1efecd85f 100644 --- a/temporalio/contrib/deepagents/_plugin.py +++ b/temporalio/contrib/deepagents/_plugin.py @@ -226,52 +226,21 @@ def configure_worker(self, config: Any) -> Any: return config -async def _temporal_aio_to_thread( - default_aio_to_thread: Callable[..., Any], - ctx: Any, - func: Callable[..., Any], - /, - *args: Any, - **kwargs: Any, -) -> Any: - """Run LangSmith's ``aio_to_thread`` seam safely inside a workflow. - - Outside a workflow (activities, clients) we defer to LangSmith's default - thread hop. Inside a workflow the deterministic event loop cannot spawn a - thread, so we run the setup inline in the requested context. ``func`` here is - LangSmith's own run-tree bookkeeping — cheap and side-effect-free with respect - to workflow determinism. - """ - from temporalio import workflow - - if not workflow.in_workflow(): - return await default_aio_to_thread(ctx, func, *args, **kwargs) - with workflow.unsafe.sandbox_unrestricted(): - return ctx.run(func, *args, **kwargs) - - def _install_langsmith_temporal_override() -> None: """Route LangSmith's thread hop inline while in a workflow. - ``set_runtime_overrides`` mutates a module global in the sandbox-passthrough - ``langsmith`` package, so the in-workflow tracer sees it too. No-op (and - harmless) when LangSmith is not installed or too old to expose + Delegates to the shared installer in ``temporalio.contrib._langchain`` + (the same override ``temporalio.contrib.langsmith`` installs, once per + process, never uninstalled — see its docstring for the rationale). No-op + (and harmless) when LangSmith is not installed or too old to expose ``set_runtime_overrides``. - - The override is deliberately installed for the life of the process and - never uninstalled: LangSmith exposes a single process-wide override slot - (each ``set_runtime_overrides`` call replaces it wholesale), and - ``temporalio.contrib.langsmith`` installs a behaviorally identical override - exactly once, without reinstalling. Resetting the slot on this worker's - shutdown would therefore permanently strip a composed LangSmith plugin's - workflow-safety override. Leaving it installed is safe: the override defers - to LangSmith's default thread hop whenever ``workflow.in_workflow()`` is - false, so it is inert outside workflows. """ try: - import langsmith + from temporalio.contrib._langchain._aio_to_thread import ( + install_aio_to_thread_override, + ) - langsmith.set_runtime_overrides(aio_to_thread=_temporal_aio_to_thread) + install_aio_to_thread_override() except Exception: pass diff --git a/temporalio/contrib/deepagents/_serde.py b/temporalio/contrib/deepagents/_serde.py index 13ed0d1c6..27b472cc4 100644 --- a/temporalio/contrib/deepagents/_serde.py +++ b/temporalio/contrib/deepagents/_serde.py @@ -29,18 +29,53 @@ from __future__ import annotations -import contextvars import dataclasses -import hashlib -import json -from typing import TYPE_CHECKING, Any, cast +from typing import Any -if TYPE_CHECKING: - from langchain_core.runnables import RunnableConfig - -from temporalio.contrib.pydantic import PydanticPayloadConverter, ToJsonOptions +from temporalio.contrib._langchain import _converter, _runnable_config, _task_cache +from temporalio.contrib._langchain._converter import ( + LangChainPayloadConverter, + data_converter, +) +from temporalio.contrib._langchain._messages import ( + dump_messages, + dump_object, + load_messages, + load_object, + tool_to_schema, +) +from temporalio.contrib._langchain._passthrough import merge_passthrough_modules +from temporalio.contrib._langchain._runnable_config import ( + is_jsonish, + rebuild_runnable_config, +) from temporalio.converter import DataConverter +__all__ = [ + "DeepAgentsPayloadConverter", + "Settings", + "build_data_converter", + "cache_key", + "cache_lookup", + "cache_put", + "data_converter", + "default_passthrough_modules", + "dump_backend_result", + "dump_messages", + "dump_object", + "get_settings", + "load_backend_result", + "load_messages", + "load_object", + "rebuild_runnable_config", + "resolve_passthrough_modules", + "result_cache_snapshot", + "set_result_cache", + "set_settings", + "strip_runnable_config", + "tool_to_schema", +] + # --------------------------------------------------------------------------- # Worker-wide dispatch settings # --------------------------------------------------------------------------- @@ -85,23 +120,8 @@ def get_settings() -> Settings: # --------------------------------------------------------------------------- -class DeepAgentsPayloadConverter(PydanticPayloadConverter): - """Pydantic payload converter pinned to ``exclude_unset=True``. - - LangChain request/response types (and ``DeepAgentState``) are deeply nested - with many ``Optional[...] = None`` fields. Shipping every unset default - inflates payloads several-fold and some peers reject the explicit nulls on - round-trip, so we exclude unset fields by convention. - """ - - def __init__(self) -> None: - """Construct the converter with ``exclude_unset`` serialization.""" - super().__init__(ToJsonOptions(exclude_unset=True)) - - -data_converter = DataConverter(payload_converter_class=DeepAgentsPayloadConverter) -"""The plugin's default data converter (LangChain messages are shipped as their -``dumpd`` JSON form, so the Pydantic converter only ever sees plain containers).""" +# The family converter, kept under this plugin's historical name. +DeepAgentsPayloadConverter = LangChainPayloadConverter def build_data_converter( @@ -109,23 +129,11 @@ def build_data_converter( ) -> DataConverter: """Compose the plugin's converter with whatever the caller already set. - * ``None`` — install the plugin default. - * the SDK default converter — swap in the LangChain-aware Pydantic - converter via :func:`dataclasses.replace`. - * a custom converter — refuse rather than silently clobber it; the caller - must fold :class:`DeepAgentsPayloadConverter` into their own converter. + Delegates to the shared family implementation; see its docstring for the + None / SDK-default / custom-converter contract. """ - if user_converter is None: - return data_converter - if user_converter is DataConverter.default: - return dataclasses.replace( - user_converter, payload_converter_class=DeepAgentsPayloadConverter - ) - raise ValueError( - "DeepAgentsPlugin cannot compose with a custom data_converter " - "automatically. Set payload_converter_class=DeepAgentsPayloadConverter " - "on your own DataConverter (so LangChain messages serialize with " - "exclude_unset=True), or omit data_converter to use the plugin default." + return _converter.build_data_converter( + user_converter, plugin_name="DeepAgentsPlugin" ) @@ -134,20 +142,6 @@ def build_data_converter( # --------------------------------------------------------------------------- -def dump_object(obj: Any) -> Any: - """Serialize a single LangChain ``Serializable`` (message, tool call, …).""" - from langchain_core.load import dumpd - - return dumpd(obj) - - -def load_object(data: Any) -> Any: - """Rehydrate a value produced by :func:`dump_object`, preserving subtype.""" - from langchain_core.load import load - - return load(data) - - # --------------------------------------------------------------------------- # Backend protocol result (de)serialization # --------------------------------------------------------------------------- @@ -233,94 +227,36 @@ def load_backend_result(value: Any) -> Any: return value -def dump_messages(messages: Any) -> list[Any]: - """Serialize a sequence of LangChain messages to their ``dumpd`` form.""" - from langchain_core.load import dumpd - - return [dumpd(m) for m in messages] - - -def load_messages(dumped: list[Any]) -> list[Any]: - """Rehydrate messages serialized by :func:`dump_messages`.""" - from langchain_core.load import load - - return [load(d) for d in dumped] - - -def tool_to_schema(tool: Any) -> dict[str, Any]: - """Advertise a tool to the model as a full OpenAI tool schema. - - Carries name + description + argument JSON schema so the model can build - valid arguments, not just select the tool by name. - """ - from langchain_core.utils.function_calling import convert_to_openai_tool - - return convert_to_openai_tool(tool) - - # --------------------------------------------------------------------------- # RunnableConfig strip / rebuild # --------------------------------------------------------------------------- -def _is_jsonish(value: Any) -> bool: - if value is None or isinstance(value, (str, int, float, bool)): - return True - if isinstance(value, (list, tuple)): - return all(_is_jsonish(v) for v in value) - if isinstance(value, dict): - return all(isinstance(k, str) and _is_jsonish(v) for k, v in value.items()) - return False +# Shared with the other LangChain-family plugins; kept under the old private +# name because sibling modules reference it. +_is_jsonish = is_jsonish + + +def _keep_configurable(key: str, value: Any) -> bool: + """Keep non-dunder, JSON-safe configurable entries.""" + return not key.startswith("__") and is_jsonish(value) def strip_runnable_config(config: Any) -> dict[str, Any]: """Reduce a live ``RunnableConfig`` to its JSON-safe subset for shipping. - Keeps ``tags``, ``run_name``, ``run_id``, ``recursion_limit``, JSON-safe - ``metadata`` and the JSON-safe (non-dunder) ``configurable`` keys. Drops - callbacks, checkpointer / store / cache handles and every other live - reference — those are reconstructed activity-side. + Keeps ``tags`` and ``metadata`` (always present, values filtered to + JSON-safe ones), truthy ``run_name`` / ``run_id``, ``recursion_limit``, + and the JSON-safe non-dunder ``configurable`` keys. Drops callbacks, + checkpointer / store / cache handles and every other live reference — + those are reconstructed activity-side. Output shape follows the shared + (langgraph-vetted) implementation. """ - if not config: - return {} - out: dict[str, Any] = {} - if config.get("tags"): - out["tags"] = list(config["tags"]) - for key in ("run_id", "run_name"): - if config.get(key) is not None: - out[key] = str(config[key]) - if config.get("recursion_limit") is not None: - out["recursion_limit"] = config["recursion_limit"] - metadata = config.get("metadata") - if isinstance(metadata, dict): - out["metadata"] = {k: v for k, v in metadata.items() if _is_jsonish(v)} - configurable = config.get("configurable") - if isinstance(configurable, dict): - kept = { - k: v - for k, v in configurable.items() - if not k.startswith("__") and _is_jsonish(v) - } - if kept: - out["configurable"] = kept - return out - - -def rebuild_runnable_config(data: dict[str, Any]) -> "RunnableConfig": - """Reconstruct a minimal ``RunnableConfig`` from :func:`strip_runnable_config`.""" - config: dict[str, Any] = {"metadata": dict(data.get("metadata", {}))} - if data.get("tags"): - config["tags"] = list(data["tags"]) - for key in ("run_id", "run_name"): - if data.get(key) is not None: - config[key] = data[key] - if data.get("recursion_limit") is not None: - config["recursion_limit"] = data["recursion_limit"] - if data.get("configurable"): - config["configurable"] = dict(data["configurable"]) - # Double cast: a TypedDict and dict[str, Any] "insufficiently overlap" - # for basedpyright's reportInvalidCast; object is the sanctioned bridge. - return cast("RunnableConfig", cast(object, config)) + return _runnable_config.strip_runnable_config( + config, + configurable_filter=_keep_configurable, + metadata_filter=is_jsonish, + ) # --------------------------------------------------------------------------- @@ -328,44 +264,37 @@ def rebuild_runnable_config(data: dict[str, Any]) -> "RunnableConfig": # --------------------------------------------------------------------------- # Per-workflow state: set at the top of the workflow run, read by the model and -# tool dispatch paths, snapshotted for continue-as-new. A ContextVar (not a -# module-global keyed by run id) so update / signal handler tasks spawned by the -# workflow inherit the same cache automatically. -_result_cache: contextvars.ContextVar[dict[str, Any] | None] = contextvars.ContextVar( - "_deepagents_result_cache", default=None -) +# tool dispatch paths, snapshotted for continue-as-new. Built on the shared +# ContextVar-backed cache (update / signal handler tasks spawned by the +# workflow inherit it automatically); this plugin's instance is distinct from +# the langgraph plugin's. +_result_cache = _task_cache.TaskResultCache("_deepagents_result_cache") def set_result_cache(cache: dict[str, Any] | None) -> None: """Seed the workflow-scoped result cache (e.g. carried across CAN).""" - _result_cache.set(dict(cache) if cache else {}) + _result_cache.set_cache(dict(cache) if cache else {}) def result_cache_snapshot() -> dict[str, Any] | None: """Return a serializable copy of the cache, or ``None`` when empty.""" - cache = _result_cache.get() + cache = _result_cache.get_cache() return dict(cache) if cache else None def cache_key(kind: str, call_id: str, args: Any) -> str: """Stable key over ``(kind, call_id, args)`` for cache lookups.""" - blob = json.dumps([kind, call_id, args], default=str, sort_keys=True) - return hashlib.sha256(blob.encode("utf-8")).hexdigest() + return _task_cache.cache_key(kind, (call_id, args), {}) def cache_lookup(key: str) -> tuple[bool, Any]: """Return ``(hit, value)`` for ``key`` in the active cache.""" - cache = _result_cache.get() - if cache is not None and key in cache: - return True, cache[key] - return False, None + return _result_cache.lookup(key) def cache_put(key: str, value: Any) -> None: """Record ``value`` under ``key`` when a cache is active for this run.""" - cache = _result_cache.get() - if cache is not None: - cache[key] = value + _result_cache.put(key, value) # --------------------------------------------------------------------------- @@ -405,6 +334,4 @@ def default_passthrough_modules() -> tuple[str, ...]: def resolve_passthrough_modules(user: Any) -> tuple[str, ...]: """Merge caller-supplied passthrough modules with the plugin defaults.""" - merged = [*_DEFAULT_PASSTHROUGH, *(user or ())] - # dict.fromkeys preserves order while de-duplicating. - return tuple(dict.fromkeys(merged)) + return merge_passthrough_modules(_DEFAULT_PASSTHROUGH, user) diff --git a/temporalio/contrib/langgraph/_interceptor.py b/temporalio/contrib/langgraph/_interceptor.py index f68d9d45d..2409481f9 100644 --- a/temporalio/contrib/langgraph/_interceptor.py +++ b/temporalio/contrib/langgraph/_interceptor.py @@ -11,7 +11,7 @@ from temporalio import workflow from temporalio.contrib.langgraph._activity import clear_store_warning -from temporalio.contrib.workflow_streams._stream import _PUBLISH_SIGNAL +from temporalio.contrib.workflow_streams import current_workflow_stream from temporalio.worker import ( ExecuteWorkflowInput, Interceptor, @@ -54,10 +54,7 @@ def init(self, outbound: WorkflowOutboundInterceptor) -> None: super().init(outbound) async def execute_workflow(self, input: ExecuteWorkflowInput) -> Any: - if ( - streaming_topic is not None - and workflow.get_signal_handler(_PUBLISH_SIGNAL) is None - ): + if streaming_topic is not None and current_workflow_stream() is None: raise RuntimeError( f"LangGraphPlugin was configured with " f"streaming_topic={streaming_topic!r}, but workflow " diff --git a/temporalio/contrib/langgraph/_langgraph_config.py b/temporalio/contrib/langgraph/_langgraph_config.py index 90c6c810d..68caf30e1 100644 --- a/temporalio/contrib/langgraph/_langgraph_config.py +++ b/temporalio/contrib/langgraph/_langgraph_config.py @@ -3,7 +3,7 @@ # pyright: reportMissingTypeStubs=false import dataclasses -from typing import Any, Callable +from typing import Any, Callable, cast from langchain_core.runnables.config import var_child_runnable_config from langgraph._internal._constants import ( @@ -23,6 +23,21 @@ from langgraph.pregel._algo import LazyAtomicCounter from langgraph.runtime import ExecutionInfo, Runtime +from temporalio.contrib._langchain import _runnable_config + +# The configurable keys the langgraph plugin ships across activity +# boundaries (checkpoint/resumption state); everything else in +# ``configurable`` is a live handle that stays behind. +_KEPT_CONFIGURABLE_KEYS = ( + CONFIG_KEY_CHECKPOINT_NS, + CONFIG_KEY_CHECKPOINT_ID, + CONFIG_KEY_CHECKPOINT_MAP, + CONFIG_KEY_THREAD_ID, + CONFIG_KEY_TASK_ID, + CONFIG_KEY_RESUMING, + CONFIG_KEY_DURABILITY, +) + def strip_runnable_config(config: RunnableConfig | None) -> RunnableConfig: """Return a serializable subset of a RunnableConfig. @@ -31,38 +46,21 @@ def strip_runnable_config(config: RunnableConfig | None) -> RunnableConfig: config kwarg. The full object holds non-serializable things (callbacks, checkpointer/store/cache handles, pregel send/read callables) that can't cross an activity boundary, so we keep only primitive fields and the - serializable subset of configurable. + serializable subset of configurable. Delegates to the shared + implementation with this plugin's checkpoint-key whitelist; the output is + byte-identical to the pre-refactor behavior. """ - orig = config or {} - configurable = orig.get("configurable") or {} - - result: RunnableConfig = { - "tags": list(orig.get("tags") or []), - "metadata": dict(orig.get("metadata") or {}), - } - if run_name := orig.get("run_name"): - result["run_name"] = run_name - if run_id := orig.get("run_id"): - result["run_id"] = run_id - if (recursion_limit := orig.get("recursion_limit")) is not None: - result["recursion_limit"] = recursion_limit - - stripped_configurable: dict[str, Any] = { - key: configurable[key] - for key in ( - CONFIG_KEY_CHECKPOINT_NS, - CONFIG_KEY_CHECKPOINT_ID, - CONFIG_KEY_CHECKPOINT_MAP, - CONFIG_KEY_THREAD_ID, - CONFIG_KEY_TASK_ID, - CONFIG_KEY_RESUMING, - CONFIG_KEY_DURABILITY, - ) - if key in configurable - } - if stripped_configurable: - result["configurable"] = stripped_configurable - return result + # Double cast: a TypedDict and dict[str, Any] "insufficiently overlap" + # for basedpyright's reportInvalidCast; object is the sanctioned bridge. + return cast( + "RunnableConfig", + cast( + object, + _runnable_config.strip_runnable_config( + config, configurable_keys=_KEPT_CONFIGURABLE_KEYS + ), + ), + ) def get_langgraph_config() -> dict[str, Any]: @@ -124,17 +122,7 @@ def get_null_resume(consume: bool = False) -> Any: ) restored_configurable: dict[str, Any] = { - key: configurable[key] - for key in ( - CONFIG_KEY_CHECKPOINT_NS, - CONFIG_KEY_CHECKPOINT_ID, - CONFIG_KEY_CHECKPOINT_MAP, - CONFIG_KEY_THREAD_ID, - CONFIG_KEY_TASK_ID, - CONFIG_KEY_RESUMING, - CONFIG_KEY_DURABILITY, - ) - if key in configurable + key: configurable[key] for key in _KEPT_CONFIGURABLE_KEYS if key in configurable } restored_configurable[CONFIG_KEY_SCRATCHPAD] = PregelScratchpad( step=scratchpad.get("step", 0), diff --git a/temporalio/contrib/langgraph/_task_cache.py b/temporalio/contrib/langgraph/_task_cache.py index ab3e683d7..a2e27f020 100644 --- a/temporalio/contrib/langgraph/_task_cache.py +++ b/temporalio/contrib/langgraph/_task_cache.py @@ -3,81 +3,49 @@ Caches task results by (module.qualname, args, kwargs) hash so that previously completed tasks are not re-executed after a continue-as-new. The cache state is a plain dict that can travel through workflow.continue_as_new(). + +The mechanism lives in ``temporalio.contrib._langchain._task_cache``; this +module binds the langgraph plugin's own cache instance and preserves the +original function surface. """ from __future__ import annotations -from contextvars import ContextVar -from hashlib import sha256 -from json import dumps from typing import Any -_task_cache: ContextVar[dict[str, Any] | None] = ContextVar( - "_temporal_task_cache", default=None +from temporalio.contrib._langchain._task_cache import ( + TaskResultCache, + cache_key, + task_id, ) +__all__ = [ + "cache_key", + "cache_lookup", + "cache_put", + "get_task_cache", + "set_task_cache", + "task_id", +] + +_cache = TaskResultCache("_temporal_task_cache") + def set_task_cache(cache: dict[str, Any] | None) -> None: """Set the task result cache for the current context.""" - _task_cache.set(cache) + _cache.set_cache(cache) def get_task_cache() -> dict[str, Any] | None: """Get the task result cache for the current context.""" - return _task_cache.get() - - -def task_id(func: Any) -> str: - """Return the fully-qualified module.qualname for a function. - - Raises ValueError for functions that cannot be identified unambiguously - (lambdas, closures, __main__ functions). - """ - module = getattr(func, "__module__", None) - qualname = getattr(func, "__qualname__", None) or getattr(func, "__name__", None) - - if module is None or qualname is None: - raise ValueError( - f"Cannot identify task {func}: missing __module__ or __qualname__. " - "Tasks must be defined at module level." - ) - if module == "__main__": - raise ValueError( - f"Cannot identify task {qualname}: defined in __main__. " - "Tasks must be importable from a named module." - ) - if "" in qualname: - raise ValueError( - f"Cannot identify task {qualname}: closures/local functions are not supported. " - "Tasks must be defined at module level." - ) - return f"{module}.{qualname}" - - -def cache_key( - task_id: str, - args: tuple[Any, ...], - kwargs: dict[str, Any], - context: Any = None, -) -> str: - """Build a cache key from the full task identifier, arguments, and runtime context.""" - try: - key_str = dumps([task_id, args, kwargs, context], sort_keys=True, default=str) - except (TypeError, ValueError): - key_str = repr([task_id, args, kwargs, context]) - return sha256(key_str.encode()).hexdigest()[:32] + return _cache.get_cache() def cache_lookup(key: str) -> tuple[bool, Any]: """Return (True, value) if cached, (False, None) otherwise.""" - cache = _task_cache.get() - if cache is not None and key in cache: - return True, cache[key] - return False, None + return _cache.lookup(key) def cache_put(key: str, value: Any) -> None: """Store a value in the task result cache.""" - cache = _task_cache.get() - if cache is not None: - cache[key] = value + _cache.put(key, value) diff --git a/temporalio/contrib/langgraph/_workflow.py b/temporalio/contrib/langgraph/_workflow.py index 43b3d06ae..93d303e98 100644 --- a/temporalio/contrib/langgraph/_workflow.py +++ b/temporalio/contrib/langgraph/_workflow.py @@ -13,7 +13,7 @@ from langgraph._internal._constants import CONFIG_KEY_RUNTIME from temporalio import workflow -from temporalio.contrib.workflow_streams._stream import _PUBLISH_SIGNAL +from temporalio.contrib.workflow_streams import current_workflow_stream def wrap_workflow( @@ -65,8 +65,14 @@ async def run(stream_writer: Callable[[Any], None] | None) -> Any: if streaming_topic is None: return await run(stream_writer=None) - publish_handler = workflow.get_signal_handler(_PUBLISH_SIGNAL) - stream = getattr(publish_handler, "__self__") + stream = current_workflow_stream() + if stream is None: + # The interceptor validates stream presence at workflow start, so + # this is unreachable in a plugin-managed workflow; guard anyway. + raise RuntimeError( + "streaming_topic is configured but the workflow has no " + "WorkflowStream; construct WorkflowStream() in @workflow.init." + ) topic = stream.topic(streaming_topic) return await run(stream_writer=topic.publish) diff --git a/temporalio/contrib/langsmith/_interceptor.py b/temporalio/contrib/langsmith/_interceptor.py index a7eea0714..f88045fcf 100644 --- a/temporalio/contrib/langsmith/_interceptor.py +++ b/temporalio/contrib/langsmith/_interceptor.py @@ -24,6 +24,9 @@ import temporalio.worker import temporalio.workflow from temporalio.api.common.v1 import Payload +from temporalio.contrib._langchain._aio_to_thread import ( + install_aio_to_thread_override, +) from temporalio.exceptions import ApplicationError, ApplicationErrorCategory # This logger is only used in _log_future_exception, which runs on the @@ -153,49 +156,6 @@ def _get_current_run_for_propagation() -> RunTree | None: return run -# --------------------------------------------------------------------------- -# Workflow event loop safety: override @traceable's aio_to_thread -# --------------------------------------------------------------------------- - -_aio_to_thread_override_installed = False - - -async def _temporal_aio_to_thread( - default_aio_to_thread: Callable[..., Any], - ctx: Any, - func: Callable[..., Any], - /, - *args: Any, - **kwargs: Any, -) -> Any: - """Run LangSmith's ``aio_to_thread`` synchronously inside Temporal workflows. - - The ``@traceable`` decorator on async functions uses ``aio_to_thread()`` → - ``loop.run_in_executor()`` for run setup/teardown. The Temporal workflow - event loop does not support ``run_in_executor``. This override runs those - functions synchronously on the workflow thread when inside a workflow, - and delegates to the default implementation outside workflows. - - Registered via ``langsmith.set_runtime_overrides(aio_to_thread=...)``. - """ - if not temporalio.workflow.in_workflow(): - return await default_aio_to_thread(ctx, func, *args, **kwargs) - with temporalio.workflow.unsafe.sandbox_unrestricted(): - return ctx.run(func, *args, **kwargs) - - -def _install_aio_to_thread_override() -> None: - """Install the ``aio_to_thread`` override via LangSmith's official API. - - Safe to call multiple times; the override is only installed once. - """ - global _aio_to_thread_override_installed # noqa: PLW0603 - if _aio_to_thread_override_installed: - return - langsmith.set_runtime_overrides(aio_to_thread=_temporal_aio_to_thread) - _aio_to_thread_override_installed = True - - # --------------------------------------------------------------------------- # Replay safety # --------------------------------------------------------------------------- @@ -609,7 +569,7 @@ def workflow_interceptor_class( self, input: temporalio.worker.WorkflowInterceptorClassInput ) -> type[_LangSmithWorkflowInboundInterceptor]: """Return the workflow interceptor class with config bound.""" - _install_aio_to_thread_override() + install_aio_to_thread_override() config = self class InterceptorWithConfig(_LangSmithWorkflowInboundInterceptor): diff --git a/temporalio/contrib/workflow_streams/__init__.py b/temporalio/contrib/workflow_streams/__init__.py index 41f670f0c..5a756281b 100644 --- a/temporalio/contrib/workflow_streams/__init__.py +++ b/temporalio/contrib/workflow_streams/__init__.py @@ -13,7 +13,10 @@ """ from temporalio.contrib.workflow_streams._client import WorkflowStreamClient -from temporalio.contrib.workflow_streams._stream import WorkflowStream +from temporalio.contrib.workflow_streams._stream import ( + WorkflowStream, + current_workflow_stream, +) from temporalio.contrib.workflow_streams._topic_handle import ( TopicHandle, WorkflowTopicHandle, @@ -40,4 +43,5 @@ "WorkflowStreamItem", "WorkflowStreamState", "WorkflowTopicHandle", + "current_workflow_stream", ] diff --git a/temporalio/contrib/workflow_streams/_stream.py b/temporalio/contrib/workflow_streams/_stream.py index ae8608c3b..33a5346b9 100644 --- a/temporalio/contrib/workflow_streams/_stream.py +++ b/temporalio/contrib/workflow_streams/_stream.py @@ -59,6 +59,24 @@ T = TypeVar("T") +def current_workflow_stream() -> WorkflowStream | None: + """Return the current workflow's :py:class:`WorkflowStream`, if any. + + Must be called inside a workflow. A workflow that constructed a + ``WorkflowStream`` (in ``@workflow.init``) registers the stream's publish + signal handler; that registration is the marker this accessor reads. + Returns ``None`` when no stream was constructed. + + .. warning:: + This function is experimental and may change in future versions. + """ + handler = workflow.get_signal_handler(_PUBLISH_SIGNAL) + if handler is None: + return None + stream = getattr(handler, "__self__", None) + return stream if isinstance(stream, WorkflowStream) else None + + def _payload_wire_size(payload: Payload, topic: str) -> int: """Approximate poll-response contribution of a single item. diff --git a/tests/contrib/deepagents/test_failures.py b/tests/contrib/deepagents/test_failures.py index 6cfe6e5ff..21cf5de52 100644 --- a/tests/contrib/deepagents/test_failures.py +++ b/tests/contrib/deepagents/test_failures.py @@ -22,8 +22,10 @@ ) from temporalio import workflow from temporalio.client import WorkflowFailureError +from temporalio.contrib._langchain._activity_helpers import ( + translate_api_error as _translate_api_error, +) from temporalio.contrib.deepagents import DeepAgentsPlugin, DeepAgentsWorkflowError -from temporalio.contrib.deepagents._activity import _translate_api_error from temporalio.exceptions import ApplicationError from temporalio.worker import Worker diff --git a/tests/contrib/langchain_shared/__init__.py b/tests/contrib/langchain_shared/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/contrib/langchain_shared/test_aio_override.py b/tests/contrib/langchain_shared/test_aio_override.py new file mode 100644 index 000000000..42ac32d84 --- /dev/null +++ b/tests/contrib/langchain_shared/test_aio_override.py @@ -0,0 +1,70 @@ +"""The shared LangSmith ``aio_to_thread`` override installer.""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import pytest + +from temporalio.contrib._langchain import _aio_to_thread as mod + +pytest.importorskip("langsmith") + + +@pytest.fixture(autouse=True) +def reset_installed_flag(monkeypatch: pytest.MonkeyPatch): + """Isolate the module-level install-once flag per test.""" + monkeypatch.setattr(mod, "_installed", False) + yield + + +def test_installs_exactly_once(monkeypatch: pytest.MonkeyPatch) -> None: + import langsmith + + calls: list[Any] = [] + monkeypatch.setattr( + langsmith, "set_runtime_overrides", lambda **kw: calls.append(kw) + ) + mod.install_aio_to_thread_override() + mod.install_aio_to_thread_override() + mod.install_aio_to_thread_override() + assert len(calls) == 1 + assert calls[0] == {"aio_to_thread": mod._temporal_aio_to_thread} + + +def test_failure_leaves_flag_unset_for_retry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import langsmith + + boom = RuntimeError("boom") + + def _raise(**_kw: Any) -> None: + raise boom + + monkeypatch.setattr(langsmith, "set_runtime_overrides", _raise) + with pytest.raises(RuntimeError): + mod.install_aio_to_thread_override() + assert mod._installed is False + + # A later call retries and succeeds. + calls: list[Any] = [] + monkeypatch.setattr( + langsmith, "set_runtime_overrides", lambda **kw: calls.append(kw) + ) + mod.install_aio_to_thread_override() + assert len(calls) == 1 + assert mod._installed is True + + +def test_override_delegates_outside_workflows() -> None: + """Outside a workflow the default thread hop is used unchanged.""" + + async def default(_ctx: Any, func: Any, *args: Any, **kwargs: Any) -> Any: + return ("default", func(*args, **kwargs)) + + result = asyncio.run( + mod._temporal_aio_to_thread(default, object(), lambda x: x + 1, 41) + ) + assert result == ("default", 42) diff --git a/tests/contrib/langchain_shared/test_converter.py b/tests/contrib/langchain_shared/test_converter.py new file mode 100644 index 000000000..595735ded --- /dev/null +++ b/tests/contrib/langchain_shared/test_converter.py @@ -0,0 +1,59 @@ +"""The shared LangChain-family payload converter composition.""" + +from __future__ import annotations + +import dataclasses + +import pytest + +from temporalio.converter import DataConverter + +pytest.importorskip("pydantic") + +from temporalio.contrib._langchain._converter import ( # noqa: E402 + LangChainPayloadConverter, + build_data_converter, + data_converter, +) + + +def test_none_installs_family_default() -> None: + assert build_data_converter(None, plugin_name="XPlugin") is data_converter + + +def test_sdk_default_gets_converter_swapped() -> None: + result = build_data_converter(DataConverter.default, plugin_name="XPlugin") + assert result.payload_converter_class is LangChainPayloadConverter + # Everything else on the SDK default is preserved (dataclasses.replace + # rebuilds derived instance attributes, so compare the declared fields). + assert ( + result.failure_converter_class is DataConverter.default.failure_converter_class + ) + assert result.payload_codec is DataConverter.default.payload_codec + + +def test_custom_converter_is_refused_with_plugin_name() -> None: + custom = dataclasses.replace(DataConverter.default) + with pytest.raises(ValueError, match="XPlugin cannot compose"): + build_data_converter(custom, plugin_name="XPlugin") + + +def test_exclude_unset_round_trip() -> None: + from pydantic import BaseModel + + class Inner(BaseModel): + a: int = 1 + b: str | None = None + + class Outer(BaseModel): + inner: Inner + note: str | None = None + + converter = LangChainPayloadConverter() + payload = converter.to_payload(Outer(inner=Inner(a=2))) + assert payload is not None + # Unset defaults are excluded from the wire form. + assert b'"b"' not in payload.data + assert b'"note"' not in payload.data + restored = converter.from_payload(payload, Outer) + assert restored.inner.a == 2 and restored.note is None diff --git a/tests/contrib/langchain_shared/test_passthrough.py b/tests/contrib/langchain_shared/test_passthrough.py new file mode 100644 index 000000000..1844934a4 --- /dev/null +++ b/tests/contrib/langchain_shared/test_passthrough.py @@ -0,0 +1,15 @@ +"""The shared sandbox passthrough-list merge.""" + +from __future__ import annotations + +from temporalio.contrib._langchain._passthrough import merge_passthrough_modules + + +def test_merge_dedupes_preserving_first_occurrence_order() -> None: + merged = merge_passthrough_modules(("a", "b", "c"), ["b", "d", "a", "e"]) + assert merged == ("a", "b", "c", "d", "e") + + +def test_merge_tolerates_no_user_list() -> None: + assert merge_passthrough_modules(("a", "b"), None) == ("a", "b") + assert merge_passthrough_modules((), None) == () diff --git a/tests/contrib/langchain_shared/test_runnable_config.py b/tests/contrib/langchain_shared/test_runnable_config.py new file mode 100644 index 000000000..a58cd6275 --- /dev/null +++ b/tests/contrib/langchain_shared/test_runnable_config.py @@ -0,0 +1,195 @@ +"""The shared ``RunnableConfig`` strip/rebuild. + +The langgraph plugin's strip output is RELEASED behavior: stripped configs +ride in activity payloads and feed continue-as-new cache keys, so the shared +implementation must reproduce it byte-for-byte. ``_reference_strip_pre_refactor`` +below is the pre-refactor langgraph implementation embedded verbatim; the +corpus test string-compares JSON dumps (captures key order, not just dict +equality) and compares derived cache keys. +""" + +from __future__ import annotations + +import json +import uuid +from typing import Any + +import pytest + +from temporalio.contrib._langchain._runnable_config import ( + is_jsonish, + rebuild_runnable_config, + strip_runnable_config, +) +from temporalio.contrib._langchain._task_cache import cache_key + +pytest.importorskip("langgraph") + +from langgraph._internal._constants import ( # noqa: E402 # pyright: ignore[reportMissingTypeStubs] + CONFIG_KEY_CHECKPOINT_ID, + CONFIG_KEY_CHECKPOINT_MAP, + CONFIG_KEY_CHECKPOINT_NS, + CONFIG_KEY_DURABILITY, + CONFIG_KEY_RESUMING, + CONFIG_KEY_TASK_ID, + CONFIG_KEY_THREAD_ID, +) + +_KEPT = ( + CONFIG_KEY_CHECKPOINT_NS, + CONFIG_KEY_CHECKPOINT_ID, + CONFIG_KEY_CHECKPOINT_MAP, + CONFIG_KEY_THREAD_ID, + CONFIG_KEY_TASK_ID, + CONFIG_KEY_RESUMING, + CONFIG_KEY_DURABILITY, +) + + +def _reference_strip_pre_refactor(config: Any) -> dict[str, Any]: + """The langgraph plugin's strip_runnable_config as released (verbatim).""" + orig = config or {} + configurable = orig.get("configurable") or {} + + result: dict[str, Any] = { + "tags": list(orig.get("tags") or []), + "metadata": dict(orig.get("metadata") or {}), + } + if run_name := orig.get("run_name"): + result["run_name"] = run_name + if run_id := orig.get("run_id"): + result["run_id"] = run_id + if (recursion_limit := orig.get("recursion_limit")) is not None: + result["recursion_limit"] = recursion_limit + + stripped_configurable: dict[str, Any] = { + key: configurable[key] + for key in ( + CONFIG_KEY_CHECKPOINT_NS, + CONFIG_KEY_CHECKPOINT_ID, + CONFIG_KEY_CHECKPOINT_MAP, + CONFIG_KEY_THREAD_ID, + CONFIG_KEY_TASK_ID, + CONFIG_KEY_RESUMING, + CONFIG_KEY_DURABILITY, + ) + if key in configurable + } + if stripped_configurable: + result["configurable"] = stripped_configurable + return result + + +class _LiveJunk: + """Stand-in for callbacks / checkpointer handles / pregel callables.""" + + def __repr__(self) -> str: + return "" + + +def _corpus() -> list[Any]: + junk = _LiveJunk() + scrambled_configurable = { + "user_key": "kept-by-nothing", + CONFIG_KEY_DURABILITY: "sync", + "__pregel_send": junk, + CONFIG_KEY_THREAD_ID: "t-1", + CONFIG_KEY_CHECKPOINT_NS: "ns", + "another": {"nested": True}, + CONFIG_KEY_RESUMING: False, + CONFIG_KEY_CHECKPOINT_MAP: {"a": "b"}, + CONFIG_KEY_TASK_ID: "task-9", + CONFIG_KEY_CHECKPOINT_ID: "ckpt-3", + } + return [ + None, + {}, + {"tags": [], "metadata": {}}, + {"tags": ["a", "b"], "metadata": {"k": "v", "obj": junk}}, + {"run_name": ""}, # falsy run_name: dropped + {"run_name": "r", "run_id": uuid.uuid4()}, + {"recursion_limit": 0}, # zero is kept (not-None gate) + {"recursion_limit": None}, + { + "tags": ["x"], + "metadata": {"m": 1}, + "run_name": "full", + "run_id": "rid", + "recursion_limit": 25, + "callbacks": junk, + "configurable": scrambled_configurable, + }, + {"configurable": {k: f"v-{k}" for k in reversed(_KEPT)}}, + {"configurable": {"only": "dropped keys"}}, + ] + + +def test_langgraph_shape_is_byte_identical_to_released_behavior() -> None: + for cfg in _corpus(): + expected = _reference_strip_pre_refactor(cfg) + actual = strip_runnable_config(cfg, configurable_keys=_KEPT) + assert json.dumps(actual, default=str) == json.dumps(expected, default=str), cfg + # Cache keys derived from stripped configs must not shift either. + assert cache_key("t", (), {"config": actual}) == cache_key( + "t", (), {"config": expected} + ) + + +def test_langgraph_plugin_wrapper_matches_reference() -> None: + from temporalio.contrib.langgraph._langgraph_config import ( + strip_runnable_config as lg_strip, + ) + + for cfg in _corpus(): + assert json.dumps(lg_strip(cfg), default=str) == json.dumps( + _reference_strip_pre_refactor(cfg), default=str + ) + + +def test_filter_mode_matches_deepagents_semantics() -> None: + junk = _LiveJunk() + cfg = { + "tags": ["t"], + "metadata": {"ok": 1, "bad": junk}, + "run_id": "rid", + "configurable": {"keep": "v", "__dunder": "x", "unsafe": junk}, + } + out = strip_runnable_config( + cfg, + configurable_filter=lambda k, v: not k.startswith("__") and is_jsonish(v), + metadata_filter=is_jsonish, + ) + assert out == { + "tags": ["t"], + "metadata": {"ok": 1}, + "run_id": "rid", + "configurable": {"keep": "v"}, + } + + +def test_rebuild_round_trip() -> None: + stripped = { + "tags": ["t"], + "metadata": {"m": 1}, + "run_id": "rid", + "recursion_limit": 5, + "configurable": {"k": "v"}, + } + assert rebuild_runnable_config(dict(stripped)) == { + "metadata": {"m": 1}, + "tags": ["t"], + "run_id": "rid", + "recursion_limit": 5, + "configurable": {"k": "v"}, + } + # Empty input still yields the always-present metadata mapping. + assert rebuild_runnable_config({}) == {"metadata": {}} + + +def test_exactly_one_configurable_selector_required() -> None: + with pytest.raises(ValueError, match="exactly one"): + strip_runnable_config({}) + with pytest.raises(ValueError, match="exactly one"): + strip_runnable_config( + {}, configurable_keys=("a",), configurable_filter=lambda k, v: True + ) diff --git a/tests/contrib/langchain_shared/test_serde_messages.py b/tests/contrib/langchain_shared/test_serde_messages.py new file mode 100644 index 000000000..f47f12e7e --- /dev/null +++ b/tests/contrib/langchain_shared/test_serde_messages.py @@ -0,0 +1,58 @@ +"""Message and tool-schema serialization in the shared core.""" + +from __future__ import annotations + +import pytest + +from temporalio.contrib._langchain._messages import ( + dump_messages, + dump_object, + load_messages, + load_object, + tool_to_schema, +) + +pytest.importorskip("langchain_core") + +from langchain_core.messages import ( # noqa: E402 + AIMessage, + HumanMessage, + ToolMessage, +) + + +def test_object_round_trip_preserves_subtype_and_tool_calls() -> None: + msg = AIMessage( + content="", + tool_calls=[{"name": "t", "args": {"x": 1}, "id": "call-1"}], + ) + loaded = load_object(dump_object(msg)) + assert isinstance(loaded, AIMessage) + assert loaded.tool_calls[0]["id"] == "call-1" + assert loaded.tool_calls[0]["args"] == {"x": 1} + + +def test_messages_round_trip_preserves_order_and_types() -> None: + msgs = [ + HumanMessage(content="hi"), + AIMessage(content="hello"), + ToolMessage(content="result", tool_call_id="c1"), + ] + loaded = load_messages(dump_messages(msgs)) + assert [type(m) for m in loaded] == [HumanMessage, AIMessage, ToolMessage] + assert loaded[2].tool_call_id == "c1" + + +def test_tool_to_schema_carries_parameters() -> None: + from langchain_core.tools import tool + + @tool + def get_weather(city: str) -> str: + """Return the weather for a city.""" + return city + + schema = tool_to_schema(get_weather) + fn = schema["function"] + assert fn["name"] == "get_weather" + assert fn["description"] + assert "city" in fn["parameters"]["properties"] diff --git a/tests/contrib/langchain_shared/test_task_cache.py b/tests/contrib/langchain_shared/test_task_cache.py new file mode 100644 index 000000000..4638641ba --- /dev/null +++ b/tests/contrib/langchain_shared/test_task_cache.py @@ -0,0 +1,108 @@ +"""The shared continue-as-new result cache.""" + +from __future__ import annotations + +import functools +from contextvars import copy_context + +import pytest + +from temporalio.contrib._langchain._task_cache import ( + TaskResultCache, + cache_key, + task_id, +) + + +def test_instances_do_not_share_state() -> None: + a = TaskResultCache("test_cache_a") + b = TaskResultCache("test_cache_b") + a.set_cache({}) + a.put("k", "va") + assert a.lookup("k") == (True, "va") + # b has no active cache at all — a's state is invisible to it. + assert b.lookup("k") == (False, None) + b.set_cache({}) + assert b.lookup("k") == (False, None) + + +def test_isolation_across_copied_contexts() -> None: + cache = TaskResultCache("test_cache_ctx") + + def seed() -> None: + cache.set_cache({"k": "inner"}) + + # Setting inside a copied context must not leak to the outer context. + copy_context().run(seed) + assert cache.get_cache() is None + + +def test_set_cache_keeps_identity() -> None: + """The langgraph plugin exposes the live dict; the core must not copy.""" + cache = TaskResultCache("test_cache_identity") + backing: dict[str, object] = {} + cache.set_cache(backing) + assert cache.get_cache() is backing + cache.put("k", 1) + assert backing == {"k": 1} + + +def test_put_without_active_cache_is_a_noop() -> None: + cache = TaskResultCache("test_cache_noop") + cache.put("k", 1) + assert cache.lookup("k") == (False, None) + + +def test_task_id_rejects_unidentifiable_functions() -> None: + def outer(): + def inner() -> None: ... + + return inner + + with pytest.raises(ValueError, match="closures"): + task_id(outer()) + + class _MainStub: + __module__ = "__main__" + __qualname__ = "stub" + + with pytest.raises(ValueError, match="__main__"): + task_id(_MainStub()) + + # functools.partial instances carry neither __qualname__ nor __name__. + with pytest.raises(ValueError, match="module level"): + task_id(functools.partial(print)) + + +def test_task_id_happy_path() -> None: + assert task_id(test_task_id_happy_path) == f"{__name__}.test_task_id_happy_path" + + +def test_cache_key_stable_and_fallback() -> None: + k1 = cache_key("mod.fn", (1, "a"), {"b": 2}, context=None) + k2 = cache_key("mod.fn", (1, "a"), {"b": 2}, context=None) + assert k1 == k2 and len(k1) == 32 + # Different inputs, different key. + assert cache_key("mod.fn", (2, "a"), {"b": 2}) != k1 + + # Unserializable arguments fall back to repr and still produce a key. + class Weird: + def __repr__(self) -> str: + return "" + + k3 = cache_key("mod.fn", (Weird(),), {}) + assert len(k3) == 32 + + +def test_langgraph_delegation_preserves_identity_semantics() -> None: + pytest.importorskip("langgraph") + from temporalio.contrib.langgraph import _task_cache as lg + + backing: dict[str, object] = {"seed": 1} + lg.set_task_cache(backing) + assert lg.get_task_cache() is backing + lg.cache_put("k", "v") + assert lg.cache_lookup("k") == (True, "v") + assert backing["k"] == "v" + lg.set_task_cache(None) + assert lg.get_task_cache() is None diff --git a/tests/contrib/workflow_streams/test_workflow_streams.py b/tests/contrib/workflow_streams/test_workflow_streams.py index 12026ff1a..da9c1d173 100644 --- a/tests/contrib/workflow_streams/test_workflow_streams.py +++ b/tests/contrib/workflow_streams/test_workflow_streams.py @@ -45,6 +45,7 @@ WorkflowStreamItem, WorkflowStreamState, WorkflowTopicHandle, + current_workflow_stream, ) from temporalio.contrib.workflow_streams._types import _encode_payload from temporalio.converter import DataConverter @@ -2788,3 +2789,38 @@ async def broker_started() -> bool: await broker_handle.signal("close") result = await caller_handle.result() assert result == "done" + + +@workflow.defn +class StreamAccessorWorkflow: + @workflow.init + def __init__(self) -> None: + self.stream = WorkflowStream() + + @workflow.run + async def run(self) -> bool: + return current_workflow_stream() is self.stream + + +@workflow.defn +class NoStreamAccessorWorkflow: + @workflow.run + async def run(self) -> bool: + return current_workflow_stream() is None + + +async def test_current_workflow_stream_accessor(client: Client) -> None: + """The accessor returns the constructed stream instance, else None.""" + async with new_worker( + client, StreamAccessorWorkflow, NoStreamAccessorWorkflow + ) as worker: + assert await client.execute_workflow( + StreamAccessorWorkflow.run, + id=f"stream-accessor-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + assert await client.execute_workflow( + NoStreamAccessorWorkflow.run, + id=f"stream-accessor-none-{uuid.uuid4()}", + task_queue=worker.task_queue, + )