123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105 |
- 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()
|