From 609528f2a0057335a278d353542879858f8be123 Mon Sep 17 00:00:00 2001 From: Danny Mccormick Date: Wed, 1 Apr 2026 10:47:11 -0400 Subject: [PATCH 1/2] Fix lint --- .../apache_beam/ml/inference/agent_development_kit.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 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 392a46673983..196fc62c1934 100644 --- a/sdks/python/apache_beam/ml/inference/agent_development_kit.py +++ b/sdks/python/apache_beam/ml/inference/agent_development_kit.py @@ -57,6 +57,7 @@ from collections.abc import Sequence from typing import Any from typing import Optional +from typing import Union from apache_beam.ml.inference.base import ModelHandler from apache_beam.ml.inference.base import PredictionResult @@ -74,15 +75,14 @@ ADK_AVAILABLE = False genai_Content = Any # type: ignore[assignment, misc] genai_Part = Any # type: ignore[assignment, misc] - Agent = None LOGGER = logging.getLogger("ADKAgentModelHandler") # Type alias for an agent or factory that produces one -_AgentOrFactory = Agent | Callable[[], Agent] +_AgentOrFactory = Union["Agent", Callable[[], "Agent"]] -class ADKAgentModelHandler(ModelHandler[str | genai_Content, +class ADKAgentModelHandler(ModelHandler[Union[str, genai_Content], PredictionResult, "Runner"]): """ModelHandler for running ADK agents with the Beam RunInference transform. @@ -185,7 +185,7 @@ def load_model(self) -> "Runner": def run_inference( self, - batch: Sequence[str | genai_Content], + batch: Sequence[Union[str, genai_Content]], model: "Runner", inference_args: Optional[dict[str, Any]] = None, ) -> Iterable[PredictionResult]: From 5da3d2f00d4e9b911c273741a4a65ac7f4027d84 Mon Sep 17 00:00:00 2001 From: Danny Mccormick Date: Wed, 1 Apr 2026 10:54:30 -0400 Subject: [PATCH 2/2] Shrink part of revert --- sdks/python/apache_beam/ml/inference/agent_development_kit.py | 4 ++-- 1 file changed, 2 insertions(+), 2 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 196fc62c1934..2e0dcba9de1b 100644 --- a/sdks/python/apache_beam/ml/inference/agent_development_kit.py +++ b/sdks/python/apache_beam/ml/inference/agent_development_kit.py @@ -82,7 +82,7 @@ _AgentOrFactory = Union["Agent", Callable[[], "Agent"]] -class ADKAgentModelHandler(ModelHandler[Union[str, genai_Content], +class ADKAgentModelHandler(ModelHandler[str | genai_Content, PredictionResult, "Runner"]): """ModelHandler for running ADK agents with the Beam RunInference transform. @@ -185,7 +185,7 @@ def load_model(self) -> "Runner": def run_inference( self, - batch: Sequence[Union[str, genai_Content]], + batch: Sequence[str | genai_Content], model: "Runner", inference_args: Optional[dict[str, Any]] = None, ) -> Iterable[PredictionResult]: