Update backend.py

backend.py

@@ -1,16 +1,13 @@
 import csv
 import io
-import requests
 import json
 import html  # For escaping HTML characters
 from bs4 import BeautifulSoup
-from openai import OpenAI 
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# 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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"
-)
+# 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")
 
 def clean_test_case_output(text):
     """
@@ -24,7 +21,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 Nvidia's llama model API. The prompt is refined for clarity, 
+    with the Reflection-Llama-3.1-70B model. 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.
@@ -33,60 +30,49 @@ 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 as testcases possible minimum 6 ,maximum it can be anything"
+    "Generate as many test cases as 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" 
+    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)
 
     try:
-        # 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 []
+        # 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}]
 
 # Export test cases in CSV format
 def export_test_cases(test_cases, format='csv'):