This is an automated email from the ASF dual-hosted git repository. damccorm pushed a commit to branch users/damccorm/awaitSessionCreation in repository https://gitbox.apache.org/repos/asf/beam.git
commit 740b834b3089c9be4345a2024a04cfa21f9d62aa Author: Danny McCormick <[email protected]> AuthorDate: Tue May 12 16:33:24 2026 -0400 Make sure session creation happens before starting agent --- .../ml/inference/agent_development_kit.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/agent_development_kit.py b/sdks/python/apache_beam/ml/inference/agent_development_kit.py index 1130598f06f..8bcca74ee14 100644 --- a/sdks/python/apache_beam/ml/inference/agent_development_kit.py +++ b/sdks/python/apache_beam/ml/inference/agent_development_kit.py @@ -229,17 +229,6 @@ class ADKAgentModelHandler(ModelHandler[str | genai_Content, for element in batch: session_id: str = inference_args.get("session_id", str(uuid.uuid4())) - # Ensure a session exists for this invocation - try: - model.session_service.create_session( - app_name=self._app_name, - user_id=user_id, - session_id=session_id, - ) - except sessions.SessionExistsError: - # It's okay if the session already exists for shared session IDs. - pass - # Wrap plain strings in a Content object if isinstance(element, str): # pyrefly: ignore[bad-instantiation] @@ -288,6 +277,16 @@ class ADKAgentModelHandler(ModelHandler[str | genai_Content, The text of the agent's final response, or ``None`` if the agent produced no final text response. """ + # Ensure a session exists for this invocation + try: + await model.session_service.create_session( + app_name=self._app_name, + user_id=user_id, + session_id=session_id, + ) + except sessions.SessionExistsError: + # It's okay if the session already exists for shared session IDs. + pass async for event in runner.run_async( user_id=user_id, session_id=session_id,
