Update app.py

app.py

@@ -1,4 +1,3 @@
-import unittest
 import streamlit as st
 import requests
 from fpdf import FPDF
@@ -6,57 +5,142 @@ import os
 import time
 from datetime import datetime
 
-# Mock function to call Mistral API (replace with real call for actual usage)
-def call_mistral_api(prompt):
-    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!")
+# Mistral API key (replace with your key or use environment variable)
+api_key = os.getenv("MISTRAL_API_KEY", "gz6lDXokxgR6cLY72oomALWcm7vhjRzQ")
 
-        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)}")
+# Function to call Mistral API with rate limit handling
+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):
+    # Define test case expectations
+    expected_outcomes = {
+        "Type": "Functional or Non-Functional",
+        "Stakeholders": "List of stakeholders for the requirement",
+        "Domain": "Classify the domain (e.g., Bank, Healthcare)",
+        "Defects": "List of major defects or 'No defects' if none",
+        "Rewritten": "Concise, defect-free rewritten 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 both the result and expected outcome comparison
+    return {
+        "Requirement": requirement,
+        "Type": {"Expected": expected_outcomes["Type"], "Actual": req_type},
+        "Stakeholders": {"Expected": expected_outcomes["Stakeholders"], "Actual": stakeholders},
+        "Domain": {"Expected": expected_outcomes["Domain"], "Actual": domain},
+        "Defects": {"Expected": expected_outcomes["Defects"], "Actual": defects},
+        "Rewritten": {"Expected": expected_outcomes["Rewritten"], "Actual": 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)
+        
+        # Display expected vs actual results
+        for key in ["Type", "Stakeholders", "Domain", "Defects", "Rewritten"]:
+            pdf.multi_cell(200, 10, txt=f"{key}:\nExpected: {result[key]['Expected']}\nActual: {result[key]['Actual']}", 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
 
 # Streamlit app
 def main():
@@ -71,7 +155,7 @@ def main():
         # Split by periods or newlines
         requirements = [req.strip() for req in input_text.replace("\n", ".").split(".") if req.strip()]
 
-    # Analyze requirements and display results
+    # Analyze requirements
     if st.button("Analyze Requirements"):
         if not requirements:
             st.warning("Please enter requirements.")
@@ -85,11 +169,8 @@ def main():
             st.subheader("Analysis Results")
             for i, result in enumerate(results, start=1):
                 st.write(f"### Requirement R{i}: {result['Requirement']}")
-                st.write(f"**Type:** {result['Type']}")
-                st.write(f"**Stakeholders:** {result['Stakeholders']}")
-                st.write(f"**Domain:** {result['Domain']}")
-                st.write(f"**Defects:** {result['Defects']}")
-                st.write(f"**Rewritten:** {result['Rewritten']}")
+                for key in ["Type", "Stakeholders", "Domain", "Defects", "Rewritten"]:
+                    st.write(f"**{key}:** Expected: {result[key]['Expected']} | Actual: {result[key]['Actual']}")
                 st.write("---")
 
             # Generate and download PDF report
@@ -102,12 +183,6 @@ def main():
                     mime="application/pdf"
                 )
 
-            # 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__":
     main()