Python: AG-UI protocol support (#1826)

* Add AG-UI integration * Fix tests. PR feedback * Cleanup * PR Feedback * Improve README and getting started experience * Fix links
2026-06-16 21:04:09 +08:00 · 2025-11-05 14:25:24 +09:00
parent 0c862e97a6
commit 35a8565495
51 changed files with 7677 additions and 163 deletions
@@ -0,0 +1,122 @@
+# 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())