diff --git a/sdk/python/adrian/anthropic_handler.py b/sdk/python/adrian/anthropic_handler.py index a405a3f..f2f2194 100644 --- a/sdk/python/adrian/anthropic_handler.py +++ b/sdk/python/adrian/anthropic_handler.py @@ -77,7 +77,7 @@ from uuid import uuid4 from adrian.config import AdrianConfig -from adrian.context import get_invocation_id, set_invocation_id +from adrian.context import get_invocation_id, in_chat_model, set_invocation_id from adrian.format.types import AgentContext, LlmPairData, PairedEvent from adrian.hooks import HookRegistry from adrian.types import ChatMessage, EventData, TokenUsage, ToolCallRecord @@ -1055,6 +1055,14 @@ def _patched_sync_create( **kwargs: Any, # noqa: ANN401 ) -> Any: # noqa: ANN401 response = _original_sync(self, *args, **kwargs) + + # ChatAnthropic routes through here, but the LangChain + # callbacks already emit and gate this call. Injecting + # thinking above still applies, so the LangChain event + # keeps its reasoning. + if in_chat_model(): + return response + # Emit + gate together on the WS loop (BLOCK/HITL); degrades to # audit-only emission when gating isn't possible. return _emit_and_gate_sync(response, kwargs) @@ -1069,10 +1077,14 @@ def _patched_sync_stream( # Sampled now, not at consumption time: the caller may read the # final message after leaving the anthropic_invocation() block. captured = get_invocation_id() + inner = _original_sync_stream(self, *args, **kwargs) - return _GatedMessageStreamManager( - _original_sync_stream(self, *args, **kwargs), kwargs, captured - ) + # Owned by the LangChain callbacks: hand back the SDK's own + # manager so nothing is emitted or gated twice. + if in_chat_model(): + return inner + + return _GatedMessageStreamManager(inner, kwargs, captured) sync_cls.create = _patched_sync_create # type: ignore[method-assign] sync_cls.stream = _patched_sync_stream # type: ignore[method-assign] @@ -1097,6 +1109,12 @@ async def _patched_async_create( **kwargs: Any, # noqa: ANN401 ) -> Any: # noqa: ANN401 response = await _original_async(self, *args, **kwargs) + + # See the sync wrapper: the LangChain callbacks own this + # call, so emitting here would double-count it. + if in_chat_model(): + return response + # Emit first so the verdict future is registered, then gate: # under BLOCK/HITL this holds the response until the verdict # arrives and rewrites blocked tool calls to a [BLOCKED] block. @@ -1113,10 +1131,13 @@ def _patched_async_stream( # Not async: stream() returns an async context manager without # being awaited itself. captured = get_invocation_id() + inner = _original_async_stream(self, *args, **kwargs) - return _GatedAsyncMessageStreamManager( - _original_async_stream(self, *args, **kwargs), kwargs, captured - ) + # See the sync stream wrapper. + if in_chat_model(): + return inner + + return _GatedAsyncMessageStreamManager(inner, kwargs, captured) async_cls.create = _patched_async_create # type: ignore[method-assign] async_cls.stream = _patched_async_stream # type: ignore[method-assign] diff --git a/sdk/python/adrian/context.py b/sdk/python/adrian/context.py index ec8dc77..7f4484c 100644 --- a/sdk/python/adrian/context.py +++ b/sdk/python/adrian/context.py @@ -19,6 +19,8 @@ from __future__ import annotations +from collections.abc import Generator +from contextlib import contextmanager from contextvars import ContextVar, Token from adrian.format.types import AgentContext, ParentContext @@ -28,6 +30,11 @@ default=None, ) +_in_chat_model: ContextVar[bool] = ContextVar( + "adrian_in_chat_model", + default=False, +) + def get_invocation_id() -> str | None: """Get the current invocation ID from async context. @@ -54,6 +61,35 @@ def set_invocation_id(invocation_id: str) -> Token[str | None]: return _invocation_id.set(invocation_id) +@contextmanager +def chat_model_scope() -> Generator[None]: + """Mark the calling frame as a LangChain chat model invocation. + + ``ChatAnthropic`` reaches Anthropic through the same SDK entry points + Adrian patches, so without this both the LangChain callbacks and the + Anthropic patch would describe one call as two events. The patched + ``BaseChatModel`` methods hold this scope for the duration of the + request; :func:`in_chat_model` lets the Anthropic patch see it and + stand down. + """ + token = _in_chat_model.set(True) + + try: + yield + finally: + _in_chat_model.reset(token) + + +def in_chat_model() -> bool: + """Report whether a LangChain chat model call is on the stack. + + Returns: + ``True`` when the LangChain callbacks already own this call, so + provider-level instrumentation should not emit its own event. + """ + return _in_chat_model.get() + + class AgentContextTracker: """Tracks agent prompt context for parent-child enrichment. diff --git a/sdk/python/adrian/langchain_handler.py b/sdk/python/adrian/langchain_handler.py index 4b3c977..4596a20 100644 --- a/sdk/python/adrian/langchain_handler.py +++ b/sdk/python/adrian/langchain_handler.py @@ -39,7 +39,7 @@ from langchain_core.runnables.base import Runnable from langchain_core.runnables.config import ensure_config -from adrian.context import get_invocation_id, set_invocation_id +from adrian.context import chat_model_scope, get_invocation_id, set_invocation_id from adrian.ws import should_halt if TYPE_CHECKING: @@ -262,6 +262,9 @@ def _patch_chat_model() -> None: original_astream = BaseChatModel.astream original_stream = BaseChatModel.stream + # chat_model_scope suppresses the provider-level patch for the + # duration of the request. ChatAnthropic calls the same Anthropic SDK + # methods Adrian patches, so without it one call yields two events. def patched_invoke( self: Any, # noqa: ANN401 input: Any, # noqa: A002, ANN401 @@ -269,7 +272,8 @@ def patched_invoke( **kwargs: Any, ) -> Any: # noqa: ANN401 config = _inject_callbacks(config) - return original_invoke(self, input, config=config, **kwargs) + with chat_model_scope(): + return original_invoke(self, input, config=config, **kwargs) async def patched_ainvoke( self: Any, # noqa: ANN401 @@ -278,7 +282,8 @@ async def patched_ainvoke( **kwargs: Any, ) -> Any: # noqa: ANN401 config = _inject_callbacks(config) - return await original_ainvoke(self, input, config=config, **kwargs) + with chat_model_scope(): + return await original_ainvoke(self, input, config=config, **kwargs) async def patched_astream( self: Any, # noqa: ANN401 @@ -287,8 +292,9 @@ async def patched_astream( **kwargs: Any, ) -> Any: # noqa: ANN401 config = _inject_callbacks(config) - async for chunk in original_astream(self, input, config=config, **kwargs): - yield chunk + with chat_model_scope(): + async for chunk in original_astream(self, input, config=config, **kwargs): + yield chunk def patched_stream( self: Any, # noqa: ANN401 @@ -297,7 +303,8 @@ def patched_stream( **kwargs: Any, ) -> Any: # noqa: ANN401 config = _inject_callbacks(config) - yield from original_stream(self, input, config=config, **kwargs) + with chat_model_scope(): + yield from original_stream(self, input, config=config, **kwargs) BaseChatModel.invoke = patched_invoke # type: ignore[assignment] BaseChatModel.ainvoke = patched_ainvoke # type: ignore[assignment] diff --git a/sdk/python/tests/test_anthropic_handler.py b/sdk/python/tests/test_anthropic_handler.py index eb926de..ff2848f 100644 --- a/sdk/python/tests/test_anthropic_handler.py +++ b/sdk/python/tests/test_anthropic_handler.py @@ -38,7 +38,12 @@ patch_anthropic, ) from adrian.config import AdrianConfig -from adrian.context import get_invocation_id, set_invocation_id +from adrian.context import ( + chat_model_scope, + get_invocation_id, + in_chat_model, + set_invocation_id, +) from adrian.format.types import LlmPairData, PairedEvent from adrian.hooks import HookRegistry from adrian.proto import event_pb2 as pb @@ -1868,3 +1873,104 @@ def test_patched_stream_captures_invocation_id_eagerly(self) -> None: # Sampled inside the block, so it survives the context exit. assert get_invocation_id() is None assert manager._state._invocation_id == expected # pyright: ignore[reportPrivateUsage] + + +# ------------------------------------------------------------------ +# LangChain co-instrumentation +# ------------------------------------------------------------------ + + +class TestChatModelScopeSuppression: + """ChatAnthropic drives the same SDK methods Adrian patches. + + Without suppression one LangChain call produces two LLM events: one + from the callbacks and one from the Anthropic patch. + """ + + @pytest.fixture(autouse=True) # pyright: ignore[reportUntypedFunctionDecorator] + def _restore_sdk(self) -> Any: # noqa: ANN401 + from anthropic.resources.messages import AsyncMessages, Messages + + saved = [ + (Messages, "create", Messages.create), + (Messages, "stream", Messages.stream), + (AsyncMessages, "create", AsyncMessages.create), + (AsyncMessages, "stream", AsyncMessages.stream), + ] + flags = [ + (cls, getattr(cls, "_adrian_patched", None)) + for cls in (Messages, AsyncMessages) + ] + + for cls, _flag in flags: + if hasattr(cls, "_adrian_patched"): + delattr(cls, "_adrian_patched") + + yield + + for cls, name, original in saved: + setattr(cls, name, original) + + for cls, flag in flags: + if hasattr(cls, "_adrian_patched"): + delattr(cls, "_adrian_patched") + if flag is not None: + cls._adrian_patched = flag # type: ignore[attr-defined] + + _ah._hooks_getter = None + _ah._config_getter = None + _ah._ws_getter = None + _ah._handler_getter = None + + def test_no_event_emitted_inside_scope(self) -> None: + from anthropic.resources.messages import Messages + + response = _make_text_response() + Messages.create = lambda _self, **_kw: response # type: ignore[method-assign, assignment] + hooks, collector = _wired_hooks(AdrianConfig(session_id="s")) + patch_anthropic(lambda: hooks, lambda: AdrianConfig(session_id="s")) + + with chat_model_scope(): + returned = Messages.create(MagicMock(), **_KWARGS) + + assert returned is response + assert collector.events == [] + + def test_event_emitted_outside_scope(self) -> None: + from anthropic.resources.messages import Messages + + Messages.create = lambda _self, **_kw: _make_text_response() # type: ignore[method-assign, assignment] + hooks, collector = _wired_hooks(AdrianConfig(session_id="s")) + patch_anthropic(lambda: hooks, lambda: AdrianConfig(session_id="s")) + + Messages.create(MagicMock(), **_KWARGS) + + assert len(collector.events) == 1 + + def test_stream_returns_sdk_manager_inside_scope(self) -> None: + from anthropic.resources.messages import Messages + + inner = _FakeStreamManager(_FakeMessageStream(_make_text_response())) + Messages.stream = lambda _self, **_kw: inner # type: ignore[method-assign, assignment] + hooks, _ = _wired_hooks(AdrianConfig(session_id="s")) + patch_anthropic(lambda: hooks, lambda: AdrianConfig(session_id="s")) + + with chat_model_scope(): + manager = Messages.stream(MagicMock(), **_KWARGS) + + assert manager is inner + assert not isinstance(manager, _GatedMessageStreamManager) + + def test_scope_resets_on_exit(self) -> None: + assert in_chat_model() is False + + with chat_model_scope(): + assert in_chat_model() is True + + assert in_chat_model() is False + + def test_scope_resets_on_exception(self) -> None: + with pytest.raises(RuntimeError), chat_model_scope(): + raise RuntimeError("boom") + + assert in_chat_model() is False