Spaces:
Running
Running
File size: 5,072 Bytes
6918275 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | from typing import Any
from helpers.tool import Tool, Response
from helpers.print_style import PrintStyle
from helpers.fasta2a_client import connect_to_agent, is_client_available
A2A_EMPTY_RESPONSE_ERROR = (
"A2A chat failed: the remote task completed but no assistant text was found. "
"Expected final.result.history to include an assistant message with a text "
"part, or a text artifact/status message. Treat this as a failed remote "
"response, not success."
)
def _session_key(agent_url: str) -> str:
"""Keep root and explicit /a2a URLs in the same conversation cache."""
normalized = agent_url.rstrip("/")
if normalized.endswith("/a2a"):
return normalized[:-4].rstrip("/")
return normalized
def _text_from_part(part: Any) -> str:
if not isinstance(part, dict):
return ""
for key in ("text", "content"):
value = part.get(key)
if isinstance(value, str) and value.strip():
return value.strip()
return ""
def _text_from_message(message: Any) -> str:
if isinstance(message, str):
return message.strip()
if not isinstance(message, dict):
return ""
parts = message.get("parts")
if isinstance(parts, list):
texts = [_text_from_part(part) for part in parts]
text = "\n".join(text for text in texts if text)
if text:
return text
for key in ("text", "content", "message", "output"):
value = message.get(key)
if isinstance(value, str) and value.strip():
return value.strip()
return ""
def _extract_latest_assistant_text(task_response: Any) -> str:
if not isinstance(task_response, dict):
return ""
result = task_response.get("result", task_response)
if not isinstance(result, dict):
return ""
history = result.get("history")
if isinstance(history, list):
for message in reversed(history):
if isinstance(message, dict) and message.get("role") == "user":
continue
text = _text_from_message(message)
if text:
return text
status = result.get("status")
if isinstance(status, dict):
text = _text_from_message(status.get("message"))
if text:
return text
artifacts = result.get("artifacts")
if isinstance(artifacts, list):
for artifact in reversed(artifacts):
text = _text_from_message(artifact)
if text:
return text
return _text_from_message(result)
class A2AChatTool(Tool):
"""Communicate with another FastA2A-compatible agent."""
async def execute(self, **kwargs):
if not is_client_available():
return Response(message="FastA2A client not available on this instance.", break_loop=False)
agent_url: str | None = kwargs.get("agent_url") # required
user_message: str | None = kwargs.get("message") # required
attachments = kwargs.get("attachments", None) # optional list[str]
reset = bool(kwargs.get("reset", False))
if not agent_url or not isinstance(agent_url, str):
return Response(message="agent_url argument missing", break_loop=False)
if not user_message or not isinstance(user_message, str):
return Response(message="message argument missing", break_loop=False)
# Retrieve or create session cache on the Agent instance
sessions: dict[str, str] = self.agent.get_data("_a2a_sessions") or {}
cache_key = _session_key(agent_url)
# Handle reset flag: start fresh conversation
if reset and cache_key in sessions:
sessions.pop(cache_key, None)
context_id = None if reset else sessions.get(cache_key)
try:
async with await connect_to_agent(agent_url) as conn:
task_resp = await conn.send_message(user_message, attachments=attachments, context_id=context_id)
task_id = task_resp.get("result", {}).get("id") # type: ignore[index]
if not task_id:
return Response(message="Remote agent failed to create task.", break_loop=False)
final = await conn.wait_for_completion(task_id)
new_context_id = final["result"].get("context_id") # type: ignore[index]
if isinstance(new_context_id, str):
sessions[cache_key] = new_context_id
# persist back to agent data
self.agent.set_data("_a2a_sessions", sessions)
assistant_text = _extract_latest_assistant_text(final)
if not assistant_text:
return Response(
message=A2A_EMPTY_RESPONSE_ERROR,
break_loop=False,
)
return Response(message=assistant_text, break_loop=False)
except Exception as e:
PrintStyle.error(f"A2A chat error: {e}")
return Response(message=f"A2A chat error: {e}", break_loop=False)
|