Skip to content

[Agent Engine][ADK] Add methods for managing artifact to AdkApp #5352

Open
@soundTricker

Description

@soundTricker

Is your feature request related to a problem? Please describe.
When adk runner uses InMemoryArtifactService, there is no way to extract the artifact.
Also, even if we are using GcsArtifactService, since it is the Agent that manages the artifacts, you should be able to obtain the artifacts from the Agent (Engine) in the same way as with the Session.

Describe the solution you'd like
Add all artifact management method load_artifact save_artifact list_artifact_keys delete_artifact list_versions to AdkApp like session.

Additional context
Current workaround is creating custom adk app class like below.

class CustomAdkApp(reasoning_engines.AdkApp):
    def clone(self):
        """Returns a clone of the ADK application."""
        import copy

        return CustomAdkApp(
            agent=copy.deepcopy(self._tmpl_attrs.get("agent")),
            enable_tracing=self._tmpl_attrs.get("enable_tracing"),
            session_service_builder=self._tmpl_attrs.get("session_service_builder"),
            artifact_service_builder=self._tmpl_attrs.get("artifact_service_builder"),
        )

    def load_artifact(self, user_id: str, session_id: str, artifact_id: str, **kwargs):

        import nest_asyncio
        nest_asyncio.apply()

        artifact_service: BaseArtifactService = self._tmpl_attrs["artifact_service"]
        return asyncio.run(artifact_service.load_artifact(
            app_name=self._tmpl_attrs.get("app_name"), user_id=user_id, session_id=session_id, filename=artifact_id)
        )

    def register_operations(self) -> dict[str, list[str]]:
        """Registers the operations of the ADK application."""
        return {
            "": [
                "get_session",
                "list_sessions",
                "create_session",
                "delete_session",
                "load_artifact"
            ],
            "stream": ["stream_query", "streaming_agent_run_with_events"],
            "async_stream": ["async_stream_query"],
        }

Metadata

Metadata

Assignees

Labels

api: vertex-aiIssues related to the googleapis/python-aiplatform API.

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions