Update models/openai_image_to_json.py

models/openai_image_to_json.py

@@ -6,90 +6,96 @@ import os
 from dotenv import load_dotenv
 import json
 
-# Load the .env file
-load_dotenv()
+# Using Gemini instead of OpenAI for this deployment only
+from gemini_image_to_json import fetch_gemini_response
 
-# Get the API key from the environment
-api_key = os.getenv('OPENAI_API_KEY')
+openaiprocess_image_to_json = fetch_gemini_response
 
-# Function to encode the image
-def encode_image(image):
-    # Convert the image to RGB if it has an alpha channel
-    if image.mode == 'RGBA':
-        image = image.convert('RGB')
+# # Load the .env file
+# load_dotenv()
+
+# # Get the API key from the environment
+# api_key = os.getenv('OPENAI_API_KEY')
+
+# # Function to encode the image
+# def encode_image(image):
+#     # Convert the image to RGB if it has an alpha channel
+#     if image.mode == 'RGBA':
+#         image = image.convert('RGB')
     
-    buffered = BytesIO()
-    image.save(buffered, format="JPEG")
-    return base64.b64encode(buffered.getvalue()).decode('utf-8')
+#     buffered = BytesIO()
+#     image.save(buffered, format="JPEG")
+#     return base64.b64encode(buffered.getvalue()).decode('utf-8')
 
-def openaiprocess_image_to_json(image):
+# def openaiprocess_image_to_json(image):
     
-    print(f'fetching openai response')
+#     print(f'fetching openai response')
     
-    # Encode the image
-    base64_image = encode_image(image)
+#     # Encode the image
+#     base64_image = encode_image(image)
 
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {api_key}"
-    }
+#     headers = {
+#         "Content-Type": "application/json",
+#         "Authorization": f"Bearer {api_key}"
+#     }
 
-    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 as follows:
-    {
-    "nodes": [
-        {"id": "1", "label": string},
-        {"id": "2", "label": string},...
-        ],
-    "edges": [
-        {"source": SOURCE_ID, "target": TARGET_ID, "type": "->"},
-        {"source": SOURCE_ID, "target": TARGET_ID, "type": "->"},...
-        ]
-    }
-    Only return valid python dictionary, dont include (line jump)n in it, dont include spaces, only a dictionary. Do not include any other text outside the Dictionary structure. Make sure that i will get a valid Python dictionary.
-    make sure that what you return as json_string i can use it in python in this function: json.loads(json_string)
-    Now extract the entities and relationships from this image:
-    '''
+#     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 as follows:
+#     {
+#     "nodes": [
+#         {"id": "1", "label": string},
+#         {"id": "2", "label": string},...
+#         ],
+#     "edges": [
+#         {"source": SOURCE_ID, "target": TARGET_ID, "type": "->"},
+#         {"source": SOURCE_ID, "target": TARGET_ID, "type": "->"},...
+#         ]
+#     }
+#     Only return valid python dictionary, dont include (line jump)n in it, dont include spaces, only a dictionary. Do not include any other text outside the Dictionary structure. Make sure that i will get a valid Python dictionary.
+#     make sure that what you return as json_string i can use it in python in this function: json.loads(json_string)
+#     Now extract the entities and relationships from this image:
+#     '''
 
-    payload = {
-        "model": "gpt-4o",
-        "messages": [
-            {
-                "role": "user",
-                "content": [
-                    {
-                    "type": "text",
-                    "text": PROMPT
-                    },
-                    {
-                    "type": "image_url",
-                    "image_url": {
-                        "url": f"data:image/jpeg;base64,{base64_image}"
-                    }
-                    }
-                ]
-            }
-        ]
-    }
+#     payload = {
+#         "model": "gpt-4o",
+#         "messages": [
+#             {
+#                 "role": "user",
+#                 "content": [
+#                     {
+#                     "type": "text",
+#                     "text": PROMPT
+#                     },
+#                     {
+#                     "type": "image_url",
+#                     "image_url": {
+#                         "url": f"data:image/jpeg;base64,{base64_image}"
+#                     }
+#                     }
+#                 ]
+#             }
+#         ]
+#     }
 
-    # Send the request to the OpenAI API
-    response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
+#     # Send the request to the OpenAI API
+#     response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
     
-    # Parse the response
-    response_data = response.json()
-    print(response_data)
+#     # Parse the response
+#     response_data = response.json()
+#     print(response_data)
+
+#     # Extract the JSON graph data from the response
+#     if "choices" in response_data and response_data["choices"]:
+#         content = response_data["choices"][0]["message"]["content"]
+#         try:
+#             graph_data = content
+#         except json.JSONDecodeError as e:
+#             print("Failed:", e)
+#             graph_data = None
+#     else:
+#         raise ValueError("No valid response from OpenAI API")
 
-    # Extract the JSON graph data from the response
-    if "choices" in response_data and response_data["choices"]:
-        content = response_data["choices"][0]["message"]["content"]
-        try:
-            graph_data = content
-        except json.JSONDecodeError as e:
-            print("Failed:", e)
-            graph_data = None
-    else:
-        raise ValueError("No valid response from OpenAI API")
+#     return graph_data
 
-    return graph_data