models/gemini_image_to_json.py

import google.generativeai as genai

from dotenv import load_dotenv
import os

load_dotenv()
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

genai.configure(api_key=GOOGLE_API_KEY)

# gemini-1.5-pro only gives 50 requests per day. check https://ai.google.dev/pricing for more details
# model = genai.GenerativeModel('gemini-1.5-pro',
model = genai.GenerativeModel(
    "gemini-1.5-flash",
    # Set the `response_mime_type` to output JSON
    # Pass the schema object to the `response_schema` field
    generation_config={
        "response_mime_type": "application/json",
        "temperature": 0.0,
    },
)
#  "response_schema": Recipe, 'max_output_tokens':4000})

PROMPT = """
You are responsible for extracting the entities (nodes) and relationships (edges) from the images of mind maps. The mind maps are for Object Oriented Programming.
Don't make up facts, just extracts them. Do not create new entity types that aren't mentioned in the image, and at the same time don't miss anything.
Give the output in JSON format with this schema:
{
  "nodes": [{"id": "1", "label": string},{"id": "2", "label": string}],"edges": [{"source": SOURCE_ID, "target": TARGET_ID, "type": "->"},{"source": SOURCE_ID, "target": TARGET_ID, "type": "->"}]
}

Now extract the entities and relationships from this image:
"""


def fetch_gemini_response(mind_map_image):

    print("fetching gemini response")
    response = model.generate_content([PROMPT, mind_map_image], stream=False)

    return response.text