diff --git a/samples/agent/adk/orchestrator/agent.py b/samples/agent/adk/orchestrator/agent.py
index e59ca5c0..4e8c441a 100644
--- a/samples/agent/adk/orchestrator/agent.py
+++ b/samples/agent/adk/orchestrator/agent.py
@@ -15,7 +15,8 @@
 import json
 import logging
 import os
-from a2a.client import A2ACardResolver
+from typing import Any, List
+from a2a.client import A2ACardResolver, Client, ClientCallContext, Consumer
 from a2a.extensions.common import HTTP_EXTENSION_HEADER
 from google.adk.models.lite_llm import LiteLlm
 from google.adk.agents.llm_agent import LlmAgent
@@ -29,8 +30,9 @@
 from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
 from subagent_route_manager import SubagentRouteManager
-from typing import Any, override, List
-from a2a.types import TransportProtocol as A2ATransport
+from a2ui.a2ui_extension import is_a2ui_part, A2UI_EXTENSION_URI
+from typing import override
+from a2a.types import AgentCard, TransportProtocol as A2ATransport
 
 from a2a.client.client import Consumer, Client
 from a2a.client.middleware import ClientCallContext, ClientCallInterceptor