theme5_ex4_fastapi.py 9.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231
  1. import os
  2. from dotenv import load_dotenv
  3. from agno.agent import Agent, RunResponseEvent
  4. from agno.tools import tool
  5. from agno.memory.v2.db.sqlite import SqliteMemoryDb
  6. from agno.memory.v2.memory import Memory
  7. from agno.storage.sqlite import SqliteStorage
  8. from agno.models.openai.like import OpenAILike
  9. from typing import Iterator, Optional, Dict, Any
  10. from enum import Enum
  11. import json
  12. import uvicorn
  13. from fastapi import FastAPI, Request
  14. from fastapi.responses import StreamingResponse
  15. from sse_starlette.sse import EventSourceResponse
  16. import asyncio
  17. import json
  18. from typing import AsyncGenerator
  19. load_dotenv()
  20. base_url = os.getenv('BAILIAN_API_BASE_URL')
  21. api_key = os.getenv('BAILIAN_API_KEY')
  22. SYSTEM_PROMPT = """
  23. 你是一个信息收集助手。你需要依次收集用户的姓名、年龄和感兴趣的行业。
  24. 如果用户输入其他无关内容,请提醒用户需要先完成信息收集。
  25. 开始对话前,请先介绍自己。
  26. 工具使用规则:
  27. 1. 每当用户提供有效信息时,必须立即使用 add_info 工具保存信息
  28. 2. 使用 get_info 工具检查特定信息是否已收集
  29. 3. 使用 list_info 工具确认所有已收集的信息
  30. 4. 当用户需要重新收集信息时,使用 reset_info 工具清空所有信息,再重新开始收集
  31. 信息收集策略:
  32. 1. 使用开放式问题引导用户提供更多信息
  33. 2. 对模糊的信息进行澄清和确认
  34. 3. 根据上下文,智能地推断和补充相关信息
  35. 4. 保持对话的自然性和连贯性
  36. 5. 务必保证信息的准确性和完整性
  37. 6. 只有当信息符合需要收集的内容的规则时(例如姓名必须是符合姓名规则的,年龄必须是1~100内的数字等等),才将其记录到已收集信息中
  38. 7. 如果用户输入年龄时给出的是出生年份,请帮用户计算实际年龄(今年是2025年)。
  39. 8. 如果用户输入的年龄没有明确的数字,则此输入无效。
  40. 9. 如果用户一次性将所有信息给出,请将其拆分为单独的字段进行处理,不必重复提问。
  41. 对话要求:
  42. 1. 保持友好、专业的语调
  43. 2. 每次回复要简洁明了
  44. 3. 每个信息收集成功后,确认用户输入的内容,格式为:"好的,已经收集到您的(此处为收集的项目)是:",并询问下一个问题。
  45. 4. 当用户询问一些其他问题时,请礼貱地告诉用户,你是一个信息收集助手,不回答其它不相关问题。提醒用户还没收集完成。
  46. 5. 所有问题回答并收集完毕后,请务必要总结用户的所有信息,并告知用户"信息已收集完毕"。
  47. 6. 如果没有对话历史,请先介绍自己,再进行询问。
  48. 7. 如果用户回答与问题无关,请礼貌且友好地提醒用户需要回答相关问题。
  49. 8. 请不要直接将问题内容直接输出,而是要根据用户的回答进行智能化的处理和回复。
  50. 注意事项:
  51. 除非用户输入与信息收集无关,否则一定要调用对应的tools再给予回复。
  52. 除非用户输入与信息收集无关,否则一定要调用对应的tools再给予回复。
  53. 除非用户输入与信息收集无关,否则一定要调用对应的tools再给予回复。
  54. 除非用户输入与信息收集无关,否则一定要调用对应的tools再给予回复。
  55. 除非用户输入与信息收集无关,否则一定要调用对应的tools再给予回复。
  56. 除非用户输入与信息收集无关,否则一定要调用对应的tools再给予回复。
  57. """
  58. # 初始化空的session状态
  59. initial_state = {
  60. "name": None,
  61. "age": None,
  62. "industry": None
  63. }
  64. class DataField(str, Enum):
  65. NAME = "name"
  66. AGE = "age"
  67. INDUSTRY = "industry"
  68. @tool(name="add_info", description="添加或修改用户信息")
  69. def add_info(agent: Agent, field: DataField, value: str) -> str:
  70. if field == DataField.NAME:
  71. agent.session_state["name"] = value
  72. return json.dumps({"result": f"已收集姓名: {value}"})
  73. elif field == DataField.AGE:
  74. try:
  75. age = int(value)
  76. agent.session_state["age"] = age
  77. return json.dumps({"result": f"已收集年龄: {age}"})
  78. except ValueError:
  79. return json.dumps({"result": "年龄输入无效,请输入一个正确的数字。"})
  80. elif field == DataField.INDUSTRY:
  81. agent.session_state["industry"] = value
  82. return json.dumps({"result": f"已收集感兴趣的行业: {value}"})
  83. @tool(name="add_all_info", description="一次性添加所有用户信息")
  84. def add_all_info(agent: Agent, name: str, age: str, industry: str) -> str:
  85. results = []
  86. # 姓名
  87. if name and isinstance(name, str):
  88. agent.session_state["name"] = name
  89. results.append(f"已收集姓名: {name}")
  90. else:
  91. results.append("姓名输入无效")
  92. # 年龄
  93. try:
  94. age_int = int(age)
  95. if 1 <= age_int <= 100:
  96. agent.session_state["age"] = age_int
  97. results.append(f"已收集年龄: {age_int}")
  98. else:
  99. results.append("年龄输入无效")
  100. except Exception:
  101. results.append("年龄输入无效")
  102. # 行业
  103. if industry and isinstance(industry, str):
  104. agent.session_state["industry"] = industry
  105. results.append(f"已收集感兴趣的行业: {industry}")
  106. else:
  107. results.append("行业输入无效")
  108. # 检查是否全部收集
  109. if all([
  110. agent.session_state["name"],
  111. agent.session_state["age"],
  112. agent.session_state["industry"]
  113. ]):
  114. summary = {
  115. "姓名": agent.session_state["name"],
  116. "年龄": agent.session_state["age"],
  117. "感兴趣的行业": agent.session_state["industry"]
  118. }
  119. results.append(f"信息已收集完毕: {summary}")
  120. return json.dumps({"result": results})
  121. @tool(name="get_info", description="获取用户信息")
  122. def get_info(agent:Agent, field: DataField) -> str:
  123. info = agent.session_state.get(field)
  124. if info is None:
  125. info = "未收集" if field != DataField.AGE else "未收集"
  126. return json.dumps({"result": str(info)})
  127. @tool(name="list_info", description="列出所有用户信息")
  128. def list_info(agent:Agent) -> str:
  129. user_info = agent.session_state
  130. info = {
  131. "姓名": user_info.get("name") or "未收集",
  132. "年龄": user_info.get("age") or "未收集",
  133. "感兴趣的行业": user_info.get("industry") or "未收集"
  134. }
  135. return json.dumps({"result": info})
  136. @tool(name="reset_info", description="重置所有用户信息")
  137. def reset_info(agent: Agent) -> str:
  138. agent.session_state["name"] = None
  139. agent.session_state["age"] = None
  140. agent.session_state["industry"] = None
  141. return json.dumps({"result": "所有信息已重置"})
  142. # 配置
  143. user_id = "test"
  144. session_id = "325"
  145. db_file = "tmp/agent.db"
  146. # 模型配置
  147. model = OpenAILike(
  148. id="qwen3-32b",
  149. api_key=api_key,
  150. base_url=base_url,
  151. request_params={
  152. "extra_body": {"enable_thinking": False},
  153. # "response_format": {"type": "json_object"}
  154. },
  155. )
  156. # 初始化内存和存储
  157. memory = Memory(
  158. model=model,
  159. db=SqliteMemoryDb(table_name="user_memories", db_file=db_file),
  160. )
  161. storage = SqliteStorage(table_name="agent_sessions", db_file=db_file)
  162. # 创建agent实例
  163. agent = Agent(
  164. model=model,
  165. tools=[add_info, add_all_info, get_info, list_info, reset_info],
  166. instructions=SYSTEM_PROMPT,
  167. add_datetime_to_instructions=True,
  168. # show_tool_calls=True,
  169. markdown=True,
  170. memory=memory,
  171. enable_user_memories=True,
  172. storage=storage,
  173. add_history_to_messages=True,
  174. enable_session_summaries=True,
  175. session_state=initial_state,
  176. add_state_in_messages=True,
  177. session_id=session_id,
  178. )
  179. app = FastAPI()
  180. async def generate_agent_response(prompt: str, user_id: str) -> AsyncGenerator[str, None]:
  181. """生成Agent响应的异步生成器"""
  182. try:
  183. response_generator: Iterator[RunResponseEvent] = agent.run(prompt, user_id=user_id, stream=True)
  184. for chunk in response_generator:
  185. # RunResponseEvent通常是一个对象,需根据其属性获取内容
  186. if hasattr(chunk, "content"):
  187. yield json.dumps({"type": "message", "content": chunk.content})
  188. else:
  189. yield json.dumps({"type": "message", "content": str(chunk)})
  190. # await asyncio.sleep(0.05)
  191. # 发送当前状态
  192. state_info = json.loads(list_info(agent))
  193. yield json.dumps({"type": "state", "content": state_info["result"]})
  194. except Exception as e:
  195. yield json.dumps({"type": "error", "content": str(e)})
  196. @app.get("/")
  197. async def root():
  198. return {"message": "信息收集助手 API"}
  199. @app.get("/chat/sse")
  200. async def chat_sse(prompt: str, user_id: str = "default"):
  201. """SSE聊天端点"""
  202. return EventSourceResponse(
  203. generate_agent_response(prompt, user_id),
  204. media_type="text/event-stream"
  205. )
  206. if __name__ == "__main__":
  207. uvicorn.run(app, host="localhost", port=8000)