Update app.py

app.py

@@ -1,164 +1,63 @@
+import unittest
 import streamlit as st
 import requests
 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")
 
-# Function to call Mistral API with rate limit handling
+# Mock function to call Mistral API (replace with real call for actual usage)
 def call_mistral_api(prompt):
-    url = "https://api.mistral.ai/v1/chat/completions"
-    headers = {
-        "Authorization": f"Bearer {api_key}",
-        "Content-Type": "application/json"
-    }
-    payload = {
-        "model": "mistral-medium",
-        "messages": [
-            {"role": "user", "content": prompt}
-        ]
-    }
-    try:
-        response = requests.post(url, headers=headers, json=payload)
-        response.raise_for_status()  # Raise an error for bad status codes
-        return response.json()['choices'][0]['message']['content']
-    except requests.exceptions.HTTPError as err:
-        if response.status_code == 429:  # Rate limit exceeded
-            st.warning("Rate limit exceeded. Please wait a few seconds and try again.")
-            time.sleep(5)  # Wait for 5 seconds before retrying
-            return call_mistral_api(prompt)  # Retry the request
-        return f"HTTP Error: {err}"
-    except Exception as err:
-        return f"Error: {err}"
-
-# Function to analyze a single requirement
-def analyze_requirement(requirement):
-    # Detect requirement type
-    type_prompt = f"Classify the following requirement as Functional or Non-Functional:\n\n{requirement}\n\nType:"
-    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 = 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 = 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:"""
-    defects = call_mistral_api(defects_prompt)
-
-    # Rewrite requirement
-    rewritten = rewrite_requirement(requirement, defects)
-
-    return {
-        "Requirement": requirement,
-        "Type": req_type,
-        "Stakeholders": stakeholders,
-        "Domain": domain,
-        "Defects": defects,
-        "Rewritten": rewritten
-    }
-
-# Function to rewrite requirement concisely
-def rewrite_requirement(requirement, defects):
-    if "no defects" in defects.lower():
-        return "No modification needed."
-
-    # If defects are found, generate a concise and clear rewritten requirement
-    prompt = f"""Rewrite the following requirement to address the defects listed below. Ensure the rewritten requirement is clear, concise, and free of defects. It should be no more than 1-2 sentences.
-    
-    Original Requirement: {requirement}
-    
-    Defects: {defects}
-    
-    Rewritten Requirement:"""
-    response = call_mistral_api(prompt)
-    return response.strip()
-
-# Function to generate a PDF report with professional formatting
-def generate_pdf_report(results):
-    pdf = FPDF()
-    pdf.add_page()
-    pdf.set_font("Arial", size=12)
-
-    # Add watermark
-    pdf.set_font("Arial", 'B', 50)
-    pdf.set_text_color(230, 230, 230)  # Light gray color for watermark
-    pdf.rotate(45)  # Rotate the text for watermark effect
-    pdf.text(60, 150, "AI Powered Requirement Analysis")
-    pdf.rotate(0)  # Reset rotation
-
-    # Add title and date/time
-    pdf.set_font("Arial", 'B', 16)
-    pdf.set_text_color(0, 0, 0)  # Black color for title
-    pdf.cell(200, 10, txt="AI Powered Requirement Analysis and Defect Detection using LLM Model Mistral", ln=True, align='C')
-    pdf.set_font("Arial", size=12)
-    pdf.cell(200, 10, txt=f"Report Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", ln=True, align='C')
-    pdf.ln(10)  # Add some space
-
-    # Add requirements analysis
-    pdf.set_font("Arial", size=12)
-    for i, result in enumerate(results, start=1):
-        # Check if we need a new page
-        if pdf.get_y() > 250:  # If the content is near the bottom of the page
-            pdf.add_page()  # Add a new page
-            pdf.set_font("Arial", 'B', 16)
-            pdf.cell(200, 10, txt="AI Powered Requirement Analysis and Defect Detection using LLM Model Mistral", ln=True, align='C')
-            pdf.set_font("Arial", size=12)
-            pdf.cell(200, 10, txt=f"Report Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", ln=True, align='C')
-            pdf.ln(10)  # Add some space
-
-        # Add requirement details
-        pdf.set_font("Arial", 'B', 14)
-        pdf.multi_cell(200, 10, txt=f"Requirement R{i}: {result['Requirement']}", align='L')
-        pdf.set_font("Arial", size=12)
-        pdf.multi_cell(200, 10, txt=f"Type: {result['Type']}", align='L')
-        pdf.multi_cell(200, 10, txt=f"Stakeholders: {result['Stakeholders']}", align='L')
-        pdf.multi_cell(200, 10, txt=f"Domain: {result['Domain']}", align='L')
-        pdf.multi_cell(200, 10, txt=f"Defects: {result['Defects']}", align='L')
-        pdf.multi_cell(200, 10, txt=f"Rewritten: {result['Rewritten']}", align='L')
-        pdf.multi_cell(200, 10, txt="-" * 50, align='L')
-        pdf.ln(5)  # Add some space between requirements
-
-    pdf_output = "requirements_report.pdf"
-    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))
-    
+    return "Functional"  # Mock response for testing
+
+# Test case class to check requirement analysis functionality
+class RequirementTestCase(unittest.TestCase):
+
+    def test_requirement_type(self):
+        requirement = "The system shall allow users to login."
+        expected = "Functional"
+        actual = call_mistral_api(f"Classify the following requirement as Functional or Non-Functional:\n\n{requirement}\n\nType:")
+        self.assertEqual(actual, expected, f"Test failed for requirement type. Expected: {expected}, but got: {actual}")
+
+    def test_stakeholders_identification(self):
+        requirement = "The system shall allow users to login."
+        expected = "User, System Administrator"
+        actual = call_mistral_api(f"Identify the stakeholders for the following requirement:\n\n{requirement}\n\nStakeholders:")
+        self.assertEqual(actual, expected, f"Test failed for stakeholders. Expected: {expected}, but got: {actual}")
+
+    def test_defects_identification(self):
+        requirement = "The system shall allow users to login."
+        expected = "No defects"
+        actual = call_mistral_api(f"Analyze the following requirement and identify ONLY MAJOR defects:\n\n{requirement}\n\nDefects:")
+        self.assertEqual(actual, expected, f"Test failed for defects. Expected: {expected}, but got: {actual}")
+
+    # Add more test cases as needed...
+
+# Custom runner to display detailed results
+class DetailedTextTestRunner(unittest.TextTestRunner):
+    def getDescription(self, test):
+        return f"{test.__class__.__name__} ({test._testMethodName})"
+
+    def run(self, test):
+        result = super().run(test)
+        self.display_results(result)
+        return result
+
+    def display_results(self, result):
+        if result.errors or result.failures:
+            st.subheader("Test Results")
+            for test, traceback in result.errors + result.failures:
+                st.write(f"**Test Case:** {test}")
+                st.write(f"**Error/Failure:** {traceback}")
+        else:
+            st.success("All tests passed successfully!")
+
+        st.subheader("Test Summary")
+        st.write(f"Total Tests Run: {result.testsRun}")
+        st.write(f"Failures: {len(result.failures)}")
+        st.write(f"Errors: {len(result.errors)}")
+        st.write(f"Skipped: {len(result.skipped)}")
+
 # Streamlit app
 def main():
     st.title("AI Powered Requirement Analysis and Defect Detection using Large Language Model Mistral")
@@ -172,7 +71,7 @@ def main():
         # Split by periods or newlines
         requirements = [req.strip() for req in input_text.replace("\n", ".").split(".") if req.strip()]
 
-    # Analyze requirements
+    # Analyze requirements and display results
     if st.button("Analyze Requirements"):
         if not requirements:
             st.warning("Please enter requirements.")
@@ -203,11 +102,11 @@ 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 unit tests
+            st.subheader("Running Unit Tests...")
+            test_suite = unittest.defaultTestLoader.loadTestsFromTestCase(RequirementTestCase)
+            test_runner = DetailedTextTestRunner(verbosity=2)
+            test_runner.run(test_suite)
 
 # Run the app
 if __name__ == "__main__":