Update app.py

app.py

@@ -1,4 +1,3 @@
-
 import streamlit as st
 import requests
 from fpdf import FPDF
@@ -16,6 +15,7 @@ groq_client = groq.Client(api_key=groq_api_key)
 
 # Function to call Mistral API
 def call_mistral_api(prompt):
+    start_time = time.time()
     url = "https://api.mistral.ai/v1/chat/completions"
     headers = {
         "Authorization": f"Bearer {mistral_api_key}",
@@ -30,18 +30,23 @@ def call_mistral_api(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']
+        end_time = time.time()
+        speed = end_time - start_time
+        content = response.json()['choices'][0]['message']['content']
+        confidence = len(content.split())  # Simulate confidence with response length
+        return content, speed, confidence
     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}"
+        return f"HTTP Error: {err}", 0, 0
     except Exception as err:
-        return f"Error: {err}"
+        return f"Error: {err}", 0, 0
 
 # Function to call Groq API
 def call_groq_api(prompt):
+    start_time = time.time()
     try:
         response = groq_client.chat.completions.create(
             model="llama-3.3-70b-versatile",  # Correct model name
@@ -49,33 +54,49 @@ def call_groq_api(prompt):
                 {"role": "user", "content": prompt}
             ]
         )
-        return response.choices[0].message.content
+        end_time = time.time()
+        speed = end_time - start_time
+        content = response.choices[0].message.content
+        confidence = len(content.split())  # Simulate confidence with response length
+        return content, speed, confidence
     except Exception as err:
         st.error(f"Error: {err}")
-        return f"Error: {err}"
+        return f"Error: {err}", 0, 0
 
 # Function to analyze a single requirement using both models
 def analyze_requirement(requirement):
     # Use Mistral for classification and domain identification
     type_prompt = f"Classify the following requirement as Functional or Non-Functional in one word:\n\n{requirement}\n\nType:"
-    req_type = call_mistral_api(type_prompt).strip()
+    req_type, type_speed, type_confidence = call_mistral_api(type_prompt)
+    req_type = req_type.strip()
 
     domain_prompt = f"Classify the domain for the following requirement in one word (e.g., E-commerce, Education, etc.):\n\n{requirement}\n\nDomain:"
-    domain = call_mistral_api(domain_prompt).strip()
+    domain, domain_speed, domain_confidence = call_mistral_api(domain_prompt)
+    domain = domain.strip()
 
     # Use Groq for defect analysis and rewriting
     defects_prompt = f"""List ONLY the major defects in the following requirement (e.g., Ambiguity, Incompleteness, etc.) in 1-2 words each:\n\n{requirement}\n\nDefects:"""
-    defects = call_groq_api(defects_prompt).strip()
+    defects, defects_speed, defects_confidence = call_groq_api(defects_prompt)
+    defects = defects.strip()
 
     rewritten_prompt = f"""Rewrite the following requirement in 1-2 sentences to address the defects:\n\n{requirement}\n\nRewritten:"""
-    rewritten = call_groq_api(rewritten_prompt).strip()
+    rewritten, rewritten_speed, rewritten_confidence = call_groq_api(rewritten_prompt)
+    rewritten = rewritten.strip()
 
     return {
         "Requirement": requirement,
         "Type": req_type,
         "Domain": domain,
         "Defects": defects,
-        "Rewritten": rewritten
+        "Rewritten": rewritten,
+        "Type_Speed": type_speed,
+        "Type_Confidence": type_confidence,
+        "Domain_Speed": domain_speed,
+        "Domain_Confidence": domain_confidence,
+        "Defects_Speed": defects_speed,
+        "Defects_Confidence": defects_confidence,
+        "Rewritten_Speed": rewritten_speed,
+        "Rewritten_Confidence": rewritten_confidence
     }
 
 # Function to generate a PDF report
@@ -118,6 +139,14 @@ def generate_pdf_report(results):
         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=f"Type Speed: {result['Type_Speed']:.2f}s", align='L')
+        pdf.multi_cell(200, 10, txt=f"Type Confidence: {result['Type_Confidence']}", align='L')
+        pdf.multi_cell(200, 10, txt=f"Domain Speed: {result['Domain_Speed']:.2f}s", align='L')
+        pdf.multi_cell(200, 10, txt=f"Domain Confidence: {result['Domain_Confidence']}", align='L')
+        pdf.multi_cell(200, 10, txt=f"Defects Speed: {result['Defects_Speed']:.2f}s", align='L')
+        pdf.multi_cell(200, 10, txt=f"Defects Confidence: {result['Defects_Confidence']}", align='L')
+        pdf.multi_cell(200, 10, txt=f"Rewritten Speed: {result['Rewritten_Speed']:.2f}s", align='L')
+        pdf.multi_cell(200, 10, txt=f"Rewritten Confidence: {result['Rewritten_Confidence']}", align='L')
         pdf.multi_cell(200, 10, txt="-" * 50, align='L')
         pdf.ln(5)  # Add some space between requirements
 
@@ -155,6 +184,14 @@ def main():
                 st.write(f"**Domain:** {result['Domain']}")
                 st.write(f"**Defects:** {result['Defects']}")
                 st.write(f"**Rewritten:** {result['Rewritten']}")
+                st.write(f"**Type Speed:** {result['Type_Speed']:.2f}s")
+                st.write(f"**Type Confidence:** {result['Type_Confidence']}")
+                st.write(f"**Domain Speed:** {result['Domain_Speed']:.2f}s")
+                st.write(f"**Domain Confidence:** {result['Domain_Confidence']}")
+                st.write(f"**Defects Speed:** {result['Defects_Speed']:.2f}s")
+                st.write(f"**Defects Confidence:** {result['Defects_Confidence']}")
+                st.write(f"**Rewritten Speed:** {result['Rewritten_Speed']:.2f}s")
+                st.write(f"**Rewritten Confidence:** {result['Rewritten_Confidence']}")
                 st.write("---")
 
             # Generate and download PDF report
@@ -169,5 +206,4 @@ def main():
 
 # Run the app
 if __name__ == "__main__":
-    main()
-
+    main()