| 1 | from mimetypes import guess_type |
| 2 | |
| 3 | from langchain_core.messages import HumanMessage |
| 4 | |
| 5 | from helpers import ( |
| 6 | chat_media, |
| 7 | ephemeral_images, |
| 8 | files, |
| 9 | history, |
| 10 | images, |
| 11 | parallel_tools, |
| 12 | runtime, |
| 13 | ) |
| 14 | from helpers.tool import Response, Tool |
| 15 | from plugins._model_config.helpers.model_config import ( |
| 16 | build_vision_model, |
| 17 | get_chat_model_config, |
| 18 | get_vision_model_config, |
| 19 | ) |
| 20 | |
| 21 | # image token estimation for context window |
| 22 | TOKENS_ESTIMATE = 1500 |
| 23 | |
| 24 | |
| 25 | class VisionLoad(Tool): |
| 26 | async def execute( |
| 27 | self, |
| 28 | paths: list[str] | str = [], |
| 29 | query: str = "", |
| 30 | **kwargs, |
| 31 | ) -> Response: |
| 32 | |
| 33 | self.images_dict = {} |
| 34 | self.loaded_paths: list[str] = [] |
| 35 | self.skipped_paths: list[str] = [] |
| 36 | self.vision_config = get_vision_model_config(self.agent) |
| 37 | if isinstance(paths, str): |
| 38 | paths = [paths] |
| 39 | if not isinstance(paths, list): |
| 40 | return Response( |
| 41 | message="vision_load error: `paths` must be a string or an array.", |
| 42 | break_loop=False, |
| 43 | ) |
| 44 | |
| 45 | max_embeds = self._get_max_embeds() |
| 46 | requested = [ |
| 47 | (str(path or "").strip(), self._display_input_path(str(path or "").strip(), idx + 1)) |
| 48 | for idx, path in enumerate(paths) |
| 49 | ] |
| 50 | limited_paths = requested if max_embeds <= 0 else requested[-max_embeds:] |
| 51 | self.skipped_paths = ( |
| 52 | [display for _, display in requested[:-max_embeds]] |
| 53 | if max_embeds > 0 and len(requested) > max_embeds |
| 54 | else [] |
| 55 | ) |
| 56 | |
| 57 | for idx, (path, display_path) in enumerate(limited_paths): |
| 58 | if not path: |
| 59 | continue |
| 60 | if ephemeral_images.is_ref(path): |
| 61 | image = ephemeral_images.consume_image( |
| 62 | path, |
| 63 | context_id=self._context_id(), |
| 64 | ) |
| 65 | if image is None: |
| 66 | continue |
| 67 | display = image.display_name or display_path |
| 68 | stored_ref = self._store_ephemeral_image(image) |
| 69 | if stored_ref: |
| 70 | self.images_dict[display] = stored_ref |
| 71 | self.loaded_paths.append(display) |
| 72 | continue |
| 73 | if self._is_data_image_url(path): |
| 74 | stored_ref = self._store_data_url(path, preferred_name=f"vision-load-{idx + 1}.png") |
| 75 | if stored_ref: |
| 76 | self.images_dict[display_path] = stored_ref |
| 77 | self.loaded_paths.append(display_path) |
| 78 | continue |
| 79 | if not await runtime.call_development_function(files.exists, str(path)): |
| 80 | continue |
| 81 | |
| 82 | if path not in self.images_dict: |
| 83 | mime_type, _ = guess_type(str(path)) |
| 84 | if mime_type and mime_type.startswith("image/"): |
| 85 | try: |
| 86 | stored_ref = self._store_local_image(path, preferred_name=files.basename(path)) |
| 87 | self.images_dict[display_path] = stored_ref |
| 88 | self.loaded_paths.append(display_path) |
| 89 | except (FileNotFoundError, OSError, ValueError): |
| 90 | continue |
| 91 | |
| 92 | message = self._summary() if self.images_dict or self.skipped_paths else "No images processed" |
| 93 | if self.vision_config and self.images_dict: |
| 94 | try: |
| 95 | capsule = await self._call_vision_model( |
| 96 | list(self.images_dict.values()), |
| 97 | query, |
| 98 | ) |
| 99 | message = ( |
| 100 | f"Analyzed {len(self.images_dict)} image(s)" |
| 101 | f"; {len(self.skipped_paths)} skipped.\n\n{capsule.strip()}" |
| 102 | ) |
| 103 | except Exception as exc: |
| 104 | message = f"Image analysis error: {str(exc)[:1000]}" |
| 105 | return Response(message=message, break_loop=False) |
| 106 | |
| 107 | def _get_max_embeds(self) -> int: |
| 108 | cfg = self.vision_config or get_chat_model_config(self.agent) |
| 109 | return int(cfg.get("max_embeds", 10) or 0) |
| 110 | |
| 111 | def _context_id(self) -> str: |
| 112 | context = getattr(self.agent, "context", None) |
| 113 | get_data = getattr(context, "get_data", None) |
| 114 | parent_id = ( |
| 115 | get_data(parallel_tools.PARALLEL_WORKER_PARENT_CONTEXT_KEY) |
| 116 | if get_data |
| 117 | else "" |
| 118 | ) |
| 119 | return str(parent_id or getattr(context, "id", "") or "").strip() |
| 120 | |
| 121 | async def _call_vision_model(self, image_paths: list[str], query: str) -> str: |
| 122 | user_message = getattr(self.agent, "last_user_message", None) |
| 123 | output_text = getattr(user_message, "output_text", None) |
| 124 | request = str(output_text() if callable(output_text) else "").strip() |
| 125 | content = [ |
| 126 | { |
| 127 | "type": "text", |
| 128 | "text": self.agent.read_prompt( |
| 129 | "fw.vision_load.md", |
| 130 | request=request, |
| 131 | query=str(query or "").strip(), |
| 132 | ), |
| 133 | } |
| 134 | ] |
| 135 | content.extend( |
| 136 | {"type": "image_url", "image_url": {"url": path}} |
| 137 | for path in image_paths |
| 138 | ) |
| 139 | response, _ = await build_vision_model(self.agent).unified_call( |
| 140 | messages=[HumanMessage(content=content)], |
| 141 | ) |
| 142 | if not str(response or "").strip(): |
| 143 | raise RuntimeError("Vision Model returned an empty response.") |
| 144 | return str(response) |
| 145 | |
| 146 | def _store_ephemeral_image(self, image: ephemeral_images.EphemeralImage) -> str: |
| 147 | context_id = self._context_id() |
| 148 | if not context_id: |
| 149 | return image.data_url |
| 150 | source = chat_media.infer_source(image.ref, image.display_name) |
| 151 | category = chat_media.category_for_source(source) |
| 152 | saved = chat_media.save_image_base64( |
| 153 | context_id=context_id, |
| 154 | data=image.data, |
| 155 | mime_type=image.mime, |
| 156 | category=category, |
| 157 | source=source, |
| 158 | preferred_name=image.display_name, |
| 159 | ) |
| 160 | return saved.a0_path |
| 161 | |
| 162 | def _store_data_url(self, data_url: str, *, preferred_name: str = "") -> str: |
| 163 | context_id = self._context_id() |
| 164 | if not context_id: |
| 165 | return data_url |
| 166 | source = chat_media.infer_source(data_url, preferred_name) |
| 167 | category = chat_media.category_for_source(source) |
| 168 | saved = chat_media.save_image_data_url( |
| 169 | context_id=context_id, |
| 170 | data_url=data_url, |
| 171 | category=category, |
| 172 | source=source, |
| 173 | preferred_name=preferred_name, |
| 174 | ) |
| 175 | return saved.a0_path |
| 176 | |
| 177 | def _store_local_image(self, path: str, *, preferred_name: str = "") -> str: |
| 178 | context_id = self._context_id() |
| 179 | if not context_id: |
| 180 | return images.to_data_url(path) |
| 181 | return chat_media.materialize_image_ref( |
| 182 | context_id=context_id, |
| 183 | url=path, |
| 184 | source=chat_media.infer_source(path, preferred_name), |
| 185 | preferred_name=preferred_name, |
| 186 | ) |
| 187 | |
| 188 | def _summary(self) -> str: |
| 189 | loaded = "\n".join(self.loaded_paths) if self.loaded_paths else "none" |
| 190 | summary = f"Loaded images ({len(self.loaded_paths)}):\n{loaded}" |
| 191 | if self.skipped_paths: |
| 192 | summary += ( |
| 193 | f"\n\nSkipped images ({len(self.skipped_paths)}, max {self._get_max_embeds()}):\n" |
| 194 | + "\n".join(self.skipped_paths) |
| 195 | ) |
| 196 | return summary |
| 197 | |
| 198 | @staticmethod |
| 199 | def _is_data_image_url(value: str) -> bool: |
| 200 | normalized = str(value or "").strip().lower() |
| 201 | return normalized.startswith("data:image/") and ";base64," in normalized |
| 202 | |
| 203 | @classmethod |
| 204 | def _display_input_path(cls, value: str, index: int) -> str: |
| 205 | if ephemeral_images.is_ref(value): |
| 206 | return ephemeral_images.display_ref(value) |
| 207 | if cls._is_data_image_url(value): |
| 208 | prefix = value.split(",", 1)[0] |
| 209 | return f"{prefix},<ephemeral-image-{index}>" |
| 210 | return value |
| 211 | |
| 212 | async def after_execution(self, response: Response, **kwargs): |
| 213 | await super().after_execution(response, **kwargs) |
| 214 | if self.images_dict and not self.vision_config: |
| 215 | content = [ |
| 216 | {"type": "image_url", "image_url": {"url": image_path}} |
| 217 | for image_path in self.images_dict.values() |
| 218 | ] |
| 219 | raw_message = history.RawMessage( |
| 220 | raw_content=content, |
| 221 | preview="<Image attachments loaded by path>", |
| 222 | ) |
| 223 | tokens = TOKENS_ESTIMATE * len(content) |
| 224 | if not parallel_tools.queue_parallel_parent_history( |
| 225 | self.agent, |
| 226 | content=raw_message, |
| 227 | tokens=tokens, |
| 228 | ): |
| 229 | self.agent.hist_add_message( |
| 230 | False, |
| 231 | content=raw_message, |
| 232 | tokens=tokens, |
| 233 | ) |