| 1 | import uuid |
| 2 | from typing import Any, Dict, List, Optional |
| 3 | from helpers.print_style import PrintStyle |
| 4 | |
| 5 | try: |
| 6 | from fasta2a.client import A2AClient # type: ignore |
| 7 | import httpx # type: ignore |
| 8 | FASTA2A_CLIENT_AVAILABLE = True |
| 9 | except ImportError: |
| 10 | FASTA2A_CLIENT_AVAILABLE = False |
| 11 | PrintStyle.warning("FastA2A client not available. Agent-to-agent communication disabled.") |
| 12 | |
| 13 | _PRINTER = PrintStyle(italic=True, font_color="cyan", padding=False) |
| 14 | |
| 15 | |
| 16 | class AgentConnection: |
| 17 | """Helper class for connecting to and communicating with other Agent Zero instances via FastA2A.""" |
| 18 | |
| 19 | def __init__(self, agent_url: str, timeout: int = 30, token: Optional[str] = None): |
| 20 | """Initialize connection to an agent. |
| 21 | |
| 22 | Args: |
| 23 | agent_url: The base URL of the agent (e.g., "https://agent.example.com") |
| 24 | timeout: Request timeout in seconds |
| 25 | """ |
| 26 | if not FASTA2A_CLIENT_AVAILABLE: |
| 27 | raise RuntimeError("FastA2A client not available") |
| 28 | |
| 29 | # Ensure scheme is present |
| 30 | if not agent_url.startswith(('http://', 'https://')): |
| 31 | agent_url = 'http://' + agent_url |
| 32 | |
| 33 | self.agent_url = agent_url.rstrip('/') |
| 34 | self.timeout = timeout |
| 35 | # Auth headers |
| 36 | if token is None: |
| 37 | import os |
| 38 | token = os.getenv("A2A_TOKEN") |
| 39 | headers = {} |
| 40 | if token: |
| 41 | headers["Authorization"] = f"Bearer {token}" |
| 42 | headers["X-API-KEY"] = token |
| 43 | self._http_client = httpx.AsyncClient(timeout=timeout, headers=headers) # type: ignore |
| 44 | self._a2a_client = A2AClient(base_url=self.agent_url, http_client=self._http_client) # type: ignore |
| 45 | self._agent_card: Optional[Dict[str, Any]] = None |
| 46 | # Track conversation context automatically |
| 47 | self._context_id: Optional[str] = None |
| 48 | |
| 49 | async def get_agent_card(self) -> Dict[str, Any]: |
| 50 | """Retrieve the agent card from the remote agent.""" |
| 51 | if self._agent_card is None: |
| 52 | try: |
| 53 | response = await self._http_client.get(f"{self.agent_url}/.well-known/agent.json") |
| 54 | response.raise_for_status() |
| 55 | self._agent_card = response.json() |
| 56 | _PRINTER.print(f"Retrieved agent card from {self.agent_url}") |
| 57 | _PRINTER.print(f"Agent: {self._agent_card.get('name', 'Unknown')}") # type: ignore |
| 58 | _PRINTER.print(f"Description: {self._agent_card.get('description', 'No description')}") # type: ignore |
| 59 | except Exception as e: |
| 60 | # Fallback: if URL contains '/a2a', try root path without it |
| 61 | if "/a2a" in self.agent_url: |
| 62 | root_url = self.agent_url.split("/a2a", 1)[0] |
| 63 | try: |
| 64 | response = await self._http_client.get(f"{root_url}/.well-known/agent.json") |
| 65 | response.raise_for_status() |
| 66 | self._agent_card = response.json() |
| 67 | _PRINTER.print(f"Retrieved agent card from {root_url}") |
| 68 | except Exception: |
| 69 | pass # swallow, will re-raise below |
| 70 | _PRINTER.print(f"[!] Could not connect to {self.agent_url}\n → Ensure the server is running and reachable.\n → Full error: {e}") |
| 71 | raise RuntimeError(f"Could not retrieve agent card: {e}") |
| 72 | |
| 73 | return self._agent_card # type: ignore |
| 74 | |
| 75 | async def send_message( |
| 76 | self, |
| 77 | message: str, |
| 78 | attachments: Optional[List[str]] = None, |
| 79 | context_id: Optional[str] = None, |
| 80 | metadata: Optional[Dict[str, Any]] = None |
| 81 | ) -> Dict[str, Any]: |
| 82 | """Send a message to the remote agent and return task response.""" |
| 83 | if not self._agent_card: |
| 84 | await self.get_agent_card() |
| 85 | |
| 86 | # Re-use context automatically if caller did not supply one |
| 87 | if context_id is None: |
| 88 | context_id = self._context_id |
| 89 | |
| 90 | # Build message parts |
| 91 | parts = [{'kind': 'text', 'text': message}] |
| 92 | |
| 93 | if attachments: |
| 94 | for attachment in attachments: |
| 95 | file_part = {'kind': 'file', 'file': {'uri': attachment}} |
| 96 | parts.append(file_part) # type: ignore |
| 97 | |
| 98 | # Construct A2A message |
| 99 | a2a_message = { |
| 100 | 'role': 'user', |
| 101 | 'parts': parts, |
| 102 | 'kind': 'message', |
| 103 | 'message_id': str(uuid.uuid4()) |
| 104 | } |
| 105 | |
| 106 | if context_id is not None: |
| 107 | a2a_message['context_id'] = context_id |
| 108 | |
| 109 | # Send using the message/send method (not send_task) |
| 110 | try: |
| 111 | response = await self._a2a_client.send_message( |
| 112 | message=a2a_message, # type: ignore |
| 113 | metadata=metadata, |
| 114 | configuration={'accepted_output_modes': ['application/json', 'text/plain'], 'blocking': True} # type: ignore |
| 115 | ) |
| 116 | |
| 117 | # Persist context id for subsequent calls |
| 118 | try: |
| 119 | ctx = response.get('result', {}).get('context_id') # type: ignore[index] |
| 120 | if isinstance(ctx, str): |
| 121 | self._context_id = ctx |
| 122 | except Exception: |
| 123 | pass # ignore if structure differs |
| 124 | return response # type: ignore |
| 125 | except Exception as e: |
| 126 | _PRINTER.print(f"[A2A] Error sending message: {e}") |
| 127 | raise |
| 128 | |
| 129 | async def get_task(self, task_id: str) -> Dict[str, Any]: |
| 130 | """Get the status and results of a task. |
| 131 | |
| 132 | Args: |
| 133 | task_id: The ID of the task to query |
| 134 | |
| 135 | Returns: |
| 136 | Dictionary containing the task information |
| 137 | """ |
| 138 | try: |
| 139 | response = await self._a2a_client.get_task(task_id) # type: ignore |
| 140 | return response # type: ignore |
| 141 | except Exception as e: |
| 142 | _PRINTER.print(f"Failed to get task {task_id}: {e}") |
| 143 | raise RuntimeError(f"Failed to get task: {e}") |
| 144 | |
| 145 | async def wait_for_completion(self, task_id: str, poll_interval: int = 2, max_wait: int = 300) -> Dict[str, Any]: |
| 146 | """Wait for a task to complete and return the final result. |
| 147 | |
| 148 | Args: |
| 149 | task_id: The ID of the task to wait for |
| 150 | poll_interval: How often to check task status (seconds) |
| 151 | max_wait: Maximum time to wait (seconds) |
| 152 | |
| 153 | Returns: |
| 154 | Dictionary containing the completed task information |
| 155 | """ |
| 156 | import asyncio |
| 157 | |
| 158 | waited = 0 |
| 159 | while waited < max_wait: |
| 160 | task_info = await self.get_task(task_id) |
| 161 | |
| 162 | if 'result' in task_info: |
| 163 | task = task_info['result'] |
| 164 | status = task.get('status', {}) |
| 165 | state = status.get('state', 'unknown') |
| 166 | |
| 167 | if state in ['completed', 'failed', 'canceled']: |
| 168 | _PRINTER.print(f"Task {task_id} finished with state: {state}") |
| 169 | return task_info |
| 170 | else: |
| 171 | _PRINTER.print(f"Task {task_id} status: {state}") |
| 172 | |
| 173 | await asyncio.sleep(poll_interval) |
| 174 | waited += poll_interval |
| 175 | |
| 176 | raise TimeoutError(f"Task {task_id} did not complete within {max_wait} seconds") |
| 177 | |
| 178 | async def close(self): |
| 179 | """Close the HTTP client connection.""" |
| 180 | await self._http_client.aclose() |
| 181 | |
| 182 | async def __aenter__(self): |
| 183 | """Async context manager entry.""" |
| 184 | return self |
| 185 | |
| 186 | async def __aexit__(self, exc_type, exc_val, exc_tb): |
| 187 | """Async context manager exit.""" |
| 188 | await self.close() |
| 189 | |
| 190 | |
| 191 | async def connect_to_agent(agent_url: str, timeout: int = 30) -> AgentConnection: |
| 192 | """Create a connection to a remote agent. |
| 193 | |
| 194 | Args: |
| 195 | agent_url: The base URL of the agent |
| 196 | timeout: Request timeout in seconds |
| 197 | |
| 198 | Returns: |
| 199 | AgentConnection instance |
| 200 | """ |
| 201 | connection = AgentConnection(agent_url, timeout) |
| 202 | # Verify connection by retrieving agent card |
| 203 | await connection.get_agent_card() |
| 204 | return connection |
| 205 | |
| 206 | |
| 207 | def is_client_available() -> bool: |
| 208 | """Check if FastA2A client is available.""" |
| 209 | return FASTA2A_CLIENT_AVAILABLE |