import requests import json import sys import os from typing import Iterator class SSEClient: def __init__(self, response): self.response = response self.buffer = "" def events(self) -> Iterator[str]: """解析和生成SSE事件""" for chunk in self.response.iter_lines(decode_unicode=True): if chunk: if chunk.startswith("data:"): yield chunk[5:].strip() # 移除 "data:" 前缀 def print_with_flush(text: str, end: str = ""): """实现打字机效果的打印函数""" if text is None or " completed in " in text: return print(text, end=end, flush=True) def format_state_info(state_info: dict) -> str: """格式化状态信息""" return ( "\n当前已收集的信息:\n" f"姓名: {state_info.get('姓名', '未收集')}\n" f"年龄: {state_info.get('年龄', '未收集')}\n" f"感兴趣的行业: {state_info.get('感兴趣的行业', '未收集')}\n" ) def chat_with_agent(): """与Agent进行SSE对话""" url = "http://localhost:8000/chat/sse" user_id = "test_user" headers = { 'Accept': 'text/event-stream', 'Cache-Control': 'no-cache', 'Connection': 'keep-alive' } print("欢迎使用信息收集助手!(输入 'exit' 或 'quit' 退出)") # 发送初始问候 try: response = requests.get( f"{url}?prompt=你好&user_id={user_id}", stream=True, headers=headers ) client = SSEClient(response) # 处理初始响应 for event in client.events(): try: data = json.loads(event) if data["type"] == "message": print_with_flush(data["content"]) elif data["type"] == "state": print(format_state_info(data["content"])) except json.JSONDecodeError: print(f"\n警告: 无法解析消息: {event}") # 主对话循环 while True: user_input = input("\n请输入(exit/quit以退出): ").strip() if user_input.lower() in ['exit', 'quit']: print("对话结束。") break if not user_input: print("输入不能为空,请重新输入。") continue # 发送用户输入并获取流式响应 response = requests.get( f"{url}?prompt={user_input}&user_id={user_id}", stream=True, headers=headers ) client = SSEClient(response) # 处理响应 for event in client.events(): try: data = json.loads(event) if data["type"] == "message": print_with_flush(data["content"]) elif data["type"] == "state": print(format_state_info(data["content"])) # elif data["type"] == "error": # print(f"\n错误: {data['content']}") except json.JSONDecodeError: print(f"\n警告: 无法解析消息: {event}") except KeyboardInterrupt: print("\n对话被用户中断。") except Exception as e: print(f"\n发生错误: {str(e)}") finally: print("\n感谢使用!") if __name__ == "__main__": chat_with_agent()