user_profile_extractor.py 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151
  1. #! /usr/bin/env python
  2. # -*- coding: utf-8 -*-
  3. # vim:fenc=utf-8
  4. import json
  5. from typing import Dict, Any, Optional
  6. from prompt_templates import USER_PROFILE_EXTRACT_PROMPT
  7. from openai import OpenAI
  8. import logging
  9. import global_flags
  10. class UserProfileExtractor:
  11. def __init__(self):
  12. self.llm_client = OpenAI(
  13. api_key='5e275c38-44fd-415f-abcf-4b59f6377f72',
  14. base_url="https://ark.cn-beijing.volces.com/api/v3"
  15. )
  16. self.model_name = 'ep-20250307150409-4blz9'
  17. def get_extraction_function(self) -> Dict:
  18. """
  19. 定义用于用户画像信息提取的Function Calling函数
  20. """
  21. return {
  22. "type": "function",
  23. "function": {
  24. "name": "update_user_profile",
  25. "description": "从用户对话中提取并更新用户的个人信息",
  26. "parameters": {
  27. "type": "object",
  28. "properties": {
  29. "name": {
  30. "type": "string",
  31. "description": "用户的姓名,如果能够准确识别"
  32. },
  33. "preferred_nickname": {
  34. "type": "string",
  35. "description": "用户希望对其的称呼,如果能够准确识别"
  36. },
  37. "age": {
  38. "type": "integer",
  39. "description": "用户的年龄,如果能够准确识别"
  40. },
  41. "region": {
  42. "type": "string",
  43. "description": "用户常驻的地区,不是用户临时所在地"
  44. },
  45. "interests": {
  46. "type": "array",
  47. "items": {"type": "string"},
  48. "description": "用户提到的自己的兴趣爱好"
  49. },
  50. "health_conditions": {
  51. "type": "array",
  52. "items": {"type": "string"},
  53. "description": "用户提及的健康状况"
  54. }
  55. },
  56. "required": []
  57. }
  58. }
  59. }
  60. def generate_extraction_prompt(self, user_profile: Dict, dialogue_history: str) -> str:
  61. """
  62. 生成用于信息提取的系统提示词
  63. """
  64. context = user_profile.copy()
  65. context['dialogue_history'] = dialogue_history
  66. return USER_PROFILE_EXTRACT_PROMPT.format(**context)
  67. def extract_profile_info(self, user_profile, dialogue_history: str) -> Optional[Dict]:
  68. """
  69. 使用Function Calling提取用户画像信息
  70. """
  71. if global_flags.DISABLE_LLM_API_CALL:
  72. return None
  73. try:
  74. logging.debug("try to extract profile from message: {}".format(dialogue_history))
  75. response = self.llm_client.chat.completions.create(
  76. model=self.model_name,
  77. messages=[
  78. {"role": "system", "content": '你是一个专业的用户画像分析助手。'},
  79. {"role": "user", "content": self.generate_extraction_prompt(user_profile, dialogue_history)}
  80. ],
  81. tools=[self.get_extraction_function()],
  82. temperature=0
  83. )
  84. # 解析Function Call的参数
  85. tool_calls = response.choices[0].message.tool_calls
  86. logging.debug(response)
  87. if tool_calls:
  88. function_call = tool_calls[0]
  89. if function_call.function.name == 'update_user_profile':
  90. try:
  91. profile_info = json.loads(function_call.function.arguments)
  92. return {k: v for k, v in profile_info.items() if v}
  93. except json.JSONDecodeError:
  94. logging.error("无法解析提取的用户信息")
  95. return None
  96. except Exception as e:
  97. logging.error(f"用户画像提取出错: {e}")
  98. return None
  99. def merge_profile_info(self, existing_profile: Dict, new_info: Dict) -> Dict:
  100. """
  101. 合并新提取的用户信息到现有资料
  102. """
  103. merged_profile = existing_profile.copy()
  104. merged_profile.update(new_info)
  105. return merged_profile
  106. if __name__ == '__main__':
  107. extractor = UserProfileExtractor()
  108. current_profile = {
  109. 'name': '',
  110. 'preferred_nickname': '李叔',
  111. 'age': 0,
  112. 'region': '北京',
  113. 'health_conditions': [],
  114. 'medications': [],
  115. 'interests': []
  116. }
  117. message = "我回天津老家了"
  118. resp = extractor.extract_profile_info(current_profile, message)
  119. print(resp)
  120. message = "好的,孩子,我是老李头,今年68啦,住在北京海淀区。平时喜欢在微信上跟老伙伴们聊聊养生、下下象棋,偶尔也跟年轻人学学新鲜事儿。\n" \
  121. "你叫我李叔就行,有啥事儿咱们慢慢聊啊\n" \
  122. "哎,今儿个天气不错啊,我刚才还去楼下小公园溜达了一圈儿。碰到几个老伙计在打太极,我也跟着比划了两下,这老胳膊老腿的,原来老不舒服,活动活动舒坦多了!\n" \
  123. "你吃饭了没?我们这儿中午吃的打卤面,老伴儿做的,香得很!这人老了就爱念叨些家长里短的,你可别嫌我啰嗦啊。\n" \
  124. "对了,最近我孙子教我发语音,比打字方便多啦!就是有时候一激动,说话声音太大,把手机都给震得嗡嗡响\n"
  125. resp = extractor.extract_profile_info(current_profile, message)
  126. print(resp)
  127. print(extractor.merge_profile_info(current_profile, resp))
  128. current_profile = {
  129. 'name': '李老头',
  130. 'preferred_nickname': '李叔',
  131. 'age': 68,
  132. 'region': '北京市海淀区',
  133. 'health_conditions': [],
  134. 'medications': [],
  135. 'interests': ['养生', '下象棋']
  136. }
  137. resp = extractor.extract_profile_info(current_profile, message)
  138. print(resp)
  139. print(extractor.merge_profile_info(current_profile, resp))