ayankumar/SDLC-Assistant-MultiAgent

1import google.generativeai as genai2from PIL import Image3import io4 5with open("keys/.gemini_key", "r") as f:6    genai.configure(api_key=f.read().strip())7 8model = genai.GenerativeModel("gemini-1.5-flash")9 10def generate_questions(name: str, prompt: str) -> list:11    print(f"{name} agent is running the prompt : {prompt}")12    questions = model.generate_content(prompt).text.strip().split("\n")13    return [q.strip("-• ") for q in questions if q]14 15def generate_output(name: str, prompt: str) -> dict:16    print(f"{name} agent is running the prompt : {prompt}")17    response = model.generate_content(prompt)18    # 🖼️ Extract image if available19    image_data = None20    if hasattr(response, "media"):21        for media in response.media:22            if hasattr(media, "mime_type") and media.mime_type.startswith("image"):23                image_bytes = media.data24                image_data = Image.open(io.BytesIO(image_bytes))25                break26    return {27        "text": response.text.strip() if hasattr(response, "text") else "",28        "image": image_data  # Reserved for future diagram generation29    }