269 lines
8.6 KiB
Python
269 lines
8.6 KiB
Python
# pyright: reportMissingImports=false
|
|
import asyncio
|
|
import uuid
|
|
from typing import Any
|
|
|
|
import httpx
|
|
from rich.console import Console
|
|
from rich.markdown import Markdown
|
|
from rich.prompt import Prompt
|
|
|
|
try:
|
|
from fasta2a.client import A2AClient as FastA2AClient
|
|
from fasta2a.schema import Message, TextPart
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
"A2A support requires the 'a2a' extra. "
|
|
"Install with: uv pip install 'haiku.rag[a2a]'"
|
|
) from e
|
|
|
|
|
|
class A2AClient:
|
|
"""Interactive A2A protocol client."""
|
|
|
|
def __init__(self, base_url: str = "http://localhost:8000"):
|
|
"""Initialize A2A client.
|
|
|
|
Args:
|
|
base_url: Base URL of the A2A server
|
|
"""
|
|
self.base_url = base_url.rstrip("/")
|
|
http_client = httpx.AsyncClient(timeout=60.0)
|
|
self._client = FastA2AClient(base_url=base_url, http_client=http_client)
|
|
|
|
async def close(self):
|
|
"""Close the HTTP client."""
|
|
await self._client.http_client.aclose()
|
|
|
|
async def get_agent_card(self) -> dict[str, Any]:
|
|
"""Fetch the agent card from the A2A server.
|
|
|
|
Returns:
|
|
Agent card dictionary with agent capabilities and metadata
|
|
"""
|
|
response = await self._client.http_client.get(
|
|
f"{self.base_url}/.well-known/agent-card.json"
|
|
)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
|
|
async def send_message(
|
|
self,
|
|
text: str,
|
|
context_id: str | None = None,
|
|
skill_id: str | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Send a message to the A2A agent and wait for completion.
|
|
|
|
Args:
|
|
text: Message text to send
|
|
context_id: Optional conversation context ID (creates new if None)
|
|
skill_id: Optional skill ID to use (defaults to document-qa)
|
|
|
|
Returns:
|
|
Completed task with response messages and artifacts
|
|
"""
|
|
if context_id is None:
|
|
context_id = str(uuid.uuid4())
|
|
|
|
message = Message(
|
|
kind="message",
|
|
role="user",
|
|
message_id=str(uuid.uuid4()),
|
|
parts=[TextPart(kind="text", text=text)],
|
|
)
|
|
|
|
metadata: dict[str, Any] = {"contextId": context_id}
|
|
if skill_id:
|
|
metadata["skillId"] = skill_id
|
|
|
|
response = await self._client.send_message(message, metadata=metadata)
|
|
|
|
if "error" in response:
|
|
return {"error": response["error"]}
|
|
|
|
result = response.get("result")
|
|
if not result:
|
|
return {"result": result}
|
|
|
|
# Result can be either Task or Message - check if it's a Task with an id
|
|
if result.get("kind") == "task":
|
|
task_id = result.get("id")
|
|
if task_id:
|
|
# Poll for task completion
|
|
return await self.wait_for_task(task_id)
|
|
|
|
# Return the message directly
|
|
return {"result": result}
|
|
|
|
async def wait_for_task(
|
|
self, task_id: str, max_wait: int = 120, poll_interval: float = 0.5
|
|
) -> dict[str, Any]:
|
|
"""Poll for task completion.
|
|
|
|
Args:
|
|
task_id: Task ID to poll for
|
|
max_wait: Maximum time to wait in seconds
|
|
poll_interval: Interval between polls in seconds
|
|
|
|
Returns:
|
|
Completed task result
|
|
"""
|
|
import time
|
|
|
|
start_time = time.time()
|
|
|
|
while time.time() - start_time < max_wait:
|
|
task_response = await self._client.get_task(task_id)
|
|
|
|
if "error" in task_response:
|
|
return {"error": task_response["error"]}
|
|
|
|
task = task_response.get("result")
|
|
if not task:
|
|
raise Exception("No task in response")
|
|
|
|
state = task.get("status", {}).get("state")
|
|
|
|
if state == "completed":
|
|
return {"result": task}
|
|
elif state == "failed":
|
|
raise Exception(f"Task failed: {task}")
|
|
|
|
await asyncio.sleep(poll_interval)
|
|
|
|
raise TimeoutError(f"Task {task_id} did not complete within {max_wait}s")
|
|
|
|
|
|
def print_agent_card(card: dict[str, Any], console: Console):
|
|
"""Pretty print the agent card using Rich."""
|
|
console.print()
|
|
console.print("[bold]Agent Card[/bold]")
|
|
console.rule()
|
|
|
|
console.print(f" [repr.attrib_name]name[/repr.attrib_name]: {card.get('name')}")
|
|
console.print(
|
|
f" [repr.attrib_name]description[/repr.attrib_name]: {card.get('description')}"
|
|
)
|
|
console.print(
|
|
f" [repr.attrib_name]version[/repr.attrib_name]: {card.get('version')}"
|
|
)
|
|
console.print(
|
|
f" [repr.attrib_name]protocol version[/repr.attrib_name]: {card.get('protocolVersion')}"
|
|
)
|
|
|
|
skills = card.get("skills", [])
|
|
console.print(f"\n[bold cyan]Skills ({len(skills)}):[/bold cyan]")
|
|
for skill in skills:
|
|
console.print(f" • {skill.get('id')}: {skill.get('name')}")
|
|
console.print(f" [dim]{skill.get('description')}[/dim]")
|
|
examples = skill.get("examples", [])
|
|
if examples:
|
|
console.print(f" [dim]Examples: {', '.join(examples[:2])}[/dim]")
|
|
console.print()
|
|
|
|
|
|
def print_response(response: dict[str, Any], console: Console):
|
|
"""Pretty print the A2A response using Rich."""
|
|
if "error" in response:
|
|
console.print(f"[red]Error: {response['error']}[/red]")
|
|
return
|
|
|
|
result = response.get("result", {})
|
|
|
|
# Get messages from history and artifacts from completed task
|
|
history = result.get("history", [])
|
|
artifacts = result.get("artifacts", [])
|
|
|
|
# Print agent messages from history with markdown rendering
|
|
for msg in history:
|
|
if msg.get("role") == "agent":
|
|
for part in msg.get("parts", []):
|
|
if part.get("kind") == "text":
|
|
text = part.get("text", "")
|
|
# Render as markdown
|
|
console.print()
|
|
console.print("[bold green]Answer:[/bold green]")
|
|
console.print(Markdown(text))
|
|
|
|
# Print artifacts summary with details
|
|
if artifacts:
|
|
console.rule("[dim]Artifacts generated[/dim]")
|
|
summary_lines = []
|
|
|
|
for artifact in artifacts:
|
|
name = artifact.get("name", "")
|
|
parts = artifact.get("parts", [])
|
|
|
|
if name == "search_results" and parts:
|
|
data = parts[0].get("data", {})
|
|
query = data.get("query", "")
|
|
results = data.get("results", [])
|
|
summary_lines.append(f"🔍 search: '{query}' ({len(results)} results)")
|
|
|
|
elif name == "document" and parts:
|
|
part = parts[0]
|
|
if part.get("kind") == "text":
|
|
text = part.get("text", "")
|
|
length = len(text)
|
|
summary_lines.append(f"📄 document ({length} chars)")
|
|
|
|
elif name == "qa_result" and parts:
|
|
data = parts[0].get("data", {})
|
|
skill = data.get("skill", "unknown")
|
|
summary_lines.append(f"💬 {skill}")
|
|
|
|
if summary_lines:
|
|
console.print(f"[dim]{' • '.join(summary_lines)}[/dim]")
|
|
|
|
console.print()
|
|
|
|
|
|
async def run_interactive_client(url: str = "http://localhost:8000"):
|
|
"""Run the interactive A2A client.
|
|
|
|
Args:
|
|
url: Base URL of the A2A server
|
|
"""
|
|
console = Console()
|
|
client = A2AClient(url)
|
|
|
|
console.print("[bold]haiku.rag A2A interactive client[/bold]")
|
|
console.print()
|
|
|
|
# Fetch and display agent card
|
|
console.print("[dim]Fetching agent card...[/dim]")
|
|
try:
|
|
card = await client.get_agent_card()
|
|
print_agent_card(card, console)
|
|
except Exception as e:
|
|
console.print(f"[red]Error fetching agent card: {e}[/red]")
|
|
await client.close()
|
|
return
|
|
|
|
# Create a conversation context
|
|
context_id = str(uuid.uuid4())
|
|
console.print(f"[dim]context id: {context_id}[/dim]")
|
|
console.print("[dim]Type your questions (or 'quit' to exit)[/dim]\n")
|
|
|
|
try:
|
|
while True:
|
|
try:
|
|
question = Prompt.ask("[bold blue]Question[/bold blue]").strip()
|
|
if not question:
|
|
continue
|
|
|
|
if question.lower() in ("quit", "exit", "q"):
|
|
console.print("\n[dim]Goodbye![/dim]")
|
|
break
|
|
|
|
response = await client.send_message(question, context_id=context_id)
|
|
print_response(response, console)
|
|
|
|
except KeyboardInterrupt:
|
|
console.print("\n\n[dim]Exiting...[/dim]")
|
|
break
|
|
except Exception as e:
|
|
console.print(f"\n[red]Error: {e}[/red]\n")
|
|
finally:
|
|
await client.close()
|