Skip to content

Commit f6f9bc2

Browse files
committed
feat(skills): Support non-blocking skill loading in async runtimes
Add asynchronous counterparts for loading and listing local and GCS skills: - load_skill_from_dir_async - load_skill_from_gcs_dir_async - list_skills_in_dir_async - list_skills_in_gcs_dir_async These use asyncio.to_thread to run blocking I/O in a background thread pool, keeping the main event loop free. Closes #6057
1 parent 342b59d commit f6f9bc2

3 files changed

Lines changed: 210 additions & 0 deletions

File tree

src/google/adk/skills/__init__.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,13 @@
1818
import warnings
1919

2020
from ._utils import _list_skills_in_dir as list_skills_in_dir
21+
from ._utils import _list_skills_in_dir_async as list_skills_in_dir_async
2122
from ._utils import _list_skills_in_gcs_dir as list_skills_in_gcs_dir
23+
from ._utils import _list_skills_in_gcs_dir_async as list_skills_in_gcs_dir_async
2224
from ._utils import _load_skill_from_dir as load_skill_from_dir
25+
from ._utils import _load_skill_from_dir_async as load_skill_from_dir_async
2326
from ._utils import _load_skill_from_gcs_dir as load_skill_from_gcs_dir
27+
from ._utils import _load_skill_from_gcs_dir_async as load_skill_from_gcs_dir_async
2428
from .models import Frontmatter
2529
from .models import Resources
2630
from .models import Script
@@ -35,9 +39,13 @@
3539
"Skill",
3640
"SkillRegistry",
3741
"list_skills_in_dir",
42+
"list_skills_in_dir_async",
3843
"list_skills_in_gcs_dir",
44+
"list_skills_in_gcs_dir_async",
3945
"load_skill_from_dir",
46+
"load_skill_from_dir_async",
4047
"load_skill_from_gcs_dir",
48+
"load_skill_from_gcs_dir_async",
4149
]
4250

4351

src/google/adk/skills/_utils.py

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
from __future__ import annotations
1818

19+
import asyncio
1920
import io
2021
import logging
2122
import pathlib
@@ -550,3 +551,86 @@ def _load_files_in_dir(subdir: str) -> Dict[str, Union[str, bytes]]:
550551
instructions=body,
551552
resources=resources,
552553
)
554+
555+
556+
async def _load_skill_from_dir_async(
557+
skill_dir: Union[str, pathlib.Path],
558+
) -> models.Skill:
559+
"""Load a complete skill from a directory asynchronously.
560+
561+
Args:
562+
skill_dir: Path to the skill directory.
563+
564+
Returns:
565+
Skill object with all components loaded.
566+
"""
567+
return await asyncio.to_thread(_load_skill_from_dir, skill_dir)
568+
569+
570+
async def _load_skill_from_gcs_dir_async(
571+
bucket_name: str,
572+
skill_id: str,
573+
skills_base_path: str = "",
574+
project_id: str | None = None,
575+
credentials: auth.Credentials | None = None,
576+
) -> models.Skill:
577+
"""Load a complete skill from a GCS directory asynchronously.
578+
579+
Args:
580+
bucket_name: Name of the GCS bucket.
581+
skill_id: The ID of the skill (directory name).
582+
skills_base_path: Base directory within the bucket (e.g., 'path/to/skills').
583+
project_id: Project ID to use for GCS client.
584+
credentials: Credentials to use for GCS client.
585+
586+
Returns:
587+
Skill object with all components loaded.
588+
"""
589+
return await asyncio.to_thread(
590+
_load_skill_from_gcs_dir,
591+
bucket_name,
592+
skill_id,
593+
skills_base_path,
594+
project_id,
595+
credentials,
596+
)
597+
598+
599+
async def _list_skills_in_dir_async(
600+
skills_base_path: Union[str, pathlib.Path],
601+
) -> dict[str, models.Frontmatter]:
602+
"""List skills in a local directory asynchronously.
603+
604+
Args:
605+
skills_base_path: Path to the base directory containing skills.
606+
607+
Returns:
608+
Dictionary mapping skill IDs to their frontmatter.
609+
"""
610+
return await asyncio.to_thread(_list_skills_in_dir, skills_base_path)
611+
612+
613+
async def _list_skills_in_gcs_dir_async(
614+
bucket_name: str,
615+
skills_base_path: str = "",
616+
project_id: str | None = None,
617+
credentials: auth.Credentials | None = None,
618+
) -> dict[str, models.Frontmatter]:
619+
"""List skills in a GCS directory asynchronously.
620+
621+
Args:
622+
bucket_name: Name of the GCS bucket.
623+
skills_base_path: Base directory within the bucket (e.g., 'path/to/skills').
624+
project_id: Project ID to use for GCS client.
625+
credentials: Credentials to use for GCS client.
626+
627+
Returns:
628+
Dictionary mapping skill IDs to their frontmatter.
629+
"""
630+
return await asyncio.to_thread(
631+
_list_skills_in_gcs_dir,
632+
bucket_name,
633+
skills_base_path,
634+
project_id,
635+
credentials,
636+
)

tests/unittests/skills/test__utils.py

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,13 @@
2121
import zipfile
2222

2323
from google.adk.skills import list_skills_in_dir
24+
from google.adk.skills import list_skills_in_dir_async as _list_skills_in_dir_async
2425
from google.adk.skills import list_skills_in_gcs_dir as _list_skills_in_gcs_dir
26+
from google.adk.skills import list_skills_in_gcs_dir_async as _list_skills_in_gcs_dir_async
2527
from google.adk.skills import load_skill_from_dir as _load_skill_from_dir
28+
from google.adk.skills import load_skill_from_dir_async as _load_skill_from_dir_async
2629
from google.adk.skills import load_skill_from_gcs_dir as _load_skill_from_gcs_dir
30+
from google.adk.skills import load_skill_from_gcs_dir_async as _load_skill_from_gcs_dir_async
2731
from google.adk.skills._utils import _load_skill_from_zip_bytes
2832
from google.adk.skills._utils import _read_skill_properties
2933
from google.adk.skills._utils import _validate_skill_dir
@@ -393,3 +397,117 @@ def mock_import(name, globals=None, locals=None, fromlist=(), level=0):
393397
with mock.patch("builtins.__import__", mock_import):
394398
with pytest.raises(ImportError, match="google-cloud-storage is required"):
395399
_load_skill_from_gcs_dir("my-bucket", "skills/my-skill/")
400+
401+
402+
@pytest.mark.asyncio
403+
async def test_load_skill_from_dir_async(tmp_path):
404+
"""Tests loading a skill from a directory asynchronously."""
405+
skill_dir = tmp_path / "test-skill"
406+
skill_dir.mkdir()
407+
408+
skill_md_content = """---
409+
name: test-skill
410+
description: Test description
411+
---
412+
Test instructions
413+
"""
414+
(skill_dir / "SKILL.md").write_text(skill_md_content)
415+
416+
# Create references
417+
ref_dir = skill_dir / "references"
418+
ref_dir.mkdir()
419+
(ref_dir / "ref1.md").write_text("ref1 content")
420+
421+
skill = await _load_skill_from_dir_async(skill_dir)
422+
423+
assert skill.name == "test-skill"
424+
assert skill.description == "Test description"
425+
assert skill.instructions == "Test instructions"
426+
assert skill.resources.get_reference("ref1.md") == "ref1 content"
427+
428+
429+
@pytest.mark.asyncio
430+
async def test_list_skills_in_dir_async(tmp_path):
431+
"""Tests listing skills in a directory asynchronously."""
432+
skills_dir = tmp_path / "skills"
433+
skills_dir.mkdir()
434+
435+
# Valid skill 1
436+
skill1_dir = skills_dir / "skill1"
437+
skill1_dir.mkdir()
438+
(skill1_dir / "SKILL.md").write_text(
439+
"---\nname: skill1\ndescription: desc1\n---\nbody"
440+
)
441+
442+
skills = await _list_skills_in_dir_async(skills_dir)
443+
444+
assert len(skills) == 1
445+
assert "skill1" in skills
446+
assert skills["skill1"].name == "skill1"
447+
448+
449+
@pytest.mark.asyncio
450+
@mock.patch("google.cloud.storage.Client")
451+
async def test_load_skill_from_gcs_dir_async(mock_client_class):
452+
"""Tests loading a skill from GCS asynchronously."""
453+
mock_client = mock.MagicMock()
454+
mock_client_class.return_value = mock_client
455+
mock_bucket = mock.MagicMock()
456+
mock_client.bucket.return_value = mock_bucket
457+
458+
def mock_blob_side_effect(path):
459+
m = mock.MagicMock()
460+
if path.endswith("SKILL.md"):
461+
m.exists.return_value = True
462+
m.download_as_text.return_value = (
463+
"---\nname: my-skill\ndescription: Test description\n---\nTest"
464+
" instructions"
465+
)
466+
else:
467+
m.exists.return_value = False
468+
return m
469+
470+
mock_bucket.blob.side_effect = mock_blob_side_effect
471+
472+
# For resources
473+
def list_blobs_side_effect(prefix=None):
474+
if prefix.endswith("references/"):
475+
m = mock.MagicMock()
476+
m.name = prefix + "ref1.md"
477+
m.download_as_text.return_value = "ref1 content"
478+
return [m]
479+
return []
480+
481+
mock_bucket.list_blobs.side_effect = list_blobs_side_effect
482+
483+
skill = await _load_skill_from_gcs_dir_async("my-bucket", "skills/my-skill/")
484+
485+
assert skill.name == "my-skill"
486+
assert skill.description == "Test description"
487+
assert skill.instructions == "Test instructions"
488+
assert skill.resources.get_reference("ref1.md") == "ref1 content"
489+
490+
491+
@pytest.mark.asyncio
492+
@mock.patch("google.cloud.storage.Client")
493+
async def test_list_skills_in_gcs_dir_async(mock_client_class):
494+
"""Tests listing skills in GCS asynchronously."""
495+
mock_client = mock.MagicMock()
496+
mock_client_class.return_value = mock_client
497+
mock_bucket = mock.MagicMock()
498+
mock_client.bucket.return_value = mock_bucket
499+
500+
mock_iterator = mock.MagicMock()
501+
mock_iterator.prefixes = ["skills/my-skill/"]
502+
mock_bucket.list_blobs.return_value = mock_iterator
503+
504+
mock_blob = mock.MagicMock()
505+
mock_blob.exists.return_value = True
506+
mock_blob.download_as_text.return_value = (
507+
"---\nname: my-skill\ndescription: A skill\n---\nBody"
508+
)
509+
mock_bucket.blob.return_value = mock_blob
510+
511+
skills = await _list_skills_in_gcs_dir_async("my-bucket", "skills/")
512+
assert "my-skill" in skills
513+
assert skills["my-skill"].name == "my-skill"

0 commit comments

Comments
 (0)