Interactive client
This commit is contained in:
parent
c562aa0551
commit
34b3e9969f
7 changed files with 386 additions and 2 deletions
|
|
@ -149,10 +149,14 @@ Provides tools for document management and search directly in your AI assistant.
|
||||||
Run as a conversational agent with the Agent-to-Agent protocol:
|
Run as a conversational agent with the Agent-to-Agent protocol:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
|
# Start the A2A server
|
||||||
haiku-rag serve --a2a
|
haiku-rag serve --a2a
|
||||||
|
|
||||||
|
# Connect with the interactive client (in another terminal)
|
||||||
|
haiku-rag a2aclient
|
||||||
```
|
```
|
||||||
|
|
||||||
Provides a conversational interface with:
|
The A2A agent provides:
|
||||||
- Multi-turn dialogue with context
|
- Multi-turn dialogue with context
|
||||||
- Intelligent multi-search for complex questions
|
- Intelligent multi-search for complex questions
|
||||||
- Source citations with titles and URIs
|
- Source citations with titles and URIs
|
||||||
|
|
|
||||||
26
docs/a2a.md
26
docs/a2a.md
|
|
@ -29,6 +29,32 @@ Example:
|
||||||
haiku-rag serve --a2a --a2a-host 0.0.0.0 --a2a-port 8080
|
haiku-rag serve --a2a --a2a-host 0.0.0.0 --a2a-port 8080
|
||||||
```
|
```
|
||||||
|
|
||||||
|
## Interactive A2A Client
|
||||||
|
|
||||||
|
Test and interact with haiku.rag's A2A server using the built-in interactive client:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
haiku-rag a2aclient
|
||||||
|
```
|
||||||
|
|
||||||
|
Client options:
|
||||||
|
- `--url` - Base URL of the A2A server (default: http://localhost:8000)
|
||||||
|
|
||||||
|
Example:
|
||||||
|
```bash
|
||||||
|
# Connect to local server
|
||||||
|
haiku-rag a2aclient
|
||||||
|
|
||||||
|
# Connect to remote server
|
||||||
|
haiku-rag a2aclient --url https://example.com:8000
|
||||||
|
```
|
||||||
|
|
||||||
|
The interactive client provides:
|
||||||
|
- Rich markdown rendering of agent responses
|
||||||
|
- Conversation context across multiple turns
|
||||||
|
- Agent card discovery and display
|
||||||
|
- Compact artifact summaries
|
||||||
|
|
||||||
## Requirements
|
## Requirements
|
||||||
|
|
||||||
A2A support requires the `a2a` extra:
|
A2A support requires the `a2a` extra:
|
||||||
|
|
|
||||||
20
docs/cli.md
20
docs/cli.md
|
|
@ -148,6 +148,26 @@ haiku-rag serve --mcp --mcp-port 9000 --a2a --a2a-port 9001
|
||||||
|
|
||||||
See [Server Mode](server.md) for details on available services.
|
See [Server Mode](server.md) for details on available services.
|
||||||
|
|
||||||
|
### A2A Interactive Client
|
||||||
|
|
||||||
|
Connect to and chat with haiku.rag's A2A server:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Connect to local server
|
||||||
|
haiku-rag a2aclient
|
||||||
|
|
||||||
|
# Connect to remote server
|
||||||
|
haiku-rag a2aclient --url https://example.com:8000
|
||||||
|
```
|
||||||
|
|
||||||
|
The interactive client provides:
|
||||||
|
- Rich markdown rendering of agent responses
|
||||||
|
- Multi-turn conversation with context
|
||||||
|
- Agent card discovery and display
|
||||||
|
- Compact artifact summaries
|
||||||
|
|
||||||
|
See [A2A documentation](a2a.md) for more details.
|
||||||
|
|
||||||
## Settings
|
## Settings
|
||||||
|
|
||||||
View current configuration settings:
|
View current configuration settings:
|
||||||
|
|
|
||||||
|
|
@ -73,7 +73,7 @@ Configure which LLM provider to use for question answering. Any provider and mod
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
QA_PROVIDER="ollama"
|
QA_PROVIDER="ollama"
|
||||||
QA_MODEL="qwen3"
|
QA_MODEL="gpt-oss"
|
||||||
OLLAMA_BASE_URL="http://localhost:11434"
|
OLLAMA_BASE_URL="http://localhost:11434"
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|
|
||||||
41
examples/README.md
Normal file
41
examples/README.md
Normal file
|
|
@ -0,0 +1,41 @@
|
||||||
|
# haiku.rag Examples
|
||||||
|
|
||||||
|
This directory contains example scripts demonstrating various features of haiku.rag.
|
||||||
|
|
||||||
|
## A2A Security Examples
|
||||||
|
|
||||||
|
**Directory:** `a2a-security/`
|
||||||
|
|
||||||
|
Three examples showing how to add authentication to haiku.rag's A2A server:
|
||||||
|
|
||||||
|
### API Key Authentication
|
||||||
|
|
||||||
|
**File:** `a2a-security/apikey_example.py`
|
||||||
|
|
||||||
|
Simple header-based authentication suitable for internal services and development.
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python examples/a2a-security/apikey_example.py /path/to/database.lancedb
|
||||||
|
```
|
||||||
|
|
||||||
|
### OAuth2 GitHub Authentication
|
||||||
|
|
||||||
|
**File:** `a2a-security/oauth2_github.py`
|
||||||
|
|
||||||
|
GitHub Personal Access Token authentication for GitHub-integrated services.
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python examples/a2a-security/oauth2_github.py /path/to/database.lancedb
|
||||||
|
```
|
||||||
|
|
||||||
|
### OAuth2 Enterprise Authentication
|
||||||
|
|
||||||
|
**File:** `a2a-security/oauth2_example.py`
|
||||||
|
|
||||||
|
Full OAuth2 with JWT verification for enterprise environments.
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python examples/a2a-security/oauth2_example.py /path/to/database.lancedb
|
||||||
|
```
|
||||||
|
|
||||||
|
See individual files for detailed setup instructions and usage examples.
|
||||||
271
src/haiku/rag/a2a/client.py
Normal file
271
src/haiku/rag/a2a/client.py
Normal file
|
|
@ -0,0 +1,271 @@
|
||||||
|
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
|
||||||
|
|
||||||
|
|
||||||
|
class A2AClient:
|
||||||
|
"""Simple 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("/")
|
||||||
|
self.client = httpx.AsyncClient(timeout=60.0)
|
||||||
|
|
||||||
|
async def close(self):
|
||||||
|
"""Close the HTTP client."""
|
||||||
|
await self.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.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_id = str(uuid.uuid4())
|
||||||
|
|
||||||
|
payload: dict[str, Any] = {
|
||||||
|
"jsonrpc": "2.0",
|
||||||
|
"method": "message/send",
|
||||||
|
"params": {
|
||||||
|
"contextId": context_id,
|
||||||
|
"message": {
|
||||||
|
"kind": "message",
|
||||||
|
"role": "user",
|
||||||
|
"messageId": message_id,
|
||||||
|
"parts": [{"kind": "text", "text": text}],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"id": 1,
|
||||||
|
}
|
||||||
|
|
||||||
|
if skill_id:
|
||||||
|
payload["params"]["skillId"] = skill_id
|
||||||
|
|
||||||
|
response = await self.client.post(
|
||||||
|
self.base_url,
|
||||||
|
json=payload,
|
||||||
|
headers={"Content-Type": "application/json"},
|
||||||
|
)
|
||||||
|
response.raise_for_status()
|
||||||
|
initial_response = response.json()
|
||||||
|
|
||||||
|
# Extract task ID from response
|
||||||
|
result = initial_response.get("result", {})
|
||||||
|
task_id = result.get("id")
|
||||||
|
|
||||||
|
if not task_id:
|
||||||
|
return initial_response
|
||||||
|
|
||||||
|
# Poll for task completion
|
||||||
|
return await self.wait_for_task(task_id)
|
||||||
|
|
||||||
|
async def wait_for_task(
|
||||||
|
self, task_id: str, max_wait: int = 60, 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:
|
||||||
|
payload = {
|
||||||
|
"jsonrpc": "2.0",
|
||||||
|
"method": "tasks/get",
|
||||||
|
"params": {"id": task_id},
|
||||||
|
"id": 2,
|
||||||
|
}
|
||||||
|
|
||||||
|
response = await self.client.post(
|
||||||
|
self.base_url,
|
||||||
|
json=payload,
|
||||||
|
headers={"Content-Type": "application/json"},
|
||||||
|
)
|
||||||
|
response.raise_for_status()
|
||||||
|
task = response.json()
|
||||||
|
|
||||||
|
result = task.get("result", {})
|
||||||
|
status = result.get("status", {})
|
||||||
|
state = status.get("state")
|
||||||
|
|
||||||
|
if state == "completed":
|
||||||
|
return 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:
|
||||||
|
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()
|
||||||
|
|
@ -460,5 +460,27 @@ def migrate(
|
||||||
raise typer.Exit(1)
|
raise typer.Exit(1)
|
||||||
|
|
||||||
|
|
||||||
|
@cli.command(
|
||||||
|
"a2aclient", help="Run interactive client to chat with haiku.rag's A2A server"
|
||||||
|
)
|
||||||
|
def a2aclient(
|
||||||
|
url: str = typer.Option(
|
||||||
|
"http://localhost:8000",
|
||||||
|
"--url",
|
||||||
|
help="Base URL of the A2A server",
|
||||||
|
),
|
||||||
|
):
|
||||||
|
try:
|
||||||
|
from haiku.rag.a2a.client import run_interactive_client
|
||||||
|
except ImportError:
|
||||||
|
typer.echo(
|
||||||
|
"Error: A2A support requires the 'a2a' extra. "
|
||||||
|
"Install with: uv pip install 'haiku.rag[a2a]'"
|
||||||
|
)
|
||||||
|
raise typer.Exit(1)
|
||||||
|
|
||||||
|
asyncio.run(run_interactive_client(url=url))
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
cli()
|
cli()
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue