message_processor.py
11.5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
"""
消息处理和文件保存工具类
"""
import os
import base64
import uuid
import re
from pathlib import Path
from typing import List, Any, Union
class MessageProcessor:
"""消息处理和文件保存工具类"""
def __init__(self, save_dir: str = None):
"""
初始化消息处理器
Args:
save_dir: 文件保存目录，默认为环境变量 ATTACH_SAVE_DIR 或 "uploads"
"""
self.save_dir = save_dir or os.getenv("ATTACH_SAVE_DIR", "uploads")
self._saved_files: list[str] = []
def save_and_get_file_url(self, file_data: str, filename: str = None, mime_type: str = None) -> str:
"""
保存文件并返回绝对路径字符串
Args:
file_data: base64编码的文件数据
filename: 文件名
mime_type: MIME类型
Returns:
文件绝对路径字符串
"""
uploads = Path(self.save_dir)
uploads.mkdir(exist_ok=True)
# 生成安全文件名
safe_name = filename or f"{uuid.uuid4().hex}"
if mime_type == "application/pdf" and not safe_name.lower().endswith(".pdf"):
safe_name += ".pdf"
out_path = uploads / safe_name
try:
# 解码 base64 数据并保存
with open(out_path, "wb") as f:
f.write(base64.b64decode(file_data))
abs_path = str(out_path.resolve())
print(f" -> saved file: {abs_path}")
self._saved_files.append(abs_path)
return abs_path
except Exception as e:
print(f" -> save file failed: {e}")
return f"[附件保存失败: {filename or 'unknown'}]"
def process_content(self, content: Any) -> str:
"""
处理消息内容，提取文件并转换为文本
Args:
content: 消息内容，可能是字符串、列表或字典
Returns:
处理后的文本内容
"""
if isinstance(content, str):
return content
elif isinstance(content, list):
text_parts = []
for part in content:
if isinstance(part, str):
text_parts.append(part)
elif isinstance(part, dict):
part_type = part.get("type")
if part_type == "text":
text_parts.append(part.get("text", ""))
elif part_type == "file":
# 处理文件类型
file_data = part.get("data")
filename = part.get("metadata", {}).get("filename")
mime_type = part.get("mime_type")
if file_data:
file_path = self.save_and_get_file_url(file_data, filename, mime_type)
text_parts.append(file_path)
else:
text_parts.append(f"[文件缺失: {filename or 'unknown'}]")
else:
# 未知类型最小化占位
text_parts.append(f"[{part_type or 'unknown'}]")
return "\n".join([t for t in text_parts if t])
elif isinstance(content, dict):
# 单个字典内容
ctype = content.get("type")
if ctype == "text":
return content.get("text", "")
if ctype == "file":
file_data = content.get("data")
filename = content.get("metadata", {}).get("filename")
mime_type = content.get("mime_type")
if file_data:
return self.save_and_get_file_url(file_data, filename, mime_type)
else:
return f"[文件缺失: {filename or 'unknown'}]"
else:
return f"[{ctype or 'unknown'}]"
else:
return str(content)
def process_messages(self, messages: List[Any]) -> tuple[List[Any], List[str]]:
"""
处理消息列表，提取文件并转换为文本
Args:
messages: 消息列表
Returns:
(处理后的消息列表, 保存的文件路径列表)
"""
filtered_messages = []
for idx, m in enumerate(messages or []):
# 兼容 dict 或 LangChain 的消息对象
if isinstance(m, dict):
role = m.get("role")
content = m.get("content")
else:
role = getattr(m, "role", None)
content = getattr(m, "content", None)
print(f"[msg#{idx}] role={role!r}, content_type={type(content).__name__}")
# 仅当 HumanMessage 且 content 中包含 file 分片时进行处理；否则保持原样
msg_type = m.get("type") if isinstance(m, dict) else m.__class__.__name__
should_process = (msg_type == "HumanMessage") and isinstance(content, list) and any(isinstance(p, dict) and p.get("type") == "file" for p in content)
if should_process:
# 提取文本与文件路径
text_parts: list[str] = []
file_paths: list[str] = []
# 保留的参数：slip_id, token, ver
preserved_params = {}
for part in content:
if isinstance(part, dict):
if part.get("type") == "text":
text_parts.append(part.get("text", ""))
# 从文本部分提取 slip_id, token, ver 参数
if "slip_id" in part:
preserved_params["slip_id"] = part.get("slip_id")
if "token" in part:
preserved_params["token"] = part.get("token")
if "ver" in part:
preserved_params["ver"] = part.get("ver")
elif part.get("type") == "file":
# 兼容两种文件结构：
# 1) {"type":"file", "data":"<base64>", "metadata":{"filename":...}, "mime_type":"application/pdf"}
# 2) {"type":"file", "file":{"file_data":"data:application/pdf;base64,<base64>", "filename":"..."}}
file_data = part.get("data")
filename = part.get("metadata", {}).get("filename")
mime_type = part.get("mime_type")
if not file_data and isinstance(part.get("file"), dict):
f = part.get("file") or {}
file_data_url = f.get("file_data")
filename = f.get("filename") or filename
# 解析 data URL 或原始 base64
if isinstance(file_data_url, str):
if file_data_url.startswith("data:") and "," in file_data_url:
try:
header, b64_payload = file_data_url.split(",", 1)
# 格式如 data:application/pdf;base64
if header.startswith("data:") and ";" in header:
mime_type = header[5:].split(";", 1)[0] or mime_type
file_data = b64_payload
except Exception:
file_data = None
else:
# 非 data URL，当作纯 base64
file_data = file_data_url
if file_data:
saved_path = self.save_and_get_file_url(file_data, filename, mime_type)
file_paths.append(saved_path)
merged_text = "\n".join([t for t in text_parts if t])
# 如果存在 slip_id，将其拼接到文本中
if "slip_id" in preserved_params:
slip_id_value = preserved_params["slip_id"]
# 检查文本中是否已有 slip_id：xxx 的格式，如果有则替换，否则追加
slip_id_pattern = r'slip_id[：:]\s*\d+'
if re.search(slip_id_pattern, merged_text):
# 替换原有的 slip_id：xxx
merged_text = re.sub(slip_id_pattern, f'slip_id：{slip_id_value}', merged_text)
else:
# 追加 slip_id：xxx
merged_text = f"{merged_text} slip_id：{slip_id_value}"
if file_paths:
merged_text = (merged_text + "\n" + "\n".join(file_paths)).strip()
# 将保留的参数添加到content字典中，而不是additional_kwargs
new_content_dict = {"type": "text", "text": merged_text}
if preserved_params:
new_content_dict.update(preserved_params)
new_content = [new_content_dict]
# 获取原有的 additional_kwargs（不合并保留的参数）
if isinstance(m, dict):
existing_kwargs = m.get("additional_kwargs", {})
else:
existing_kwargs = getattr(m, "additional_kwargs", {})
# 保持原有的 additional_kwargs，不添加保留的参数
merged_kwargs = existing_kwargs
if isinstance(m, dict):
filtered_msg = {**m, "content": new_content, "additional_kwargs": merged_kwargs}
else:
try:
filtered_msg = m.__class__(
content=new_content,
additional_kwargs=merged_kwargs,
response_metadata=getattr(m, "response_metadata", {}),
id=getattr(m, "id", None),
)
except Exception:
# 兜底为等价字典并保留 id
filtered_msg = {
"type": m.__class__.__name__,
"role": (role or "user"),
"id": getattr(m, "id", None),
"additional_kwargs": merged_kwargs,
"response_metadata": getattr(m, "response_metadata", {}),
"content": new_content,
}
filtered_messages.append(filtered_msg)
try:
print(f" -> processed to: {len(merged_text)} chars")
if preserved_params:
print(f" -> preserved params: {preserved_params}")
except Exception:
print(" -> processed")
if preserved_params:
print(f" -> preserved params: {preserved_params}")
else:
# 不处理，其它消息保持不变
filtered_messages.append(m)
print(" -> processed (no change)")
return filtered_messages, list(self._saved_files)