feat(a2a): support remote agent auth (#223)

Author: yaozheng-fang

veadk/a2a/remote_ve_agent.py

@@ -13,21 +13,78 @@
 # limitations under the License.
 
 import json
+from typing import Literal
 
+import httpx
 import requests
 from a2a.types import AgentCard
 from google.adk.agents.remote_a2a_agent import RemoteA2aAgent
 
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
 AGENT_CARD_WELL_KNOWN_PATH = "/.well-known/agent-card.json"
 
 
+def _convert_agent_card_dict_to_obj(agent_card_dict: dict) -> AgentCard:
+    agent_card_json_str = json.dumps(agent_card_dict, ensure_ascii=False, indent=2)
+    agent_card_object = AgentCard.model_validate_json(str(agent_card_json_str))
+    return agent_card_object
+
+
 class RemoteVeAgent(RemoteA2aAgent):
-    def __init__(self, name: str, url: str):
-        agent_card_dict = requests.get(url + AGENT_CARD_WELL_KNOWN_PATH).json()
-        agent_card_dict["url"] = url
+    """Connect to remote agent on VeFaaS platform."""
+
+    def __init__(
+        self,
+        name: str,
+        url: str,
+        auth_token: str | None = None,
+        auth_method: Literal["header", "querystring"] | None = None,
+    ):
+        if not auth_token:
+            agent_card_dict = requests.get(url + AGENT_CARD_WELL_KNOWN_PATH).json()
+            # replace agent_card_url with actual host
+            agent_card_dict["url"] = url
+
+            agent_card_object = _convert_agent_card_dict_to_obj(agent_card_dict)
+
+            logger.debug(f"Agent card of {name}: {agent_card_object}")
+            super().__init__(name=name, agent_card=agent_card_object)
+        else:
+            if auth_method == "header":
+                headers = {"Authorization": f"Bearer {auth_token}"}
+                agent_card_dict = requests.get(
+                    url + AGENT_CARD_WELL_KNOWN_PATH, headers=headers
+                ).json()
+                agent_card_dict["url"] = url
+
+                agent_card_object = _convert_agent_card_dict_to_obj(agent_card_dict)
+                httpx_client = httpx.AsyncClient(
+                    base_url=url, headers=headers, timeout=600
+                )
 
-        agent_card_json_str = json.dumps(agent_card_dict, ensure_ascii=False, indent=2)
+                logger.debug(f"Agent card of {name}: {agent_card_object}")
+                super().__init__(
+                    name=name, agent_card=agent_card_object, httpx_client=httpx_client
+                )
+            elif auth_method == "querystring":
+                agent_card_dict = requests.get(
+                    url + AGENT_CARD_WELL_KNOWN_PATH + f"?token={auth_token}"
+                ).json()
+                agent_card_dict["url"] = url
 
-        agent_card_object = AgentCard.model_validate_json(str(agent_card_json_str))
+                agent_card_object = _convert_agent_card_dict_to_obj(agent_card_dict)
+                httpx_client = httpx.AsyncClient(
+                    base_url=url, params={"token": auth_token}, timeout=600
+                )
 
-        super().__init__(name=name, agent_card=agent_card_object)
+                logger.debug(f"Agent card of {name}: {agent_card_object}")
+                super().__init__(
+                    name=name, agent_card=agent_card_object, httpx_client=httpx_client
+                )
+            else:
+                raise ValueError(
+                    f"Unsupported auth method {auth_method}, use `header` or `querystring` instead."
+                )