Add stream text output

2026-02-02 17:54:55 +08:00
parent 32c491cd3f
commit 30306addee
3 changed files with 178 additions and 10 deletions
--- a/examples/stream_chat.py
+++ b/examples/stream_chat.py
@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+"""
+Simple CLI script to interact with /chat endpoint in stream mode.
+"""
+
+import asyncio
+import aiohttp
+import json
+import sys
+from datetime import datetime
+
+
+API_BASE_URL = "http://localhost:8000"
+
+
+async def stream_chat(session_id: str, text: str):
+    """Send a streaming chat request."""
+    timestamp = datetime.now().isoformat()
+
+    payload = {
+        "sessionId": session_id,
+        "timeStamp": timestamp,
+        "text": text
+    }
+
+    async with aiohttp.ClientSession() as http_session:
+        async with http_session.post(
+            f"{API_BASE_URL}/chat",
+            json=payload,
+            params={"stream": "true"},
+        ) as response:
+            print(f"Status: {response.status}")
+            print("-" * 50)
+
+            buffer = ""
+            # Use async iterator on response.content
+            async for chunk in response.content.iter_chunked(1024):
+                chunk_str = chunk.decode("utf-8")
+                buffer += chunk_str
+
+                # Process complete SSE messages
+                while "\n\n" in buffer:
+                    message, buffer = buffer.split("\n\n", 1)
+                    if message.startswith("event: "):
+                        event_type = message[7:].split("\n")[0]
+                        data_line = message.split("data: ", 1)[-1]
+                        try:
+                            data = json.loads(data_line)
+                            print(f"[{event_type}] {json.dumps(data, ensure_ascii=False)}")
+                        except json.JSONDecodeError:
+                            print(f"[{event_type}] {data_line}")
+                    elif message.startswith("data: "):
+                        data_str = message[6:]
+                        try:
+                            data = json.loads(data_str)
+                            print(f"[data] {json.dumps(data, ensure_ascii=False)}")
+                        except json.JSONDecodeError:
+                            print(f"[data] {data_str}")
+
+
+async def main():
+    if len(sys.argv) < 3:
+        print("Usage: python stream_chat.py <session_id> <message>")
+        print("Example: python stream_chat.py test-session-123 '发生了交通事故'")
+        sys.exit(1)
+
+    session_id = sys.argv[1]
+    text = " ".join(sys.argv[2:])
+
+    print(f"Session ID: {session_id}")
+    print(f"Message: {text}")
+    print("-" * 50)
+
+    await stream_chat(session_id, text)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())