Update app.py

app.py

@@ -4,6 +4,7 @@ from fpdf import FPDF
 import os
 import time
 from datetime import datetime
+import unittest
 
 # Mistral API key (replace with your key or use environment variable)
 api_key = os.getenv("MISTRAL_API_KEY", "gz6lDXokxgR6cLY72oomALWcm7vhjRzQ")
@@ -22,9 +23,7 @@ def call_mistral_api(prompt):
         ]
     }
     try:
-        start_time = time.time()
         response = requests.post(url, headers=headers, json=payload)
-        latency = round(time.time() - start_time, 2)  # in seconds
         response.raise_for_status()  # Raise an error for bad status codes
         return response.json()['choices'][0]['message']['content']
     except requests.exceptions.HTTPError as err:
@@ -38,30 +37,27 @@ def call_mistral_api(prompt):
 
 # Function to analyze a single requirement
 def analyze_requirement(requirement):
-    latencies = {}
     # Detect requirement type
     type_prompt = f"Classify the following requirement as Functional or Non-Functional:\n\n{requirement}\n\nType:"
-    req_type, latencies["Type"] = call_mistral_api(type_prompt)
+    req_type = call_mistral_api(type_prompt)
 
     # Identify stakeholders
     stakeholders_prompt = f"Identify the stakeholders for the following requirement:\n\n{requirement}\n\nStakeholders:"
-    stakeholders, latencies["Stakeholders"] = call_mistral_api(stakeholders_prompt)
+    stakeholders = call_mistral_api(stakeholders_prompt)
 
     # Classify domain
     domain_prompt = f"Classify the domain for the following requirement (e.g., Bank, Healthcare, etc.):\n\n{requirement}\n\nDomain:"
-    domain, latencies["Domain"] = call_mistral_api(domain_prompt)
+    domain = call_mistral_api(domain_prompt)
 
     # Detect defects
     defects_prompt = f"""Analyze the following requirement and identify ONLY MAJOR defects (e.g., Ambiguity, Incompleteness, etc.). 
     If the requirement is clear and complete, respond with 'No defects.' 
     Requirement: {requirement}
     Defects:"""
-    defectss, latencies["Defects"] = call_mistral_api(defects_prompt)
+    defects = call_mistral_api(defects_prompt)
 
     # Rewrite requirement
-    rewritten, latencies["Rewritten"] = rewrite_requirement(requirement, defects)
-
-    total_latency = sum(latencies.values())
+    rewritten = rewrite_requirement(requirement, defects)
 
     return {
         "Requirement": requirement,
@@ -69,8 +65,7 @@ def analyze_requirement(requirement):
         "Stakeholders": stakeholders,
         "Domain": domain,
         "Defects": defects,
-        "Rewritten": rewritten,
-        "Latency (s)": round(total_latency, 2)
+        "Rewritten": rewritten
     }
 
 # Function to rewrite requirement concisely
@@ -86,7 +81,7 @@ def rewrite_requirement(requirement, defects):
     Defects: {defects}
     
     Rewritten Requirement:"""
-    response, latency = call_mistral_api(prompt)
+    response = call_mistral_api(prompt)
     return response.strip()
 
 # Function to generate a PDF report with professional formatting
@@ -138,6 +133,32 @@ def generate_pdf_report(results):
     pdf.output(pdf_output)
     return pdf_output
 
+# Unit testing section
+class TestRequirementAnalysis(unittest.TestCase):
+    
+    def test_analyze_requirement(self):
+        requirement = "The system must allow users to log in using email and password."
+        result = analyze_requirement(requirement)
+        self.assertIn("Requirement", result)
+        self.assertIn("Type", result)
+        self.assertIn("Stakeholders", result)
+        self.assertIn("Domain", result)
+        self.assertIn("Defects", result)
+        self.assertIn("Rewritten", result)
+    
+    def test_rewrite_requirement_no_defects(self):
+        requirement = "The system should send an email notification when a user signs up."
+        defects = "No defects."
+        result = rewrite_requirement(requirement, defects)
+        self.assertEqual(result, "No modification needed.")
+    
+    def test_generate_pdf_report(self):
+        results = [
+            {"Requirement": "The system should be scalable.", "Type": "Functional", "Stakeholders": "Developers", "Domain": "IT", "Defects": "No defects", "Rewritten": "No modification needed."}
+        ]
+        report = generate_pdf_report(results)
+        self.assertTrue(os.path.exists(report))
+    
 # Streamlit app
 def main():
     st.title("AI Powered Requirement Analysis and Defect Detection using Large Language Model Mistral")
@@ -170,7 +191,6 @@ def main():
                 st.write(f"**Domain:** {result['Domain']}")
                 st.write(f"**Defects:** {result['Defects']}")
                 st.write(f"**Rewritten:** {result['Rewritten']}")
-                st.write(f"**Total Latency:** {result['Latency (s)']} seconds")
                 st.write("---")
 
             # Generate and download PDF report
@@ -183,6 +203,12 @@ def main():
                     mime="application/pdf"
                 )
 
+            # Display test results
+            st.subheader("Unit Test Results")
+            test_suite = unittest.TextTestRunner()
+            test_results = test_suite.run(unittest.TestLoader().loadTestsFromTestCase(TestRequirementAnalysis))
+            st.write(test_results)
+
 # Run the app
 if __name__ == "__main__":
-    main()
+    main()