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..2e0dcba9de1b 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,12 +75,11 @@
   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,