138 lines
4.5 KiB
Python
138 lines
4.5 KiB
Python
"""Runtime overrides for LangSmith.
|
|
|
|
This module provides hooks to override LangSmith's default runtime behavior,
|
|
primarily for environments with constrained async runtimes (e.g., Temporal).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Any, Callable, Optional
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Awaitable
|
|
|
|
|
|
AioToThread = Callable[
|
|
..., # (ctx: contextvars.Context, func, /, *args, **kwargs)
|
|
"Awaitable[Any]",
|
|
]
|
|
|
|
|
|
class RuntimeOverrides:
|
|
"""Overrides for LangSmith runtime behavior.
|
|
|
|
This class allows overriding default async implementations for environments
|
|
that don't support certain asyncio features (e.g., Temporal doesn't support
|
|
``run_in_executor``).
|
|
|
|
Example:
|
|
import langsmith
|
|
import contextvars
|
|
|
|
|
|
async def my_aio_to_thread(
|
|
default_aio_to_thread, ctx, func, /, *args, **kwargs
|
|
):
|
|
# Custom implementation
|
|
return ctx.run(func, *args, **kwargs)
|
|
|
|
|
|
langsmith.set_runtime_overrides(aio_to_thread=my_aio_to_thread)
|
|
|
|
# Reset to defaults
|
|
langsmith.set_runtime_overrides()
|
|
"""
|
|
|
|
__slots__ = ("aio_to_thread",)
|
|
|
|
def __init__(
|
|
self,
|
|
aio_to_thread: Optional[AioToThread] = None,
|
|
):
|
|
"""Initialize runtime overrides.
|
|
|
|
Args:
|
|
aio_to_thread: Custom async-to-thread implementation, with signature
|
|
``async def (default_aio_to_thread, ctx, func, /, *args, **kwargs)``.
|
|
``default_aio_to_thread`` is LangSmith's default implementation, which
|
|
the override can call to fall back to default behavior (e.g., when
|
|
outside a constrained runtime context). ``ctx`` is the
|
|
``contextvars.Context`` LangSmith wants ``func`` to run inside;
|
|
tracing state will be read back from this Context after the call.
|
|
Override for runtimes like Temporal that don't support
|
|
``asyncio.run_in_executor``.
|
|
"""
|
|
self.aio_to_thread = aio_to_thread
|
|
|
|
|
|
_runtime_overrides = RuntimeOverrides()
|
|
|
|
|
|
def set_runtime_overrides(
|
|
aio_to_thread: Optional[AioToThread] = None,
|
|
) -> None:
|
|
"""Set LangSmith runtime overrides.
|
|
|
|
This allows customizing LangSmith's async runtime behavior for environments
|
|
with constrained async runtimes (e.g., Temporal, which doesn't support
|
|
``run_in_executor``).
|
|
|
|
Args:
|
|
aio_to_thread: Custom async function to run sync functions
|
|
asynchronously. Should have signature:
|
|
``async def aio_to_thread(
|
|
default_aio_to_thread, ctx, func, /, *args, **kwargs
|
|
)``.
|
|
``default_aio_to_thread`` is LangSmith's default implementation, which
|
|
the override can call to fall back to default behavior. The
|
|
implementation must invoke ``func`` inside ``ctx`` (e.g.
|
|
``ctx.run(func, *args, **kwargs)``) so that LangSmith's tracing state,
|
|
which is read back from ``ctx`` after the call, is visible to downstream
|
|
code. Pass ``None`` to use the default implementation.
|
|
|
|
Example:
|
|
For Temporal or similar runtimes:
|
|
|
|
```python
|
|
import langsmith
|
|
|
|
|
|
async def temporal_aio_to_thread(
|
|
default_aio_to_thread, ctx, func, /, *args, **kwargs
|
|
):
|
|
# Use the default implementation when not in a workflow
|
|
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)
|
|
|
|
|
|
langsmith.set_runtime_overrides(aio_to_thread=temporal_aio_to_thread)
|
|
```
|
|
|
|
Reset to defaults:
|
|
|
|
```python
|
|
langsmith.set_runtime_overrides()
|
|
```
|
|
"""
|
|
global _runtime_overrides
|
|
_runtime_overrides = RuntimeOverrides(aio_to_thread=aio_to_thread)
|
|
|
|
|
|
def get_runtime_overrides() -> RuntimeOverrides:
|
|
"""Get the current runtime overrides."""
|
|
return _runtime_overrides
|
|
|
|
|
|
def _aio_to_thread_override_active() -> bool:
|
|
"""Return True iff an ``aio_to_thread`` override is currently installed.
|
|
|
|
Callers use this to select a loop-behavior-independent path for
|
|
re-entering a LangSmith-mutated Context (the explicit
|
|
``tracing_context(**get_tracing_context(ctx))`` fallback), rather than
|
|
relying on ``asyncio.create_task(coro, context=ctx)`` which some custom
|
|
event loops silently ignore.
|
|
"""
|
|
return _runtime_overrides.aio_to_thread is not None
|