Skip to content
Open
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
42 changes: 38 additions & 4 deletions src/google/adk/artifacts/file_artifact_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,10 @@ def _iter_artifact_dirs(root: Path) -> list[Path]:
current = Path(dirpath)
if (current / "versions").exists():
artifact_dirs.append(current)
dirnames.clear()
# An artifact directory doubles as the parent of anything nested under
# it ("doc" and "doc/nested"), so keep walking, skipping only the
# stored versions of this artifact.
dirnames[:] = [name for name in dirnames if name != "versions"]
return artifact_dirs


Expand Down Expand Up @@ -293,6 +296,28 @@ def _canonical_uri(artifact_dir: Path, version: int) -> str:
return payload_path.resolve().as_uri()


def _prune_empty_dirs(leaf: Path, stop_at: Path) -> None:
"""Removes `leaf` and any parents it leaves empty, stopping at `stop_at`.

Filenames may contain "/", so the directory of an artifact doubles as the
parent directory of every artifact nested under it: "doc" is stored at
``{scope}/doc`` and "doc/nested" at ``{scope}/doc/nested``. A directory may
therefore only be removed once it holds nothing, or deleting "doc" would
take "doc/nested" with it.

Args:
leaf: Directory to remove, if it is empty.
stop_at: Scope root. It and everything above it are never removed.
"""
current = leaf
while current != stop_at and current.is_relative_to(stop_at):
try:
current.rmdir() # Only succeeds on an empty directory.
except OSError:
return
current = current.parent


def _list_versions_on_disk(artifact_dir: Path) -> list[int]:
"""Returns sorted versions discovered under the artifact directory."""
versions_dir = _versions_dir(artifact_dir)
Expand Down Expand Up @@ -703,9 +728,18 @@ def _delete_artifact_sync(
session_id=session_id,
filename=filename,
)
if artifact_dir.exists():
shutil.rmtree(artifact_dir)
logger.debug("Deleted artifact %s at %s", filename, artifact_dir)
versions_dir = _versions_dir(artifact_dir)
if not versions_dir.exists():
return
# Only this artifact's own versions go. Its directory may also be the
# parent of a nested artifact ("doc" vs "doc/nested"), so it is pruned
# separately and only if nothing is left under it.
shutil.rmtree(versions_dir)
scope_root = self._scope_root(
self._base_root(app_name, user_id), session_id, filename
)
_prune_empty_dirs(artifact_dir, scope_root)
logger.debug("Deleted artifact %s at %s", filename, artifact_dir)

@override
async def list_versions(
Expand Down
64 changes: 51 additions & 13 deletions src/google/adk/artifacts/gcs_artifact_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,42 @@
logger = logging.getLogger("google_adk." + __name__)


def _parse_version(blob_name: str, prefix: str) -> Optional[int]:
"""Extracts the version of an artifact from one of its blob names.

GCS has a flat namespace, so listing by prefix is a plain string match with
no notion of nesting depth. Because filenames are allowed to contain "/",
the prefix of an artifact is also a prefix of every artifact nested under it:
scanning "a/" to find versions of "a" also returns "a/b/3", which is version
3 of the distinct artifact "a/b".

A blob only holds a version of the artifact denoted by ``prefix`` when its
name is exactly ``{prefix}{version}``, so anything with a further "/" in it
belongs to some other artifact and must be skipped.

Args:
blob_name: The full name of the blob, which must start with ``prefix``.
prefix: The blob prefix of the artifact, including the trailing "/".

Returns:
The version number, or None if the blob does not hold a version of this
artifact.
"""
suffix = blob_name[len(prefix) :]
if "/" in suffix:
# Belongs to a distinct artifact nested under this one.
return None
# int() also accepts surrounding whitespace, underscores and non-ASCII
# digits, none of which _get_blob_name can produce.
if not (suffix.isascii() and suffix.isdigit()):
logger.warning(
"Skipping blob %s because it does not end with a version number.",
blob_name,
)
return None
return int(suffix)


class GcsArtifactService(BaseArtifactService):
"""An artifact service implementation using Google Cloud Storage (GCS)."""

Expand Down Expand Up @@ -357,12 +393,17 @@ def _list_versions(
artifact.
Returns an empty list if no versions are found.
"""
prefix = self._get_blob_prefix(app_name, user_id, filename, session_id)
blobs = self.storage_client.list_blobs(self.bucket, prefix=f"{prefix}/")
prefix = (
f"{self._get_blob_prefix(app_name, user_id, filename, session_id)}/"
)
blobs = self.storage_client.list_blobs(self.bucket, prefix=prefix)
versions = []
for blob in blobs:
*_, version = blob.name.split("/")
versions.append(int(version))
version = _parse_version(blob.name, prefix)
if version is None:
continue

versions.append(version)
return versions

def _get_artifact_version_sync(
Expand Down Expand Up @@ -410,17 +451,14 @@ def _list_artifact_versions_sync(
filename: str,
) -> list[ArtifactVersion]:
"""Lists all versions and their metadata of an artifact."""
prefix = self._get_blob_prefix(app_name, user_id, filename, session_id)
blobs = self.storage_client.list_blobs(self.bucket, prefix=f"{prefix}/")
prefix = (
f"{self._get_blob_prefix(app_name, user_id, filename, session_id)}/"
)
blobs = self.storage_client.list_blobs(self.bucket, prefix=prefix)
artifact_versions = []
for blob in blobs:
try:
version = int(blob.name.split("/")[-1])
except ValueError:
logger.warning(
"Skipping blob %s because it does not end with a version number.",
blob.name,
)
version = _parse_version(blob.name, prefix)
if version is None:
continue

canonical_uri = f"gs://{self.bucket_name}/{blob.name}"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -64,8 +64,11 @@ async def get_skill(self, *, name: str) -> models.Skill:
# be a single path segment before it is interpolated into the resource
# path. Accept the same character set skill names are already held to; the
# snake-or-kebab pattern is the superset of the two accepted spellings.
# fullmatch, not match: the pattern ends in `$`, which outside MULTILINE
# also matches just before a trailing newline, so `match` would let
# "my-skill\n" through into the resource name.
# pylint: disable-next=protected-access
if not models._SNAKE_OR_KEBAB_NAME_PATTERN.match(name):
if not models._SNAKE_OR_KEBAB_NAME_PATTERN.fullmatch(name):
raise ValueError(
f"Invalid skill name {name!r}: name must be lowercase kebab-case"
" (a-z, 0-9, hyphens) or snake_case (a-z, 0-9, underscores), with"
Expand Down
13 changes: 11 additions & 2 deletions src/google/adk/memory/vertex_ai_rag_memory_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import binascii
from collections import OrderedDict
import json
import logging
import os
import tempfile
from typing import Optional
Expand All @@ -37,6 +38,7 @@
from ..events.event import Event
from ..sessions.session import Session

logger = logging.getLogger("google_adk." + __name__)

_SOURCE_DISPLAY_NAME_PREFIX = "adk-memory-v1."

Expand Down Expand Up @@ -178,8 +180,15 @@ async def add_session_to_memory(self, session: Session) -> None:
if temp_file_path:
try:
os.remove(temp_file_path)
except FileNotFoundError:
pass
except OSError:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Checked the other three against upstream and they line up: _prune_empty_dirs, _iter_artifact_dirs and _delete_artifact_sync are byte-identical to origin/main, as is the GCS _parse_version. Good catch on the list_versions side of the GCS conflation -- a read resolving max(versions) to a version the artifact does not own was not something I had traced.

Also, thank you for correcting the skill-name suggestion. I offered .fullmatch() or quote() as equivalent options and they are not: v1 hands a resource name to self._client.skills.get(name=...) where main builds a URL string, so percent-encoding here would double-encode. fullmatch is the right call and the comment you added says why.

This one is the exception to the pattern, and it is the only thing I would raise. It is not a port -- origin/main:283-288 is still except FileNotFoundError: pass, so this leaves v1 ahead of main and the next backport touching this file has a divergence to reconcile by hand. The description already makes the argument that would justify fixing main too: "the displacement is a property of finally rather than of one platform". The to_thread wrapper makes it easier to hit on v1, but a read-only mount or a permissions problem displaces the propagating exception on main just the same.

Non-blocking, and I would not hold this PR for it -- would it be worth sending the same three lines upstream so the two converge rather than drift?

# Best effort: this runs in a finally, so raising here would
# displace the exception already propagating, and a cancelled
# upload can leave the worker thread still holding the file.
logger.warning(
"Could not remove the temporary transcript at %s",
temp_file_path,
exc_info=True,
)

@override
async def search_memory(
Expand Down
Loading
Loading