Update backend.py

backend.py

@@ -1,13 +1,16 @@
 import csv
 import io
+import requests
 import json
 import html  # For escaping HTML characters
 from bs4 import BeautifulSoup
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from openai import OpenAI 
 
-# Load the tokenizer and model directly
-tokenizer = AutoTokenizer.from_pretrained("mattshumer/Reflection-Llama-3.1-70B")
-model = AutoModelForCausalLM.from_pretrained("mattshumer/Reflection-Llama-3.1-70B")
+# Initialize OpenAI API with Meta's Llama 3.1 405B model
+client = OpenAI(
+    base_url="https://integrate.api.nvidia.com/v1",
+    api_key="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"
+)
 
 def clean_test_case_output(text):
     """
@@ -21,7 +24,7 @@ def clean_test_case_output(text):
 def generate_testcases(user_story):
     """
     Generates advanced QA test cases based on a provided user story by interacting 
-    with the Reflection-Llama-3.1-70B model. The prompt is refined for clarity, 
+    with Nvidia's llama model API. The prompt is refined for clarity, 
     and the output is processed for better quality.
     
     :param user_story: A string representing the user story for which to generate test cases.
@@ -30,49 +33,60 @@ def generate_testcases(user_story):
 
     # Few-shot learning examples to guide the model
     few_shot_examples = """
+
     "if its not a DropBury or ODAC Portal User Story, then we perform testing in Tech360 iOS App"
-    "Generate as many test cases as possible, minimum 6, maximum it can be anything"
+    "Generate as many as testcases possible minimum 6 ,maximum it can be anything"
     "Understand the story thoroughly"
     "If it's a DropBury or ODAC Portal User Story, then we perform testing in ODAC Portal"
     """
 
     # Combine the few-shot examples with the user story for the model to process
-    prompt = few_shot_examples + f"\nUser Story: {user_story}\n"
-
-    # Tokenize the prompt
-    inputs = tokenizer(prompt, return_tensors="pt")
-
-    # Generate text with the model
-    outputs = model.generate(
-        **inputs,
-        max_length=4096,
-        temperature=0.03,
-        top_p=0.7,
-        do_sample=False
-    )
-
-    # Decode the generated text
-    test_cases_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Ensure the entire response is captured before cleaning
-    if test_cases_text.strip() == "":
-        return [{"test_case": "No test cases generated or output was empty."}]
-
-    # Clean the output by unescaping HTML entities and replacing <br> tags
-    test_cases_text = clean_test_case_output(test_cases_text)
+    prompt = few_shot_examples + f"\nUser Story: {user_story}\n" 
 
     try:
-        # Try to parse the output as JSON, assuming the model returns structured test cases
-        test_cases = json.loads(test_cases_text)
-        if isinstance(test_cases, list):
-            return test_cases  # Return structured test cases
-
-        else:
-            return [{"test_case": test_cases_text}]  # Return as a list with the text wrapped in a dict
-
-    except json.JSONDecodeError:
-        # Fallback: return the raw text if JSON parsing fails
-        return [{"test_case": test_cases_text}]
+        # Call the Nvidia llama API with the refined prompt
+        completion = client.chat.completions.create(
+            model="meta/llama-3.1-405b-instruct",  # Using llama 3.1 450b
+            messages=[
+                {"role": "user", "content": prompt}
+            ],
+            temperature=0.03,  # Further lowering temperature for precise and deterministic output
+            top_p=0.7,  # Prioritize high-probability tokens even more
+            max_tokens=4096,  # Increase max tokens to allow longer content
+            stream=True  # Streaming the response for faster retrieval
+        )
+
+        # Initialize an empty string to accumulate the response
+        test_cases_text = ""
+        
+        # Accumulate the response from the streaming chunks
+        for chunk in completion:
+            if chunk.choices[0].delta.content is not None:
+                test_cases_text += chunk.choices[0].delta.content
+
+        # Ensure the entire response is captured before cleaning
+        if test_cases_text.strip() == "":
+            return [{"test_case": "No test cases generated or output was empty."}]
+        
+        # Clean the output by unescaping HTML entities and replacing <br> tags
+        test_cases_text = clean_test_case_output(test_cases_text)
+
+        try:
+            # Try to parse the output as JSON, assuming the model returns structured test cases
+            test_cases = json.loads(test_cases_text)
+            if isinstance(test_cases, list):
+                return test_cases  # Return structured test cases
+
+            else:
+                return [{"test_case": test_cases_text}]  # Return as a list with the text wrapped in a dict
+
+        except json.JSONDecodeError:
+            # Fallback: return the raw text if JSON parsing fails
+            return [{"test_case": test_cases_text}]
+    
+    except requests.exceptions.RequestException as e:
+        print(f"API request failed: {str(e)}")
+        return []
 
 # Export test cases in CSV format
 def export_test_cases(test_cases, format='csv'):