Update app.py

app.py

@@ -22,7 +22,9 @@ 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:
@@ -36,27 +38,30 @@ 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 = call_mistral_api(type_prompt)
+    req_type, latencies["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)
+    stakeholders, latencies["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)
+    domain, latencies["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)
+    defectss, latencies["Defects"] = call_mistral_api(defects_prompt)
 
     # Rewrite requirement
-    rewritten = rewrite_requirement(requirement, defects)
+    rewritten, latencies["Rewritten"] = rewrite_requirement(requirement, defects)
+
+    total_latency = sum(latencies.values())
 
     return {
         "Requirement": requirement,
@@ -64,7 +69,8 @@ def analyze_requirement(requirement):
         "Stakeholders": stakeholders,
         "Domain": domain,
         "Defects": defects,
-        "Rewritten": rewritten
+        "Rewritten": rewritten,
+        "Latency (s)": round(total_latency, 2)
     }
 
 # Function to rewrite requirement concisely
@@ -80,7 +86,7 @@ def rewrite_requirement(requirement, defects):
     Defects: {defects}
     
     Rewritten Requirement:"""
-    response = call_mistral_api(prompt)
+    response, latency = call_mistral_api(prompt)
     return response.strip()
 
 # Function to generate a PDF report with professional formatting
@@ -164,6 +170,7 @@ 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