How to use ChatAnthropic with azure_ad_token_provider?

hi @Alisca-Chen

can you try out this?

"""ChatAnthropic subclass with Azure AI Foundry + Entra ID (Azure AD) support.

Problem
-------
``ChatAnthropic`` from ``langchain-anthropic`` creates ``anthropic.Client``
internally and only supports ``api_key`` authentication.  In production
environments that require Azure Entra ID (Azure AD) token-based auth to
access Claude models deployed in Azure AI Foundry, there is no native
``azure_ad_token_provider`` parameter.

Solution
--------
Override the two ``@cached_property`` methods that build the underlying
Anthropic clients (``_client`` and ``_async_client``) and return
``AnthropicFoundry`` / ``AsyncAnthropicFoundry`` instead of the vanilla
``anthropic.Client`` / ``anthropic.AsyncClient``.

``AnthropicFoundry`` (from ``anthropic[foundry]``) handles:
  - Azure-specific endpoint URL construction
  - ``Authorization: Bearer <token>`` headers (not ``x-api-key``)
  - Calling ``azure_ad_token_provider`` on **every request** for auto-refresh

Requirements
------------
.. code-block:: bash

    pip install langchain-anthropic "anthropic[foundry]" azure-identity

References
----------
- AnthropicFoundry source: https://github.com/anthropics/anthropic-sdk-python/blob/main/src/anthropic/lib/foundry.py
- ChatAnthropic source: https://github.com/langchain-ai/langchain/blob/master/libs/partners/anthropic/langchain_anthropic/chat_models.py
- Official Foundry auth docs: https://docs.claude.com/en/docs/build-with-claude/claude-in-microsoft-foundry
"""

from __future__ import annotations

from collections.abc import Callable
from functools import cached_property
from typing import Any

from anthropic import AnthropicFoundry, AsyncAnthropicFoundry
from langchain_anthropic import ChatAnthropic
from pydantic import Field, SecretStr


class AzureFoundryChatAnthropic(ChatAnthropic):
    """``ChatAnthropic`` that authenticates via Azure Entra ID against Azure AI Foundry.

    All standard ``ChatAnthropic`` features work unchanged — tool calling,
    structured output, streaming, extended thinking, MCP, etc. — because
    ``AnthropicFoundry`` is a drop-in subclass of ``anthropic.Anthropic``
    with the same ``messages.create(...)`` interface.

    Parameters
    ----------
    azure_ad_token_provider : Callable[[], str]
        A **synchronous** callable that returns a bearer token string.
        Typically created with ``azure.identity.get_bearer_token_provider``.
        Called automatically on every API request, so token refresh is handled
        for you.
    azure_resource : str, optional
        Your Foundry resource name (e.g. ``"my-resource"``).  The SDK builds
        the base URL as ``https://{resource}.services.ai.azure.com/anthropic/``.
        Mutually exclusive with ``base_url``.
    model : str
        The deployment name in Foundry (e.g. ``"claude-sonnet-4-5"``).
    api_key : str
        Pass any non-empty placeholder (e.g. ``"unused"``).  Required by
        ``ChatAnthropic``'s Pydantic validation but **not sent** to Azure;
        ``AnthropicFoundry`` uses the token provider instead.

    Examples
    --------
    Basic usage with ``get_bearer_token_provider``::

        from azure.identity import DefaultAzureCredential, get_bearer_token_provider

        token_provider = get_bearer_token_provider(
            DefaultAzureCredential(),
            "https://cognitiveservices.azure.com/.default",
        )

        llm = AzureFoundryChatAnthropic(
            model="claude-sonnet-4-5",
            azure_ad_token_provider=token_provider,
            azure_resource="my-foundry-resource",
            api_key="unused",            # placeholder — not sent to Azure
        )
        response = llm.invoke("Hello from Azure Foundry!")
        print(response.content)

    With LangGraph ``create_agent``::

        from langgraph.prebuilt import create_agent

        agent = create_agent(
            model=llm,
            tools=[my_tool],
        )
        result = agent.invoke({"messages": [("user", "Do something")]})
    """

    # ------------------------------------------------------------------ fields
    azure_ad_token_provider: Callable[[], str] | None = Field(
        default=None,
        description=(
            "Synchronous callable returning an Azure AD bearer token. "
            "Created via azure.identity.get_bearer_token_provider()."
        ),
    )

    azure_ad_token_provider_async: Callable[..., Any] | None = Field(
        default=None,
        description=(
            "Optional async-compatible token provider for the async client. "
            "If not provided, azure_ad_token_provider is used for both."
        ),
    )

    azure_resource: str | None = Field(
        default=None,
        description=(
            "Azure AI Foundry resource name. "
            "Builds URL: https://{resource}.services.ai.azure.com/anthropic/. "
            "Mutually exclusive with base_url."
        ),
    )

    # Pydantic config — allow arbitrary callable types
    model_config = {**ChatAnthropic.model_config, "arbitrary_types_allowed": True}

    # -------------------------------------------------------------- clients

    @cached_property
    def _client(self) -> AnthropicFoundry:  # type: ignore[override]
        """Build a synchronous ``AnthropicFoundry`` client.

        Uses ``azure_ad_token_provider`` for auth.  The provider is called on
        every request by ``AnthropicFoundry._prepare_options``, so token
        refresh is automatic.
        """
        kwargs: dict[str, Any] = {
            "azure_ad_token_provider": self.azure_ad_token_provider,
            "max_retries": self.max_retries,
            "default_headers": self.default_headers or None,
        }

        # Timeout — None is meaningful to Anthropic client (infinite),
        # so only skip if explicitly set to <= 0.
        if self.default_request_timeout is None or self.default_request_timeout > 0:
            kwargs["timeout"] = self.default_request_timeout

        # Resource vs base_url (mutually exclusive in AnthropicFoundry)
        if self.azure_resource:
            kwargs["resource"] = self.azure_resource
        elif self.anthropic_api_url and self.anthropic_api_url != "https://api.anthropic.com":
            kwargs["base_url"] = self.anthropic_api_url

        return AnthropicFoundry(**kwargs)

    @cached_property
    def _async_client(self) -> AsyncAnthropicFoundry:  # type: ignore[override]
        """Build an asynchronous ``AsyncAnthropicFoundry`` client.

        Uses ``azure_ad_token_provider_async`` if provided, otherwise falls
        back to ``azure_ad_token_provider`` (``AsyncAnthropicFoundry``
        accepts both sync and async callables).
        """
        provider = self.azure_ad_token_provider_async or self.azure_ad_token_provider

        kwargs: dict[str, Any] = {
            "azure_ad_token_provider": provider,
            "max_retries": self.max_retries,
            "default_headers": self.default_headers or None,
        }

        if self.default_request_timeout is None or self.default_request_timeout > 0:
            kwargs["timeout"] = self.default_request_timeout

        if self.azure_resource:
            kwargs["resource"] = self.azure_resource
        elif self.anthropic_api_url and self.anthropic_api_url != "https://api.anthropic.com":
            kwargs["base_url"] = self.anthropic_api_url

        return AsyncAnthropicFoundry(**kwargs)


# ======================================================================
# Standalone demo
# ======================================================================
if __name__ == "__main__":
    # ----- Prerequisites -----
    # pip install langchain-anthropic "anthropic[foundry]" azure-identity
    #
    # You need:
    #   1. An Azure AI Foundry resource with a Claude deployment
    #   2. Azure credentials configured (az login, managed identity, etc.)
    # -------------------------

    from azure.identity import DefaultAzureCredential, get_bearer_token_provider

    # The scope for Azure AI Foundry / Cognitive Services
    token_provider = get_bearer_token_provider(
        DefaultAzureCredential(),
        "https://cognitiveservices.azure.com/.default",
    )

    llm = AzureFoundryChatAnthropic(
        model="claude-sonnet-4-5",            # your Foundry deployment name
        azure_ad_token_provider=token_provider,
        azure_resource="your-resource-name",  # <-- replace with your resource
        api_key="unused",                     # placeholder for Pydantic validation
    )

    # --- Simple invocation ---
    response = llm.invoke("What is Azure AI Foundry?")
    print("Response:", response.content)

    # --- Streaming ---
    print("\nStreaming:")
    for chunk in llm.stream("Explain token-based auth in one paragraph."):
        print(chunk.content, end="", flush=True)
    print()

    # --- Tool calling ---
    from pydantic import BaseModel

    class GetWeather(BaseModel):
        """Get the current weather in a given location."""
        location: str

    llm_with_tools = llm.bind_tools([GetWeather])
    result = llm_with_tools.invoke("What's the weather in Paris?")
    print("\nTool calls:", result.tool_calls)

    # --- Structured output ---
    class CityInfo(BaseModel):
        """Information about a city."""
        name: str
        country: str
        population: int

    structured_llm = llm.with_structured_output(CityInfo)
    city = structured_llm.invoke("Tell me about Tokyo")
    print("\nStructured:", city)