在最近的人员风险行为分析系统开发过程中,需要解决一个问题:在缺乏GPU资源的情况下,如何提升智能交互能力。我们探索并研究了集成gemma2:2b模型的可行性,这一举措旨在在有限的硬件条件下,为我们的系统注入更高级别的智能性,从而增强用户体验并提高数据分析效率。
具体可分为如下几个步骤:
- 定义用户可能提出的问题类型
- 结合大模型设计接口,以便根据用户的问题提取实体信息
- 实现数据检索机制以提取相关信息返回给用户
如下,是系统页面原型,
交互时,用户输入自己的问题,我们无法预测用户输入什么,但是我们可以定义系统能处理什么类型的问题。在本功能里,我们定义如下系统可以交互的问题分类:
- 人员的基本信息
- 联网设备(应用系统、数据库资产、终端设备资产、服务器资产等)的基本信息
- 人员日常操作风险(从告警列表里读取)
- 和知识图谱的交互线索,例如共用终端风险线索(从知识图谱里查到用户使用两个以上终端设备的图数据)
接着我们要用大模型把用户问题里的实体抽取出来,这就需要我们定义prompt,并利用提示词让大模型帮我们把用户的问题解读成我们可继续执行的实体行为。代码如下:
@app.route('/chat', methods=['POST'])
def chat():
data = request.json
prompt = data.get('prompt', '')
payload = {
"model": "gemma2:2b",
"prompt": BASE_PROMPT + prompt
}
response = requests.post(f"{OLLAMA_API_BASE}/generate", json=payload)
response.headers['Content-Type'] = 'application/json; charset=utf-8'
if response.status_code == 200:
try:
result = response.content.decode('utf-8')
# 尝试解析每一行
responses = [json.loads(line) for line in result.strip().split('\n')]
# 提取所有响应中的 'response' 字段
combined_response = ''.join(r.get('response', '') for r in responses if 'response' in r)
# 使用正则表达式提取括号内的内容
pattern = r'\((.*?)\)'
matches = re.findall(pattern, combined_response)
result_list = []
for match in matches:
# 使用defaultdict来处理重复键
result_dict = defaultdict(list)
# 使用逗号分割键值对,允许逗号前后有空格
pairs = re.split(r'\s*,\s*', match)
# 遍历每个键值对
for pair in pairs:
# 使用冒号分割键和值,允许冒号前后有空格
key_value = re.split(r'\s*:\s*', pair, maxsplit=1)
if len(key_value) == 2:
key, value = key_value
# 去除键和值两端的空格,并将值添加到对应的键的列表中
result_dict[key.strip()].append(value.strip())
# 将defaultdict转换为普通字典,对于只有一个值的键,直接使用该值而不是列表
final_dict = {k: v[0] if len(v) == 1 else v for k, v in result_dict.items()}
result_list.append(final_dict)
return jsonify({"response": result_list})
except json.JSONDecodeError as e:
return jsonify({"error": f"Failed to parse JSON response: {str(e)}"}), 500
else:
return jsonify({"error": "Failed to get response from Ollama"}), 500
代码里的关键部分是BASE_PROMPT,我找了GPT帮我润色,如下,感觉可以写的更好一些。
BASE_PROMPT = """
您的任务是仔细分析用户提供的文本,并从中提取关键实体信息。请特别关注以下四类实体:
1. 用户相关信息:
- 姓名
- 身份证号码
- 所在公司
- 其他相关的个人识别信息
2. 设备相关信息:
- IP地址
- MAC地址
- 操作系统
- 机器名称
- 其他可识别的设备特征
3. 事件相关信息:
- 网络事件
- 风险事件
- 其他值得注意的事件
4. 时间相关信息:
- 时间段
- 开始时间
- 结束时间
- 其他时间
请仔细阅读文本,识别并提取上述实体信息。将提取的信息以(key:value)的格式整理,多个实体之间用逗号分隔。key只能从IP、姓名、身份证、电话号码、操作系统、事件、时间这几个词里选择。例如:
(姓名: 张三, 身份证: 310123199001011234, IP: 192.168.1.100, 操作系统: Windows 10, 事件: 异常登录, 时间: 前天)
如果某类实体在文本中未提及,则无需包含在结果中。请确保提取的信息准确且与原文相符。如果文本中包含其他重要的实体信息,也请一并提取。
请现在开始分析下面用户提供的文本,并按上述格式返回提取的实体信息:
"""
可以看看测试结果如下,
这样我们就可以根据response进行进一步的拆解,按姓名、按事件、按时间来整理数据并最终返回给用户。