| """ |
| AI-generated diagram tool. |
| """ |
|
|
| import os |
|
|
| import dotenv |
| from google import genai |
| from google.adk.tools import ToolContext |
| from google.genai import types |
|
|
| dotenv.load_dotenv() |
|
|
| client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"]) |
|
|
|
|
| async def generate_diagram( |
| image_gen_prompt: str, concept_to_explain: str, tool_context: ToolContext |
| ) -> dict: |
| """ |
| Generates a technical diagram for a research paper based on a detailed prompt describing the flow and design requirements. |
| Returns a dictionary with the status and filename or error detail. |
| |
| Args: |
| image_gen_prompt (str): The detailed prompt describing the flow and design requirements. |
| concept_to_explain (str): The concept to explain. |
| tool_context (ToolContext): The context for the tool execution. |
| |
| Returns: |
| dict: Contains 'status' ('success' or 'failed'), and either 'filename' or 'detail'. |
| """ |
| print("Generate diagram tool called!") |
| try: |
| |
| enhanced_prompt = f""" |
| Create a technical, high-quality diagram for a research paper to explain the concept "{concept_to_explain}", based on the following specifications: |
| |
| {image_gen_prompt} |
| |
| Requirements: |
| - Clear, readable, and precise design |
| - High resolution and clean design |
| - Use of appropriate colors and typography |
| - Helps explain the concept to a student |
| |
| Generate a diagram that represents the flow and design requirements. |
| """ |
|
|
| content = types.Content( |
| role="user", |
| parts=[ |
| types.Part.from_text(text=enhanced_prompt), |
| ], |
| ) |
|
|
| response = client.models.generate_content( |
| model="gemini-2.5-flash-image-preview", |
| contents=content, |
| config=types.GenerateContentConfig( |
| temperature=0.8, |
| top_p=0.95, |
| max_output_tokens=8192, |
| response_modalities=["TEXT", "IMAGE"], |
| ), |
| ) |
|
|
| if not response or not getattr(response, "candidates", None): |
| return {"status": "failed", "detail": "No response or candidates from model."} |
|
|
| image_bytes_out = None |
| candidate = response.candidates[0] if response.candidates else None |
| content_out = getattr(candidate, "content", None) if candidate is not None else None |
|
|
| if content_out is not None: |
| for part in getattr(content_out, "parts", []): |
| part_inline = getattr(part, "inline_data", None) |
| part_data = ( |
| getattr(part_inline, "data", None) |
| if part_inline is not None |
| else None |
| ) |
| if part_data: |
| image_bytes_out = part_data |
| break |
|
|
| if not image_bytes_out: |
| return {"status": "failed", "detail": "No image bytes found in model response."} |
|
|
| |
| await tool_context.save_artifact( |
| "diagram.png", |
| types.Part.from_bytes(data=image_bytes_out, mime_type="image/png"), |
| ) |
|
|
| return { |
| "status": "success", |
| "detail": "Diagram generated successfully and stored in artifacts.", |
| "filename": "diagram.png", |
| } |
|
|
| except Exception as e: |
| return {"status": "failed", "detail": f"Error generating diagram: {str(e)}"} |
|
|
