message_processor.py
9.27 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
"""
消息处理和文件保存工具类
"""
import os
import base64
import uuid
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] = []
for part in content:
if isinstance(part, dict):
if part.get("type") == "text":
text_parts.append(part.get("text", ""))
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])
if file_paths:
merged_text = (merged_text + "\n" + "\n".join(file_paths)).strip()
new_content = [{"type": "text", "text": merged_text}]
if isinstance(m, dict):
filtered_msg = {**m, "content": new_content}
else:
try:
filtered_msg = m.__class__(
content=new_content,
additional_kwargs=getattr(m, "additional_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": getattr(m, "additional_kwargs", {}),
"response_metadata": getattr(m, "response_metadata", {}),
"content": new_content,
}
filtered_messages.append(filtered_msg)
try:
print(f" -> processed to: {len(merged_text)} chars")
except Exception:
print(" -> processed")
else:
# 不处理，其它消息保持不变
filtered_messages.append(m)
print(" -> processed (no change)")
return filtered_messages, list(self._saved_files)