general_agent_tools.py
1 """ 2 Collection of Python tools that can be configured for general purpose agents. 3 """ 4 5 import logging 6 import asyncio 7 import os 8 import tempfile 9 import uuid 10 from datetime import datetime, timezone 11 from typing import Any, Dict, Optional, Tuple 12 from playwright.async_api import async_playwright 13 14 from google.adk.tools import ToolContext 15 16 from markitdown import MarkItDown, UnsupportedFormatException 17 from mermaid_cli import render_mermaid 18 19 from ...agent.utils.artifact_helpers import ensure_correct_extension 20 from google.genai import types as adk_types 21 from .tool_definition import BuiltinTool 22 from .tool_result import ToolResult, DataObject, DataDisposition 23 from .artifact_types import Artifact 24 from .registry import tool_registry 25 26 log = logging.getLogger(__name__) 27 28 def _simple_truncate_text(text: str, max_bytes: int = 2048) -> Tuple[str, bool]: 29 """Truncates text to a maximum number of bytes for preview.""" 30 truncated = False 31 preview_text = text 32 if not isinstance(text, str): 33 return "", False 34 if len(text.encode("utf-8")) > max_bytes: 35 encoded = text.encode("utf-8") 36 preview_text = encoded[:max_bytes].decode("utf-8", errors="ignore") + "..." 37 truncated = True 38 return preview_text, truncated 39 40 41 async def convert_file_to_markdown( 42 input_filename: Artifact, 43 tool_context: ToolContext = None, 44 tool_config: Optional[Dict[str, Any]] = None, 45 ) -> ToolResult: 46 """ 47 Converts an input file artifact to Markdown using the MarkItDown library. 48 The supported input types are those supported by MarkItDown (e.g., PDF, DOCX, XLSX, HTML, CSV, PPTX, ZIP). 49 The output is a new Markdown artifact. 50 51 Args: 52 input_filename: The input file artifact (pre-loaded by the framework). 53 tool_context: The context provided by the ADK framework. 54 tool_config: Optional dictionary for tool-specific configuration (unused by this tool). 55 56 Returns: 57 ToolResult with output artifact details and a preview of the result. 58 """ 59 if not tool_context: 60 return ToolResult.error("ToolContext is missing.") 61 62 log_identifier = f"[GeneralTool:convert_to_markdown:{input_filename.filename}]" 63 log.info("%s Processing request.", log_identifier) 64 65 temp_input_file = None 66 67 try: 68 # Use pre-loaded artifact data 69 artifact_filename = input_filename.filename 70 artifact_version = input_filename.version 71 input_bytes = input_filename.as_bytes() 72 source_metadata = input_filename.metadata or {} 73 74 # Get original filename from metadata or use artifact filename 75 original_input_filename = source_metadata.get("filename", artifact_filename) 76 original_input_basename, original_input_ext = os.path.splitext(original_input_filename) 77 78 # Write to temp file for MarkItDown processing 79 temp_suffix = original_input_ext if original_input_ext else None 80 temp_input_file = tempfile.NamedTemporaryFile(delete=False, suffix=temp_suffix) 81 temp_input_file.write(input_bytes) 82 temp_input_file.close() 83 log.debug( 84 "%s Input artifact content written to temporary file: %s", 85 log_identifier, 86 temp_input_file.name, 87 ) 88 89 md_converter = MarkItDown() 90 log.debug( 91 "%s Calling MarkItDown.convert() on %s", 92 log_identifier, 93 temp_input_file.name, 94 ) 95 conversion_result = await asyncio.to_thread( 96 md_converter.convert, temp_input_file.name 97 ) 98 99 markdown_text_content = ( 100 conversion_result.text_content 101 if conversion_result and conversion_result.text_content 102 else "" 103 ) 104 if not markdown_text_content: 105 log.warning( 106 "%s MarkItDown conversion resulted in empty content for %s.", 107 log_identifier, 108 artifact_filename, 109 ) 110 111 output_filename = f"{original_input_basename}_converted.md" 112 113 metadata = { 114 "description": f"Markdown conversion of '{original_input_basename}{original_input_ext}'", 115 "source_artifact": artifact_filename, 116 "source_artifact_version": artifact_version, 117 "conversion_tool": "MarkItDown", 118 "conversion_timestamp": datetime.now(timezone.utc).isoformat(), 119 } 120 121 preview_data, truncated = _simple_truncate_text(markdown_text_content) 122 preview_message = f"File converted to Markdown successfully." 123 if truncated: 124 preview_message += " Preview shows first portion." 125 126 log.info("%s Returning markdown as DataObject for artifact storage", log_identifier) 127 128 return ToolResult.ok( 129 preview_message, 130 data={ 131 "result_preview": preview_data, 132 "result_truncated": truncated, 133 }, 134 data_objects=[ 135 DataObject( 136 name=output_filename, 137 content=markdown_text_content, 138 mime_type="text/markdown", 139 disposition=DataDisposition.ARTIFACT_WITH_PREVIEW, 140 description=f"Markdown conversion of '{original_input_basename}{original_input_ext}'", 141 metadata=metadata, 142 preview=preview_data, 143 ) 144 ], 145 ) 146 147 except UnsupportedFormatException as e: 148 log.warning("%s MarkItDown unsupported format: %s", log_identifier, e) 149 return ToolResult.error(f"Unsupported file format for MarkItDown: {e}") 150 except ValueError as e: 151 log.warning("%s Value error: %s", log_identifier, e) 152 return ToolResult.error(str(e)) 153 except Exception as e: 154 log.exception( 155 "%s Unexpected error in convert_file_to_markdown: %s", log_identifier, e 156 ) 157 return ToolResult.error(f"An unexpected error occurred: {e}") 158 finally: 159 if ( 160 temp_input_file 161 and temp_input_file.name 162 and os.path.exists(temp_input_file.name) 163 ): 164 try: 165 os.remove(temp_input_file.name) 166 log.debug( 167 "%s Removed temporary input file: %s", 168 log_identifier, 169 temp_input_file.name, 170 ) 171 except OSError as e_remove: 172 log.error( 173 "%s Failed to remove temporary input file %s: %s", 174 log_identifier, 175 temp_input_file.name, 176 e_remove, 177 ) 178 179 async def _convert_svg_to_png_with_playwright(svg_data: str, scale: int = 2) -> bytes: 180 """ 181 Converts SVG data to a PNG image using Playwright. 182 183 Args: 184 svg_data (str): The SVG data to be converted. 185 scale (int, optional): The scale factor for the PNG image. Defaults to 2. 186 187 Returns: 188 bytes: The PNG image data as a byte array. 189 190 Raises: 191 ValueError: If the SVG bounding box cannot be determined. 192 193 """ 194 async with async_playwright() as p: 195 browser = await p.chromium.launch() 196 context = await browser.new_context(device_scale_factor=scale) 197 page = await context.new_page() 198 199 html_content = f""" 200 <html> 201 <body style="margin: 0; padding: 0; background: white;"> 202 <div id="container">{svg_data}</div> 203 </body> 204 </html> 205 """ 206 await page.set_content(html_content, wait_until="load") 207 await page.wait_for_timeout(50) 208 209 svg = page.locator("svg") 210 box = await svg.bounding_box() 211 if not box: 212 raise ValueError("Could not determine SVG bounding box.") 213 214 width = int(box["width"]) 215 height = int(box["height"]) 216 217 await page.set_viewport_size({"width": width, "height": height}) 218 219 image_data = await svg.screenshot(type="png") 220 221 await browser.close() 222 return image_data 223 224 async def mermaid_diagram_generator( 225 mermaid_syntax: str, 226 output_filename: Optional[str] = None, 227 tool_context: ToolContext = None, 228 tool_config: Optional[Dict[str, Any]] = None, 229 ) -> ToolResult: 230 """ 231 Generates a PNG image from Mermaid diagram syntax and saves it as an artifact. 232 The diagram must be detailed. 233 234 Args: 235 mermaid_syntax: The Mermaid diagram syntax (string). 236 output_filename: Optional desired name for the output PNG file. 237 If not provided, a unique name will be generated. 238 tool_context: The context provided by the ADK framework. 239 tool_config: Optional dictionary for tool-specific configuration (unused by this tool). 240 241 Returns: 242 ToolResult with output artifact details and a preview message. 243 """ 244 if not tool_context: 245 return ToolResult.error("ToolContext is missing.") 246 247 log_identifier = "[GeneralTool:mermaid_diagram_generator]" 248 if output_filename: 249 log_identifier += f":{output_filename}" 250 log.info("%s Processing request.", log_identifier) 251 252 try: 253 log.debug( 254 "%s Calling render_mermaid for syntax: %s", 255 log_identifier, 256 mermaid_syntax[:100] + "...", 257 ) 258 title, desc, svg_image_data = await render_mermaid( 259 mermaid_syntax, output_format="svg", background_color="white" 260 ) 261 262 if not svg_image_data: 263 log.error( 264 "%s Failed to render Mermaid diagram. No image data returned.", 265 log_identifier, 266 ) 267 return ToolResult.error( 268 "Failed to render Mermaid diagram. No image data returned." 269 ) 270 try: 271 scale = max(2, len(mermaid_syntax.splitlines()) // 10) 272 image_data = await _convert_svg_to_png_with_playwright(svg_image_data.decode("utf-8"), scale) 273 except Exception as e: 274 log.error( 275 "%s Failed to convert SVG to PNG with Playwright: %s", 276 log_identifier, 277 e, 278 ) 279 return ToolResult.error(f"Failed to convert SVG to PNG: {e}") 280 281 log.debug( 282 "%s Mermaid diagram rendered successfully with scale %d, image_data length: %d bytes", 283 log_identifier, 284 scale, 285 len(image_data), 286 ) 287 288 if output_filename: 289 final_output_filename = ensure_correct_extension(output_filename, "png") 290 else: 291 final_output_filename = f"mermaid_diagram_{uuid.uuid4()}.png" 292 293 log.debug( 294 "%s Determined final output filename: %s", 295 log_identifier, 296 final_output_filename, 297 ) 298 299 metadata = { 300 "description": f"PNG image generated from Mermaid syntax. Original requested filename: {output_filename if output_filename else 'N/A'}", 301 "source_format": "mermaid_syntax", 302 "generation_tool": "mermaid_diagram_generator (mermaid-cli)", 303 "generation_timestamp": datetime.now(timezone.utc).isoformat(), 304 "mermaid_title": title if title else "N/A", 305 "mermaid_description": desc if desc else "N/A", 306 } 307 308 log.info("%s Returning PNG as DataObject for artifact storage", log_identifier) 309 310 return ToolResult.ok( 311 "Mermaid diagram rendered successfully.", 312 data={ 313 "mermaid_title": title if title else "N/A", 314 "mermaid_description": desc if desc else "N/A", 315 }, 316 data_objects=[ 317 DataObject( 318 name=final_output_filename, 319 content=image_data, 320 mime_type="image/png", 321 disposition=DataDisposition.ARTIFACT, 322 description=f"PNG diagram from Mermaid syntax{f': {title}' if title else ''}", 323 metadata=metadata, 324 ) 325 ], 326 ) 327 328 except ValueError as e: 329 log.warning( 330 "%s Value error in mermaid_diagram_generator: %s", log_identifier, e 331 ) 332 return ToolResult.error(str(e)) 333 except Exception as e: 334 log.exception( 335 "%s Unexpected error in mermaid_diagram_generator: %s", log_identifier, e 336 ) 337 return ToolResult.error(f"An unexpected error occurred: {e}") 338 339 340 async def _continue_generation() -> Dict[str, Any]: 341 """ 342 Internal tool to signal the LLM to continue a response that was interrupted. 343 This tool is not intended to be called directly by the LLM. 344 """ 345 return { 346 "status": "continue", 347 "message": "You were interrupted due to a token limit. Please review your last output and continue *exactly where you left off with no other commentary*. \nTake care with newlines. If the last output did not finish with a newline, should the next character be a newline? This is especially important for files like csv or yaml files, where a missing newline will break the parser. \nFrom the user's point of view, this continuation needs to be seamless. This function (_continue_generation()) is an internal function that you must not call - it is just here to facilitate providing this message to you.\nIn the unlikely case you have nothing additional to output, then just return nothing.", 348 } 349 350 351 # --- Tool Definitions --- 352 353 _continue_generation_tool_def = BuiltinTool( 354 name="_continue_generation", 355 implementation=_continue_generation, 356 description="INTERNAL TOOL. This tool is used by the system to continue a response that was interrupted. You MUST NOT call this tool directly.", 357 category="internal", 358 required_scopes=[], 359 parameters=adk_types.Schema( 360 type=adk_types.Type.OBJECT, 361 properties={}, 362 required=[], 363 ), 364 examples=[], 365 ) 366 367 convert_file_to_markdown_tool_def = BuiltinTool( 368 name="convert_file_to_markdown", 369 implementation=convert_file_to_markdown, 370 description="Converts an input file artifact to Markdown using the MarkItDown library. The supported input types are those supported by MarkItDown (e.g., PDF, DOCX, XLSX, HTML, CSV, PPTX, ZIP). The output is a new Markdown artifact.", 371 category="general", 372 required_scopes=["tool:general:convert_file"], 373 parameters=adk_types.Schema( 374 type=adk_types.Type.OBJECT, 375 properties={ 376 "input_filename": adk_types.Schema( 377 type=adk_types.Type.STRING, 378 description="The filename (and optional :version) of the input artifact.", 379 ), 380 }, 381 required=["input_filename"], 382 ), 383 examples=[], 384 ) 385 386 mermaid_diagram_generator_tool_def = BuiltinTool( 387 name="mermaid_diagram_generator", 388 implementation=mermaid_diagram_generator, 389 description="Generates a PNG image from Mermaid diagram syntax and saves it as an artifact.", 390 category="general", 391 required_scopes=["tool:general:mermaid"], 392 parameters=adk_types.Schema( 393 type=adk_types.Type.OBJECT, 394 properties={ 395 "mermaid_syntax": adk_types.Schema( 396 type=adk_types.Type.STRING, 397 description="The Mermaid diagram syntax (string).", 398 ), 399 "output_filename": adk_types.Schema( 400 type=adk_types.Type.STRING, 401 description="Optional desired name for the output PNG file. If not provided, a unique name will be generated.", 402 nullable=True, 403 ), 404 }, 405 required=["mermaid_syntax"], 406 ), 407 examples=[], 408 ) 409 410 411 tool_registry.register(_continue_generation_tool_def) 412 tool_registry.register(convert_file_to_markdown_tool_def) 413 tool_registry.register(mermaid_diagram_generator_tool_def)