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
55 changes: 42 additions & 13 deletions src/google/adk/auth/auth_credential.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,10 @@
from __future__ import annotations

from enum import Enum
from typing import Annotated
from typing import Any
from typing import Dict
from typing import Iterator
from typing import List
from typing import Literal

Expand All @@ -26,22 +28,49 @@
from pydantic import Field
from pydantic import model_validator

_REDACTED = "<redacted>"


# Pydantic echoes the rejected value into ValidationError messages
# ("input_value=..."), which would put a malformed secret straight into logs and
# into the error strings surfaced to the LLM. The field name and error type are
# still reported.
class BaseModelWithConfig(BaseModel):
"""Base model for credential types, hardened against leaking secrets."""

model_config = ConfigDict(
extra="allow",
alias_generator=alias_generators.to_camel,
populate_by_name=True,
hide_input_in_errors=True,
)
"""The pydantic model config."""

def __repr_args__(self) -> Iterator[tuple[str | None, Any]]:
"""Redacts the values of extra (unmodeled) fields from repr and str.

`extra="allow"` lets callers attach arbitrary keys to these credential
models, and pydantic renders extras in repr unconditionally: marking a
declared field `repr=False` does nothing for a secret that arrives under an
unexpected key (e.g. a non-standard field in an OAuth2 token response).
Redacting the values keeps them out of logs and out of error strings that
reach the LLM, while still showing which keys were set.

Yields:
`(name, value)` pairs to render, with the values of extra fields replaced
by a redaction placeholder.
"""
extra = self.__pydantic_extra__ or {}
for key, value in super().__repr_args__():
yield key, _REDACTED if key in extra else value


class HttpCredentials(BaseModelWithConfig):
"""Represents the secret token value for HTTP authentication, like user name, password, oauth token, etc."""

username: str | None = None
password: str | None = None
token: str | None = None
password: Annotated[str | None, Field(repr=False)] = None
token: Annotated[str | None, Field(repr=False)] = None

@classmethod
def model_validate(cls, data: Dict[str, Any]) -> "HttpCredentials":
Expand All @@ -61,14 +90,14 @@ class HttpAuth(BaseModelWithConfig):
# Examples: 'basic', 'bearer'
scheme: str
credentials: HttpCredentials
additional_headers: Dict[str, str] | None = None
additional_headers: Annotated[Dict[str, str] | None, Field(repr=False)] = None


class OAuth2Auth(BaseModelWithConfig):
"""Represents credential value and its metadata for a OAuth2 credential."""

client_id: str | None = None
client_secret: str | None = None
client_secret: Annotated[str | None, Field(repr=False)] = None
# tool or adk can generate the auth_uri with the state info thus client
# can verify the state
auth_uri: str | None = None
Expand All @@ -79,15 +108,15 @@ class OAuth2Auth(BaseModelWithConfig):
state: str | None = None
# tool or adk can decide the redirect_uri if they don't want client to decide
redirect_uri: str | None = None
auth_response_uri: str | None = None
auth_code: str | None = None
access_token: str | None = None
refresh_token: str | None = None
id_token: str | None = None
auth_response_uri: Annotated[str | None, Field(repr=False)] = None
auth_code: Annotated[str | None, Field(repr=False)] = None
access_token: Annotated[str | None, Field(repr=False)] = None
refresh_token: Annotated[str | None, Field(repr=False)] = None
id_token: Annotated[str | None, Field(repr=False)] = None
expires_at: int | None = None
expires_in: int | None = None
audience: str | None = None
code_verifier: str | None = None
code_verifier: Annotated[str | None, Field(repr=False)] = None
code_challenge_method: str | None = None
token_endpoint_auth_method: (
Literal[
Expand Down Expand Up @@ -140,8 +169,8 @@ class ServiceAccountCredential(BaseModelWithConfig):

type_: str = Field("", alias="type")
project_id: str
private_key_id: str
private_key: str
private_key_id: Annotated[str, Field(repr=False)]
private_key: Annotated[str, Field(repr=False)]
client_email: str
client_id: str
auth_uri: str
Expand Down Expand Up @@ -279,7 +308,7 @@ class AuthCredential(BaseModelWithConfig):
# This will be supported in the future.
resource_ref: str | None = None

api_key: str | None = None
api_key: Annotated[str | None, Field(repr=False)] = None
http: HttpAuth | None = None
service_account: ServiceAccount | None = None
oauth2: OAuth2Auth | None = None
9 changes: 7 additions & 2 deletions src/google/adk/flows/llm_flows/base_llm_flow.py
Original file line number Diff line number Diff line change
Expand Up @@ -528,10 +528,15 @@ async def run_live(
return

llm = self.__get_llm(invocation_context)
# Only log non-sensitive request metadata. The full request carries the
# user conversation and http_options.headers, which may hold credentials.
logger.debug(
'Establishing live connection for agent: %s with llm request: %s',
'Establishing live connection for agent: %s, model: %s, contents: %s,'
' response modalities: %s',
invocation_context.agent.name,
llm_request,
llm_request.model,
len(llm_request.contents),
llm_request.live_connect_config.response_modalities,
)

try:
Expand Down
22 changes: 19 additions & 3 deletions src/google/adk/models/google_llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -464,8 +464,21 @@ async def connect(self, llm_request: LlmRequest) -> BaseLlmConnection:
llm_request.live_connect_config.safety_settings = (
llm_request.config.safety_settings
)
logger.debug('Connecting to live with llm_request:%s', llm_request)
logger.debug('Live connect config: %s', llm_request.live_connect_config)
logger.debug(
'Connecting to live with model: %s, contents: %d, response modalities:'
' %s',
llm_request.model,
len(llm_request.contents or []),
llm_request.live_connect_config.response_modalities,
)
# Callers may put credentials in per-request headers, so the transport
# options never go to the log.
logger.debug(
'Live connect config: %s',
llm_request.live_connect_config.model_copy(
update={'http_options': None}
),
)
async with self._live_api_client.aio.live.connect(
model=llm_request.model, config=llm_request.live_connect_config
) as live_session:
Expand Down Expand Up @@ -592,11 +605,14 @@ def _build_request_log(req: LlmRequest) -> str:
exclude={
'system_instruction': True,
'tools': tools_exclusion if req.config.tools else True,
# Callers may put credentials in per-request headers, so the
# transport options never go to the log.
'http_options': True,
},
)
)
except Exception:
config_log = repr(req.config)
config_log = repr(req.config.model_copy(update={'http_options': None}))

return f"""
LLM Request:
Expand Down
60 changes: 57 additions & 3 deletions src/google/adk/telemetry/tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -360,7 +360,12 @@ def trace_call_llm(

if _should_add_request_response_to_spans():
try:
llm_response_json = llm_response.model_dump_json(exclude_none=True)
response_for_trace = llm_response
if llm_response.content is not None:
response_for_trace = llm_response.model_copy(
update={'content': _summarize_inline_data(llm_response.content)}
)
llm_response_json = response_for_trace.model_dump_json(exclude_none=True)
except Exception: # pylint: disable=broad-exception-caught
llm_response_json = '<not serializable>'

Expand Down Expand Up @@ -409,6 +414,37 @@ def trace_call_llm(
)


def _summarize_inline_data(content: types.Content) -> types.Content:
"""Returns ``content`` with inline binary parts reduced to a description.

Serializing a part in JSON mode base64-encodes its ``inline_data``, so a
live session's audio chunks would otherwise be copied wholesale onto a span
attribute. Only the mime type and byte count are kept.

Args:
content: The content to summarize.

Returns:
A copy of ``content`` whose inline binary parts carry a text description
instead of the bytes.
"""
parts: list[types.Part] = []
for part in content.parts or []:
blob = part.inline_data
if blob is None:
parts.append(part)
continue
parts.append(
types.Part(
text=(
f"<inline_data: {blob.mime_type or 'unknown'},"
f" {len(blob.data or b'')} bytes>"
)
)
)
return types.Content(role=content.role, parts=parts)


def trace_send_data(
invocation_context: InvocationContext,
event_id: str,
Expand All @@ -435,7 +471,7 @@ def trace_send_data(
span.set_attribute(
'gcp.vertex.agent.data',
_safe_json_serialize([
types.Content(role=content.role, parts=content.parts).model_dump(
_summarize_inline_data(content).model_dump(
exclude_none=True, mode='json'
)
for content in data
Expand Down Expand Up @@ -514,7 +550,25 @@ def _build_llm_request_for_trace(llm_request: LlmRequest) -> dict[str, Any]:
result = {
'model': llm_request.model,
'config': llm_request.config.model_dump(
exclude_none=True, exclude='response_schema', mode='json'
exclude_none=True,
exclude={
'response_schema': True,
# `http_options` carries caller-supplied credentials: `headers`
# commonly holds an Authorization bearer token, and
# `extra_body` / `*client_args` are free-form passthroughs that
# can hold auth material too. None of it may reach an exported
# span attribute. The client fields are also unserializable.
'http_options': {
'httpx_client': True,
'httpx_async_client': True,
'aiohttp_client': True,
'headers': True,
'extra_body': True,
'client_args': True,
'async_client_args': True,
},
},
mode='json',
),
'contents': [],
}
Expand Down
3 changes: 1 addition & 2 deletions src/google/adk/tools/mcp_tool/mcp_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -576,8 +576,7 @@ async def _get_headers(
or not self._credentials_manager._auth_config
):
error_msg = (
"Cannot find corresponding auth scheme for API key credential"
f" {credential}"
"Cannot find corresponding auth scheme for API key credential."
)
logger.error(error_msg)
raise ValueError(error_msg)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -564,8 +564,7 @@ def __repr__(self):
return (
f'RestApiTool(name="{self.name}", description="{self.description}",'
f' endpoint="{self.endpoint}", operation="{self.operation}",'
f' auth_scheme="{self.auth_scheme}",'
f' auth_credential="{self.auth_credential}")'
f' auth_scheme="{self.auth_scheme}")'
)


Expand Down
Loading
Loading