theme5_ex4_sse_client.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105
  1. import requests
  2. import json
  3. import sys
  4. import os
  5. from typing import Iterator
  6. class SSEClient:
  7. def __init__(self, response):
  8. self.response = response
  9. self.buffer = ""
  10. def events(self) -> Iterator[str]:
  11. """解析和生成SSE事件"""
  12. for chunk in self.response.iter_lines(decode_unicode=True):
  13. if chunk:
  14. if chunk.startswith("data:"):
  15. yield chunk[5:].strip() # 移除 "data:" 前缀
  16. def print_with_flush(text: str, end: str = ""):
  17. """实现打字机效果的打印函数"""
  18. if text is None or " completed in " in text:
  19. return
  20. print(text, end=end, flush=True)
  21. def format_state_info(state_info: dict) -> str:
  22. """格式化状态信息"""
  23. return (
  24. "\n当前已收集的信息:\n"
  25. f"姓名: {state_info.get('姓名', '未收集')}\n"
  26. f"年龄: {state_info.get('年龄', '未收集')}\n"
  27. f"感兴趣的行业: {state_info.get('感兴趣的行业', '未收集')}\n"
  28. )
  29. def chat_with_agent():
  30. """与Agent进行SSE对话"""
  31. url = "http://localhost:8000/chat/sse"
  32. user_id = "test_user"
  33. headers = {
  34. 'Accept': 'text/event-stream',
  35. 'Cache-Control': 'no-cache',
  36. 'Connection': 'keep-alive'
  37. }
  38. print("欢迎使用信息收集助手!(输入 'exit' 或 'quit' 退出)")
  39. # 发送初始问候
  40. try:
  41. response = requests.get(
  42. f"{url}?prompt=你好&user_id={user_id}",
  43. stream=True,
  44. headers=headers
  45. )
  46. client = SSEClient(response)
  47. # 处理初始响应
  48. for event in client.events():
  49. try:
  50. data = json.loads(event)
  51. if data["type"] == "message":
  52. print_with_flush(data["content"])
  53. elif data["type"] == "state":
  54. print(format_state_info(data["content"]))
  55. except json.JSONDecodeError:
  56. print(f"\n警告: 无法解析消息: {event}")
  57. # 主对话循环
  58. while True:
  59. user_input = input("\n请输入(exit/quit以退出): ").strip()
  60. if user_input.lower() in ['exit', 'quit']:
  61. print("对话结束。")
  62. break
  63. if not user_input:
  64. print("输入不能为空,请重新输入。")
  65. continue
  66. # 发送用户输入并获取流式响应
  67. response = requests.get(
  68. f"{url}?prompt={user_input}&user_id={user_id}",
  69. stream=True,
  70. headers=headers
  71. )
  72. client = SSEClient(response)
  73. # 处理响应
  74. for event in client.events():
  75. try:
  76. data = json.loads(event)
  77. if data["type"] == "message":
  78. print_with_flush(data["content"])
  79. elif data["type"] == "state":
  80. print(format_state_info(data["content"]))
  81. # elif data["type"] == "error":
  82. # print(f"\n错误: {data['content']}")
  83. except json.JSONDecodeError:
  84. print(f"\n警告: 无法解析消息: {event}")
  85. except KeyboardInterrupt:
  86. print("\n对话被用户中断。")
  87. except Exception as e:
  88. print(f"\n发生错误: {str(e)}")
  89. finally:
  90. print("\n感谢使用!")
  91. if __name__ == "__main__":
  92. chat_with_agent()