diff --git a/pyproject.toml b/pyproject.toml index 729d42b2..83d27ee4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,9 +27,11 @@ dependencies = [ "websockets>=15.0", "asgiref>=3.11.0", "filetype>=1.2.0", - "google-auth[requests]>=2.0.0", ] +[project.optional-dependencies] +gcp = ["google-auth[requests]>=2.0.0"] + [project.urls] Homepage = "https://withceleste.ai" Documentation = "https://withceleste.ai/docs" diff --git a/src/celeste/exceptions.py b/src/celeste/exceptions.py index 08f6cb5f..f2812745 100644 --- a/src/celeste/exceptions.py +++ b/src/celeste/exceptions.py @@ -170,6 +170,16 @@ def __init__(self) -> None: super().__init__("Stream completed but no chunks were produced") +class MissingDependencyError(Error): + """Raised when a required optional dependency is not installed.""" + + def __init__(self, *, library: str, extra: str) -> None: + super().__init__( + f"Missing required dependency `{library}`. " + f'Install it with: pip install "celeste-ai[{extra}]"' + ) + + class CredentialsError(Error): """Errors related to API credentials.""" @@ -217,6 +227,7 @@ def __init__(self, parameter: str, model_id: str) -> None: "ConstraintViolationError", "Error", "MissingCredentialsError", + "MissingDependencyError", "ModalityNotFoundError", "ModelNotFoundError", "StreamEmptyError", diff --git a/src/celeste/providers/google/auth.py b/src/celeste/providers/google/auth.py index 75d7253e..24eba8d8 100644 --- a/src/celeste/providers/google/auth.py +++ b/src/celeste/providers/google/auth.py @@ -5,6 +5,10 @@ from pydantic import ConfigDict from celeste.auth import Authentication +from celeste.exceptions import MissingDependencyError + +VERTEX_BASE_URL = "https://{location}-aiplatform.googleapis.com" +VERTEX_GLOBAL_BASE_URL = "https://aiplatform.googleapis.com" class GoogleADC(Authentication): @@ -18,6 +22,7 @@ class GoogleADC(Authentication): scopes: ClassVar[list[str]] = ["https://www.googleapis.com/auth/cloud-platform"] project_id: str | None = None + location: str = "global" model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True) @@ -35,8 +40,11 @@ def get_headers(self) -> dict[str, str]: def _get_access_token(self) -> tuple[str, str | None]: """Get OAuth access token using Application Default Credentials.""" - import google.auth - import google.auth.transport.requests + try: + import google.auth + import google.auth.transport.requests + except ImportError as e: + raise MissingDependencyError(library="google-auth", extra="gcp") from e if self._credentials is None: self._credentials, self._project = google.auth.default(scopes=self.scopes) @@ -47,5 +55,18 @@ def _get_access_token(self) -> tuple[str, str | None]: return self._credentials.token, self.project_id or self._project + @property + def resolved_project_id(self) -> str | None: + """Return effective project_id (explicit or from ADC credentials).""" + if self._credentials is None: + self._get_access_token() + return self.project_id or self._project + + def get_vertex_base_url(self) -> str: + """Return the Vertex AI base URL for the configured location.""" + if self.location == "global": + return VERTEX_GLOBAL_BASE_URL + return VERTEX_BASE_URL.format(location=self.location) + __all__ = ["GoogleADC"]