Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 34 additions & 0 deletions google/genai/_gaos/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# pyformat: disable

"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT."""

from ._version import (
__title__,
__version__,
__openapi_doc_version__,
__gen_version__,
__user_agent__,
)
from .sdk import *
from .sdkconfiguration import *
from . import resources


VERSION: str = __version__
OPENAPI_DOC_VERSION = __openapi_doc_version__
SPEAKEASY_GENERATOR_VERSION = __gen_version__
USER_AGENT = __user_agent__
21 changes: 21 additions & 0 deletions google/genai/_gaos/_hooks/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# pyformat: disable

"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT."""

from .sdkhooks import *
from .types import *
from .registration import *
157 changes: 157 additions & 0 deletions google/genai/_gaos/_hooks/google_genai_auth.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,157 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# pyformat: disable

"""Google GenAI authentication hooks for the embedded Speakeasy SDK."""

from __future__ import annotations

from typing import Any, Callable, Mapping, Optional, Union, cast

import httpx

from .. import types, utils
from .types import BeforeRequestContext, BeforeRequestHook


GOOGLE_GENAI_API_REVISION = "2026-05-20"
_MANAGED_BEARER_AUTH = "google_genai_managed_bearer_auth"


class GoogleGenAISecurityProvider:
"""Callable security source that keeps generated auth refreshes retry-safe."""

def __init__(
self,
*,
access_token: Callable[[], str],
default_headers: Optional[Mapping[str, str]] = None,
) -> None:
self._access_token = access_token
self._default_headers = dict(default_headers or {})
self._pending_access_token: Optional[str] = None

def __call__(self) -> types.Security:
if self._pending_access_token is None:
self._pending_access_token = self._access_token()

return types.Security(
access_token=self._pending_access_token,
default_headers=self._default_headers or None,
)

def consume(self) -> None:
self._pending_access_token = None


class GoogleGenAIAuthHook(BeforeRequestHook):
"""Applies Google GenAI headers and custom auth before sending requests."""

def before_request(
self,
hook_ctx: BeforeRequestContext,
request: httpx.Request,
) -> Union[httpx.Request, Exception]:
security = _resolve_security(hook_ctx.security_source)

_apply_default_headers(security, request)
_apply_api_revision(hook_ctx, request)
_apply_user_project(hook_ctx, request)
_apply_auth(security, request)

return request


def _resolve_security(security_source: Any) -> Optional[types.Security]:
security = security_source

if callable(security_source):
security = security_source()
consume = getattr(security_source, "consume", None)
if callable(consume):
consume()

security = utils.get_security_from_env(security, types.Security)
if security is None:
return None

if isinstance(security, types.Security):
return security

if isinstance(security, Mapping):
return types.Security(**security)

return cast(types.Security, security)


def _apply_default_headers(
security: Optional[types.Security], request: httpx.Request
) -> None:
headers = security.default_headers if security else None
for key, value in (headers or {}).items():
if request.headers.get(key) is None:
request.headers[key] = value


def _apply_api_revision(
hook_ctx: BeforeRequestContext, request: httpx.Request
) -> None:
if request.headers.get("Api-Revision") is not None:
return

api_revision = hook_ctx.config.globals.api_revision
request.headers["Api-Revision"] = api_revision or GOOGLE_GENAI_API_REVISION


def _apply_user_project(
hook_ctx: BeforeRequestContext, request: httpx.Request
) -> None:
user_project = hook_ctx.config.globals.user_project
if user_project and request.headers.get("x-goog-user-project") is None:
request.headers["x-goog-user-project"] = user_project


def _has_auth_headers(request: httpx.Request) -> bool:
return (
request.headers.get("authorization") is not None
or request.headers.get("x-goog-api-key") is not None
)


def _has_user_auth_headers(request: httpx.Request) -> bool:
return (
not request.extensions.get(_MANAGED_BEARER_AUTH)
and _has_auth_headers(request)
)


def _apply_auth(
security: Optional[types.Security], request: httpx.Request
) -> None:
if security is None or _has_user_auth_headers(request):
return

if security.api_key:
request.extensions.pop(_MANAGED_BEARER_AUTH, None)
request.headers["x-goog-api-key"] = security.api_key
return

if security.access_token:
request.extensions[_MANAGED_BEARER_AUTH] = True
request.headers["Authorization"] = _bearer(security.access_token)


def _bearer(token: str) -> str:
return token if token.lower().startswith("bearer ") else f"Bearer {token}"
29 changes: 29 additions & 0 deletions google/genai/_gaos/_hooks/registration.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# pyformat: disable

from .types import Hooks


# This file is only ever generated once on the first generation and then is free to be modified.
# Any hooks you wish to add should be registered in the init_hooks function. Feel free to define them
# in this file or in separate files in the hooks folder.


def init_hooks(hooks: Hooks):
# pylint: disable=unused-argument
"""Add hooks by calling hooks.register{sdk_init/before_request/after_success/after_error}Hook
with an instance of a hook that implements that specific Hook interface
Hooks are registered per SDK instance, and are valid for the lifetime of the SDK instance"""
94 changes: 94 additions & 0 deletions google/genai/_gaos/_hooks/sdkhooks.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# pyformat: disable

"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT."""

import httpx
from .types import (
SDKInitHook,
BeforeRequestContext,
BeforeRequestHook,
AfterSuccessContext,
AfterSuccessHook,
AfterErrorContext,
AfterErrorHook,
Hooks,
)
from .registration import init_hooks
from .google_genai_auth import GoogleGenAIAuthHook
from typing import List, Optional, Tuple
from ..sdkconfiguration import SDKConfiguration


class SDKHooks(Hooks):
def __init__(self) -> None:
self.sdk_init_hooks: List[SDKInitHook] = []
self.before_request_hooks: List[BeforeRequestHook] = []
self.after_success_hooks: List[AfterSuccessHook] = []
self.after_error_hooks: List[AfterErrorHook] = []
init_hooks(self)
self.register_before_request_hook(GoogleGenAIAuthHook())

def register_sdk_init_hook(self, hook: SDKInitHook) -> None:
self.sdk_init_hooks.append(hook)

def register_before_request_hook(self, hook: BeforeRequestHook) -> None:
self.before_request_hooks.append(hook)

def register_after_success_hook(self, hook: AfterSuccessHook) -> None:
self.after_success_hooks.append(hook)

def register_after_error_hook(self, hook: AfterErrorHook) -> None:
self.after_error_hooks.append(hook)

def sdk_init(self, config: SDKConfiguration) -> SDKConfiguration:
for hook in self.sdk_init_hooks:
config = hook.sdk_init(config)
return config

def before_request(
self, hook_ctx: BeforeRequestContext, request: httpx.Request
) -> httpx.Request:
for hook in self.before_request_hooks:
out = hook.before_request(hook_ctx, request)
if isinstance(out, Exception):
raise out
request = out

return request

def after_success(
self, hook_ctx: AfterSuccessContext, response: httpx.Response
) -> httpx.Response:
for hook in self.after_success_hooks:
out = hook.after_success(hook_ctx, response)
if isinstance(out, Exception):
raise out
response = out
return response

def after_error(
self,
hook_ctx: AfterErrorContext,
response: Optional[httpx.Response],
error: Optional[Exception],
) -> Tuple[Optional[httpx.Response], Optional[Exception]]:
for hook in self.after_error_hooks:
result = hook.after_error(hook_ctx, response, error)
if isinstance(result, Exception):
raise result
response, error = result
return response, error
Loading
Loading