Coverage for agentos/multimodal/manager.py: 0%

1""" 

2AgentOS v0.40 Multimodal — 多模态输入支持。 

3支持：图片理解、语音转文字、PDF/文档解析。 

4""" 

5 

6from __future__ import annotations 

7 

8import base64 

9import logging 

10from dataclasses import dataclass, field 

11from enum import StrEnum 

12 

13logger = logging.getLogger(__name__) 

14 

15 

16class Modality(StrEnum): 

17 """模态类型枚举。""" 

18 

19 TEXT = "text" 

20 IMAGE = "image" 

21 AUDIO = "audio" 

22 VIDEO = "video" 

23 DOCUMENT = "document" 

24 

25 

26@dataclass 

27class MultimodalBlock: 

28 """多模态输入块 — 遵循OpenAI/Anthropic content block格式。""" 

29 

30 type: str # text | image_url | audio | image 

31 text: str = "" 

32 source: dict = field(default_factory=dict) 

33 mime_type: str = "" 

34 

35 @classmethod 

36 def text_block(cls, text: str) -> MultimodalBlock: 

37 return cls(type="text", text=text) 

38 

39 @classmethod 

40 def image_url(cls, url: str, detail: str = "auto") -> MultimodalBlock: 

41 return cls( 

42 type="image_url", 

43 source={"type": "image_url", "image_url": {"url": url, "detail": detail}}, 

44 ) 

45 

46 @classmethod 

47 def image_base64(cls, data: bytes, mime: str = "image/jpeg") -> MultimodalBlock: 

48 b64 = base64.b64encode(data).decode() 

49 return cls( 

50 type="image_url", 

51 source={"type": "image_url", "image_url": {"url": f"data:{mime};base64,{b64}"}}, 

52 ) 

53 

54 @classmethod 

55 def audio(cls, data: bytes, mime: str = "audio/wav") -> MultimodalBlock: 

56 b64 = base64.b64encode(data).decode() 

57 return cls(type="audio", mime_type=mime, source={"data": b64}) 

58 

59 def to_openai_format(self) -> dict: 

60 if self.type == "text": 

61 return {"type": "text", "text": self.text} 

62 if self.type == "image_url": 

63 return {"type": "image_url", "image_url": self.source["image_url"]} 

64 return {"type": self.type, **self.source} 

65 

66 

67class ImageProcessor: 

68 """图片处理器 — 压缩、格式转换、OCR预处理。""" 

69 

70 MAX_SIZE = 2048 

71 JPEG_QUALITY = 85 

72 

73 @staticmethod 

74 def encode_file(path: str) -> tuple[str, str]: 

75 """返回(base64, mime_type)。""" 

76 import mimetypes 

77 

78 mime = mimetypes.guess_type(path)[0] or "image/png" 

79 with open(path, "rb") as f: 

80 data = f.read() 

81 return base64.b64encode(data).decode(), mime 

82 

83 @staticmethod 

84 def encode_bytes(data: bytes, mime: str = "image/jpeg") -> str: 

85 return base64.b64encode(data).decode() 

86 

87 @staticmethod 

88 def estimate_tokens(width: int, height: int, detail: str = "auto") -> int: 

89 """估算图片token消耗（OpenAI定价模型）。""" 

90 if detail == "low": 

91 return 85 

92 # high detail 

93 short_side = min(width, height) 

94 scale = min(768 / short_side, 1.0) if short_side > 768 else 1.0 

95 w = int(width * scale) 

96 h = int(height * scale) 

97 tiles = ((w + 511) // 512) * ((h + 511) // 512) 

98 return 85 + 170 * tiles 

99 

100 @staticmethod 

101 def purge_metadata(data: bytes) -> bytes: 

102 """清除图片EXIF元数据。""" 

103 try: 

104 import io 

105 

106 from PIL import Image 

107 

108 img = Image.open(io.BytesIO(data)) 

109 data_no_exif = list(img.getdata()) 

110 cleaned = Image.new(img.mode, img.size) 

111 cleaned.putdata(data_no_exif) 

112 buf = io.BytesIO() 

113 cleaned.save(buf, format=img.format or "PNG") 

114 return buf.getvalue() 

115 except ImportError: 

116 return data 

117 

118 

119class AudioProcessor: 

120 """音频处理器 — 转文字、格式转换。""" 

121 

122 SUPPORTED_FORMATS = ["wav", "mp3", "ogg", "flac", "m4a"] 

123 

124 @staticmethod 

125 def transcribe(path: str, whisper_model: str = "base") -> str: 

126 """使用whisper转文字。""" 

127 try: 

128 import whisper 

129 

130 model = whisper.load_model(whisper_model) 

131 result = model.transcribe(path) 

132 return result["text"] 

133 except ImportError: 

134 logger.warning("whisper not installed, returning empty") 

135 return "[whisper not available]" 

136 

137 @staticmethod 

138 def encode_file(path: str) -> tuple[str, str]: 

139 import mimetypes 

140 

141 mime = mimetypes.guess_type(path)[0] or "audio/wav" 

142 with open(path, "rb") as f: 

143 data = f.read() 

144 return base64.b64encode(data).decode(), mime 

145 

146 

147class DocumentParser: 

148 """文档解析器 — PDF/Word/Markdown。""" 

149 

150 @staticmethod 

151 def parse_pdf(path: str) -> str: 

152 try: 

153 import PyPDF2 

154 

155 text = [] 

156 with open(path, "rb") as f: 

157 reader = PyPDF2.PdfReader(f) 

158 for page in reader.pages: 

159 page_text = page.extract_text() 

160 if page_text: 

161 text.append(page_text) 

162 return "\n\n".join(text) 

163 except ImportError: 

164 logger.warning("PyPDF2 not installed") 

165 return "[PyPDF2 not available]" 

166 

167 @staticmethod 

168 def parse_docx(path: str) -> str: 

169 try: 

170 from docx import Document 

171 

172 doc = Document(path) 

173 return "\n".join(p.text for p in doc.paragraphs if p.text) 

174 except ImportError: 

175 logger.warning("python-docx not installed") 

176 return "[python-docx not available]" 

177 

178 @staticmethod 

179 def parse_auto(path: str) -> tuple[str, str]: 

180 """自动检测文件类型并解析。返回 (content, format)。""" 

181 ext = path.rsplit(".", 1)[-1].lower() if "." in path else "" 

182 if ext == "pdf": 

183 return DocumentParser.parse_pdf(path), "pdf" 

184 elif ext in ("docx", "doc"): 

185 return DocumentParser.parse_docx(path), "docx" 

186 elif ext in ("md", "markdown", "txt"): 

187 with open(path) as f: 

188 return f.read(), ext 

189 else: 

190 try: 

191 with open(path) as f: 

192 return f.read(), "text" 

193 except Exception: 

194 return "", "unknown" 

195 

196 

197class MultimodalManager: 

198 """多模态管理器 — 统一入口。""" 

199 

200 def __init__(self): 

201 self.image = ImageProcessor() 

202 self.audio = AudioProcessor() 

203 self.document = DocumentParser() 

204 

205 def prepare_input(self, blocks: list[MultimodalBlock]) -> list[dict]: 

206 """转换为OpenAI兼容格式。""" 

207 return [b.to_openai_format() for b in blocks] 

208 

209 def from_files(self, paths: list[str]) -> list[MultimodalBlock]: 

210 """从文件路径自动推断模态。""" 

211 blocks = [] 

212 image_exts = {"png", "jpg", "jpeg", "gif", "webp", "bmp"} 

213 audio_exts = {"wav", "mp3", "ogg", "flac", "m4a"} 

214 doc_exts = {"pdf", "docx", "doc", "md", "txt"} 

215 

216 for p in paths: 

217 ext = p.rsplit(".", 1)[-1].lower() if "." in p else "" 

218 try: 

219 if ext in image_exts: 

220 b64, mime = ImageProcessor.encode_file(p) 

221 blocks.append( 

222 MultimodalBlock( 

223 type="image_url", 

224 source={ 

225 "type": "image_url", 

226 "image_url": {"url": f"data:{mime};base64,{b64}"}, 

227 }, 

228 ) 

229 ) 

230 elif ext in audio_exts: 

231 b64, mime = AudioProcessor.encode_file(p) 

232 blocks.append( 

233 MultimodalBlock(type="audio", mime_type=mime, source={"data": b64}) 

234 ) 

235 elif ext in doc_exts: 

236 text, fmt = DocumentParser.parse_auto(p) 

237 blocks.append(MultimodalBlock.text_block(text)) 

238 else: 

239 with open(p) as f: 

240 blocks.append(MultimodalBlock.text_block(f.read())) 

241 except Exception as e: 

242 blocks.append(MultimodalBlock.text_block(f"[Error reading {p}: {e}]")) 

243 return blocks 

244 

245 def stats(self) -> dict: 

246 return {"modalities": ["text", "image", "audio", "video", "document"]}