feat(vertex): add support for google vertex (#319)

examples/vertex.py

@@ -0,0 +1,43 @@
+import asyncio
+
+from anthropic import AnthropicVertex, AsyncAnthropicVertex
+
+
+def sync_client() -> None:
+    print("------ Sync Vertex ------")
+
+    client = AnthropicVertex()
+
+    message = client.beta.messages.create(
+        model="claude-instant-1p2",
+        max_tokens=100,
+        messages=[
+            {
+                "role": "user",
+                "content": "Say hello there!",
+            }
+        ],
+    )
+    print(message.model_dump_json(indent=2))
+
+
+async def async_client() -> None:
+    print("------ Async Vertex ------")
+
+    client = AsyncAnthropicVertex()
+
+    message = await client.beta.messages.create(
+        model="claude-instant-1p2",
+        max_tokens=1024,
+        messages=[
+            {
+                "role": "user",
+                "content": "Say hello there!",
+            }
+        ],
+    )
+    print(message.model_dump_json(indent=2))
+
+
+sync_client()
+asyncio.run(async_client())

pyproject.toml

@@ -36,7 +36,8 @@ classifiers = [
   "License :: OSI Approved :: MIT License"
 ]
 
-
+[project.optional-dependencies]
+vertex = ["google-auth >=2, <3"]
 
 [project.urls]
 Homepage = "https://github.com/anthropics/anthropic-sdk-python"

requirements-dev.lock

@@ -11,6 +11,7 @@ annotated-types==0.6.0
 anyio==4.1.0
 argcomplete==3.1.2
 attrs==23.1.0
+cachetools==5.3.2
 certifi==2023.7.22
 charset-normalizer==3.3.2
 colorlog==6.7.0
@@ -20,6 +21,7 @@ distro==1.8.0
 exceptiongroup==1.1.3
 filelock==3.12.4
 fsspec==2023.12.2
+google-auth==2.26.2
 h11==0.14.0
 httpcore==1.0.2
 httpx==0.25.2
@@ -35,6 +37,8 @@ packaging==23.2
 platformdirs==3.11.0
 pluggy==1.3.0
 py==1.11.0
+pyasn1==0.5.1
+pyasn1-modules==0.3.0
 pydantic==2.4.2
 pydantic-core==2.10.1
 pyright==1.1.332
@@ -45,6 +49,7 @@ pytz==2023.3.post1
 pyyaml==6.0.1
 requests==2.31.0
 respx==0.20.2
+rsa==4.9
 ruff==0.1.9
 six==1.16.0
 sniffio==1.3.0

requirements.lock

@@ -9,22 +9,27 @@
 -e file:.
 annotated-types==0.6.0
 anyio==4.1.0
+cachetools==5.3.2
 certifi==2023.7.22
 charset-normalizer==3.3.2
 distro==1.8.0
 exceptiongroup==1.1.3
 filelock==3.13.1
 fsspec==2023.12.0
+google-auth==2.26.2
 h11==0.14.0
 httpcore==1.0.2
 httpx==0.25.2
 huggingface-hub==0.16.4
 idna==3.4
 packaging==23.2
+pyasn1==0.5.1
+pyasn1-modules==0.3.0
 pydantic==2.4.2
 pydantic-core==2.10.1
 pyyaml==6.0.1
 requests==2.31.0
+rsa==4.9
 sniffio==1.3.0
 tokenizers==0.14.0
 tqdm==4.66.1

src/anthropic/__init__.py

@@ -69,6 +69,7 @@
     "AI_PROMPT",
 ]
 
+from .lib.vertex import *
 from .lib.streaming import *
 
 _setup_logging()

src/anthropic/lib/_extras/__init__.py

@@ -0,0 +1 @@
+from ._google_auth import google_auth as google_auth

src/anthropic/lib/_extras/_common.py

@@ -0,0 +1,13 @@
+from ..._exceptions import AnthropicError
+
+INSTRUCTIONS = """
+
+Anthropic error: missing required dependency `{library}`.
+
+    $ pip install anthropic[{extra}]
+"""
+
+
+class MissingDependencyError(AnthropicError):
+    def __init__(self, *, library: str, extra: str) -> None:
+        super().__init__(INSTRUCTIONS.format(library=library, extra=extra))

src/anthropic/lib/_extras/_google_auth.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+from typing_extensions import ClassVar, override
+
+from ._common import MissingDependencyError
+from ..._utils import LazyProxy
+
+if TYPE_CHECKING:
+    import google.auth  # type: ignore
+
+    google_auth = google.auth
+
+
+class GoogleAuthProxy(LazyProxy[Any]):
+    should_cache: ClassVar[bool] = True
+
+    @override
+    def __load__(self) -> Any:
+        try:
+            import google.auth  # type: ignore
+        except ImportError as err:
+            raise MissingDependencyError(extra="vertex", library="google-auth") from err
+
+        return google.auth
+
+
+if not TYPE_CHECKING:
+    google_auth = GoogleAuthProxy()

src/anthropic/lib/vertex/__init__.py

@@ -0,0 +1 @@
+from ._client import AnthropicVertex as AnthropicVertex, AsyncAnthropicVertex as AsyncAnthropicVertex

src/anthropic/lib/vertex/_auth.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from .._extras import google_auth
+
+if TYPE_CHECKING:
+    from google.auth.credentials import Credentials  # type: ignore[import-untyped]
+
+# pyright: reportMissingTypeStubs=false, reportUnknownVariableType=false, reportUnknownMemberType=false, reportUnknownArgumentType=false
+# google libraries don't provide types :/
+
+# Note: these functions are blocking as they make HTTP requests, the async
+# client runs these functions in a separate thread to ensure they do not
+# cause synchronous blocking issues.
+
+
+def load_auth() -> tuple[Credentials, str]:
+    from google.auth.transport.requests import Request  # type: ignore[import-untyped]
+
+    credentials, project_id = google_auth.default()
+    credentials.refresh(Request())
+
+    if not project_id:
+        raise ValueError("Could not resolve project_id")
+
+    if not isinstance(project_id, str):
+        raise TypeError(f"Expected project_id to be a str but got {type(project_id)}")
+
+    return credentials, project_id
+
+
+def refresh_auth(credentials: Credentials) -> None:
+    from google.auth.transport.requests import Request  # type: ignore[import-untyped]
+
+    credentials.refresh(Request())