Add ERD type

app.py

@@ -12,6 +12,7 @@ HF_MODEL_MISTRAL = "mistralai/Mistral-7B-Instruct-v0.3"
 HF_MODEL_LLAMA = "meta-llama/Llama-3.3-70B-Instruct"
 
 UML_PROMPTS_DOC_URL = os.environ['UML_PROMPTS_DOC_URL']
+ERD_PROMPTS_DOC_URL = os.environ['ERD_PROMPTS_DOC_URL']
 
 STEP1_SYSTEM_PROMPT = "STEP1 SYSPROMPT"
 STEP1_USER_PROMPT = "STEP1 USERPROMPT"
@@ -23,8 +24,11 @@ STEP3B_SYSTEM_PROMPT = "STEP3B SYSPROMPT"
 STEP3B_USER_PROMPT = "STEP3B USERPROMPT"
 
 
-def fetch_prompts_from_google_doc():
-    response = requests.get(UML_PROMPTS_DOC_URL)
+def fetch_prompts_from_google_doc(diagram_type="UML"):
+    if diagram_type == "UML":
+        response = requests.get(UML_PROMPTS_DOC_URL)
+    elif diagram_type == "ERD":
+        response = requests.get(ERD_PROMPTS_DOC_URL)
     if response.status_code != 200:
         raise Exception("Failed to fetch document")
     
@@ -73,17 +77,17 @@ def extract_plantuml_code(client_openai, uploaded_file, model_choice, prompts):
 
 
 # Step 2: Compare PlantUML Code
-def compare_plantuml(client_openai, client_hf_mistral, client_hf_llama, plantuml_instructor, plantuml_student, model_choice, prompts):
+def compare_plantuml(client_openai, client_hf_mistral, client_hf_llama, plantuml_instructor, plantuml_student, model_choice, prompts, diagram_type="UML"):
 
     st.write("Model: ", model_choice)
 
     user_prompt=f"""
     {prompts[STEP2_USER_PROMPT]}
     
-    **Instructor's UML:**
+    **Instructor's {diagram_type} Diagram:**
     {plantuml_instructor}
 
-    **Student's UML:**
+    **Student's {diagram_type} Diagram:**
     {plantuml_student}
     """
 
@@ -257,13 +261,20 @@ def generate_educator_feedback(client_openai, client_hf_mistral, client_hf_llama
 
 
 # Streamlit app layout
-st.title("LLM-based Analysis and Feedback of a UML Diagram")
+st.set_page_config(
+    page_title="LLM-based Analysis and Feedback of a UML or ER Diagram",
+    page_icon="📝",
+    initial_sidebar_state="expanded",
+)
+
+st.title("LLM-based Analysis and Feedback of a UML or ERD Diagram")
 st.write("The pipeline consists of three steps:")
-st.write("1. Extract PlantUML code from the uploaded UML diagrams using GPT-4o or GPT-4o Mini.")
+st.write("1. Extract PlantUML code from the uploaded UML or ER diagrams using GPT-4o or GPT-4o Mini.")
 st.write("2. Compare the extracted PlantUML code.")
 st.write("3. Analyse the differences and present them in a structured format.")
 
-prompts = fetch_prompts_from_google_doc()
+diagram_type = st.selectbox("Select the diagram type", ["UML", "ERD"])
+prompts = fetch_prompts_from_google_doc(diagram_type)
 
 openai_api_key = st.text_input("OpenAI API key", type="password")
 hf_api_key = st.text_input("Hugging Face API key", type="password")
@@ -282,22 +293,22 @@ if openai_api_key and hf_api_key:
     col1, col2 = st.columns(2)
     with col1:
         uploaded_instructor_solution = st.file_uploader(
-            "Upload Instructor UML Diagram", type=["jpg", "jpeg", "png"]
+            "Upload Instructor " + ("UML" if diagram_type == 'UML' else "ER") + " Diagram", type=["jpg", "jpeg", "png"]
         )
     with col2:
         uploaded_student_solution = st.file_uploader(
-            "Upload Student UML Diagram", type=["jpg", "jpeg", "png"]
+            "Upload Student " + ("UML" if diagram_type == 'UML' else "ER") + " Diagram", type=["jpg", "jpeg", "png"]
         )
 
     if (uploaded_instructor_solution is not None and uploaded_student_solution is not None):
         try:
             with st.spinner(
-                "Extracting PlantUML code from the uploaded UML diagrams..."
+                "Extracting PlantUML code from the uploaded " + ("UML" if diagram_type == 'UML' else "ER") + " diagrams..."
             ):
                 with col1:
                     st.image(
                         uploaded_instructor_solution,
-                        caption="Uploaded Instructor UML Diagram",
+                        caption="Uploaded Instructor " + ("UML" if diagram_type == 'UML' else "ER") + " Diagram",
                         use_container_width=True,
                     )
                     st.write("")
@@ -308,7 +319,7 @@ if openai_api_key and hf_api_key:
                     st.write("")
                     st.image(
                         uploaded_student_solution,
-                        caption="Uploaded Student UML Diagram",
+                        caption="Uploaded Student " + ("UML" if diagram_type == 'UML' else "ER") + " Diagram",
                         use_container_width=True,
                     )
                     st.write("")
@@ -331,8 +342,8 @@ if openai_api_key and hf_api_key:
                     height=600,
                 )
 
-            st.subheader("Step 2: UML Diagram Comparison")
-            with st.spinner("Comparing instructor and student UML diagrams..."):
+            st.subheader("Step 2: " + ("UML" if diagram_type == 'UML' else "ER") + " Diagram Comparison")
+            with st.spinner("Comparing instructor and student " + ("UML" if diagram_type == 'UML' else "ER") + " diagrams..."):
                 differences = compare_plantuml(
                     client_openai,
                     client_hf_mistral,
@@ -340,7 +351,8 @@ if openai_api_key and hf_api_key:
                     plantuml_instructor_solution,
                     plantuml_student_solution,
                     model_choice_step2,
-                    prompts
+                    prompts,
+                    diagram_type
                 )
                 with st.expander("View differences"):
                     for difference in differences.split("\n"):