101 lines
3.8 KiB
Python
Executable file
101 lines
3.8 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
"""Test script for AG-UI server functionality."""
|
|
|
|
import asyncio
|
|
import json
|
|
|
|
import httpx
|
|
|
|
|
|
async def test_agui_server():
|
|
"""Test the AG-UI server endpoint."""
|
|
base_url = "http://localhost:8000"
|
|
|
|
# Test health check
|
|
async with httpx.AsyncClient() as client:
|
|
print("Testing health check...")
|
|
response = await client.get(f"{base_url}/health")
|
|
print(f"Health check response: {response.status_code}")
|
|
print(f"Health check data: {response.json()}")
|
|
|
|
# Test AG-UI streaming endpoint
|
|
print("\nTesting AG-UI stream endpoint...")
|
|
request_data = {
|
|
"threadId": "test-thread-1",
|
|
"runId": "test-run-1",
|
|
"state": {"question": "What is pydantic-graph?"},
|
|
"messages": [],
|
|
"config": {},
|
|
}
|
|
|
|
print(f"Request data: {json.dumps(request_data, indent=2)}")
|
|
|
|
# Send request and stream response
|
|
async with client.stream(
|
|
"POST",
|
|
f"{base_url}/v1/agent/stream",
|
|
json=request_data,
|
|
timeout=120.0,
|
|
) as response:
|
|
print(f"Response status: {response.status_code}")
|
|
print("Streaming events...\n")
|
|
|
|
event_count = 0
|
|
async for line in response.aiter_lines():
|
|
if line.startswith("data: "):
|
|
event_data = line[6:] # Remove "data: " prefix
|
|
try:
|
|
event = json.loads(event_data)
|
|
event_type = event.get("type", "UNKNOWN")
|
|
print(f"Event {event_count}: {event_type}")
|
|
|
|
# Show specific event details
|
|
if event_type == "RUN_STARTED":
|
|
print(f" Thread ID: {event.get('threadId')}")
|
|
print(f" Run ID: {event.get('runId')}")
|
|
elif event_type == "STEP_STARTED":
|
|
print(f" Step: {event.get('stepName')}")
|
|
elif event_type == "ACTIVITY_SNAPSHOT":
|
|
print(f" Activity: {event.get('content')}")
|
|
elif event_type == "STATE_SNAPSHOT":
|
|
state = event.get("snapshot", {})
|
|
if "context" in state:
|
|
context = state["context"]
|
|
if "sub_questions" in context:
|
|
num_questions = len(context["sub_questions"])
|
|
print(f" Sub-questions: {num_questions}")
|
|
elif event_type == "RUN_FINISHED":
|
|
result = event.get("result", {})
|
|
if "title" in result:
|
|
print(f" Report Title: {result['title']}")
|
|
elif event_type == "RUN_ERROR":
|
|
print(f" Error: {event.get('message')}")
|
|
|
|
event_count += 1
|
|
except json.JSONDecodeError as e:
|
|
print(f"Failed to parse event: {e}")
|
|
print(f"Raw line: {line}")
|
|
|
|
print(f"\nTotal events received: {event_count}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
print("AG-UI Server Test")
|
|
print("=" * 50)
|
|
print("Make sure to start the server first with:")
|
|
print(" haiku-rag serve --agui --agui-port 8000")
|
|
print("=" * 50)
|
|
print()
|
|
|
|
try:
|
|
asyncio.run(test_agui_server())
|
|
except httpx.ConnectError:
|
|
print("ERROR: Could not connect to server at http://localhost:8000")
|
|
print("Make sure the AG-UI server is running.")
|
|
except KeyboardInterrupt:
|
|
print("\nTest interrupted by user")
|
|
except Exception as e:
|
|
print(f"Test failed with error: {e}")
|
|
import traceback
|
|
|
|
traceback.print_exc()
|