mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Evan Mattson
35a8565495
* Add AG-UI integration * Fix tests. PR feedback * Cleanup * PR Feedback * Improve README and getting started experience * Fix links
123 lines
4.0 KiB
Python
123 lines
4.0 KiB
Python
# Copyright (c) Microsoft. All rights reserved.
|
|
|
|
"""AG-UI client example."""
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
from collections.abc import AsyncIterator
|
|
|
|
import httpx
|
|
|
|
|
|
class AGUIClient:
|
|
"""Simple AG-UI protocol client."""
|
|
|
|
def __init__(self, server_url: str):
|
|
"""Initialize the client.
|
|
|
|
Args:
|
|
server_url: The AG-UI server endpoint URL
|
|
"""
|
|
self.server_url = server_url
|
|
self.thread_id: str | None = None
|
|
|
|
async def send_message(self, message: str) -> AsyncIterator[dict]:
|
|
"""Send a message and stream the response.
|
|
|
|
Args:
|
|
message: The user message to send
|
|
|
|
Yields:
|
|
AG-UI events from the server
|
|
"""
|
|
# Prepare the request
|
|
request_data: dict[str, object] = {
|
|
"messages": [
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{"role": "user", "content": message},
|
|
]
|
|
}
|
|
|
|
# Include thread_id if we have one (for conversation continuity)
|
|
if self.thread_id:
|
|
request_data["thread_id"] = self.thread_id
|
|
|
|
# Stream the response
|
|
async with httpx.AsyncClient(timeout=60.0) as client:
|
|
async with client.stream(
|
|
"POST",
|
|
self.server_url,
|
|
json=request_data,
|
|
headers={"Accept": "text/event-stream"},
|
|
) as response:
|
|
response.raise_for_status()
|
|
|
|
async for line in response.aiter_lines():
|
|
# Parse Server-Sent Events format
|
|
if line.startswith("data: "):
|
|
data = line[6:] # Remove "data: " prefix
|
|
try:
|
|
event = json.loads(data)
|
|
yield event
|
|
|
|
# Capture thread_id from RUN_STARTED event
|
|
if event.get("type") == "RUN_STARTED" and not self.thread_id:
|
|
self.thread_id = event.get("threadId")
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
|
|
async def main():
|
|
"""Main client loop."""
|
|
# Get server URL from environment or use default
|
|
server_url = os.environ.get("AGUI_SERVER_URL", "http://127.0.0.1:5100/")
|
|
print(f"Connecting to AG-UI server at: {server_url}\n")
|
|
|
|
client = AGUIClient(server_url)
|
|
|
|
try:
|
|
while True:
|
|
# Get user input
|
|
message = input("\nUser (:q or quit to exit): ")
|
|
if not message.strip():
|
|
print("Request cannot be empty.")
|
|
continue
|
|
|
|
if message.lower() in (":q", "quit"):
|
|
break
|
|
|
|
# Send message and display streaming response
|
|
print("\n", end="")
|
|
async for event in client.send_message(message):
|
|
event_type = event.get("type", "")
|
|
|
|
if event_type == "RUN_STARTED":
|
|
thread_id = event.get("threadId", "")
|
|
run_id = event.get("runId", "")
|
|
print(f"\033[93m[Run Started - Thread: {thread_id}, Run: {run_id}]\033[0m")
|
|
|
|
elif event_type == "TEXT_MESSAGE_CONTENT":
|
|
# Stream text content in cyan
|
|
print(f"\033[96m{event.get('delta', '')}\033[0m", end="", flush=True)
|
|
|
|
elif event_type == "RUN_FINISHED":
|
|
thread_id = event.get("threadId", "")
|
|
run_id = event.get("runId", "")
|
|
print(f"\n\033[92m[Run Finished - Thread: {thread_id}, Run: {run_id}]\033[0m")
|
|
|
|
elif event_type == "RUN_ERROR":
|
|
error_message = event.get("message", "Unknown error")
|
|
print(f"\n\033[91m[Run Error - Message: {error_message}]\033[0m")
|
|
|
|
print()
|
|
|
|
except KeyboardInterrupt:
|
|
print("\n\nExiting...")
|
|
except Exception as e:
|
|
print(f"\n\033[91mAn error occurred: {e}\033[0m")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|