File size: 2,814 Bytes
0892c77 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | import requests
import json
import os
from dotenv import load_dotenv
# Step 1: Set your Hugging Face API Key
load_dotenv()
HUGGINGFACE_API_KEY = os.getenv("HF_TOKEN")
# Step 2: Load Prompt Template
def load_prompt():
with open("prompts/basic_prompt.txt", "r") as file:
return file.read()
# Step 3: Define the function to call HF Inference API
def generate_test_cases(user_story):
prompt_template = load_prompt()
# Fill in the user story
prompt = prompt_template.replace("{user_story}", user_story)
# Define request payload
payload = {
"inputs": prompt,
"parameters": {
"max_new_tokens": 1000,
"temperature": 0.5,
"top_p": 0.9
}
}
# Call Hugging Face Inference API (you can use a model like 'mistralai/mistral-7b-instruct' or 'google/flan-t5-xxl')
response = requests.post(
f"https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.3",
headers={"Authorization": f"Bearer {HUGGINGFACE_API_KEY}"},
json=payload
)
# Print and inspect the response
response_json = response.json()
#print("Response data:", response_json) # Check the structure
# Handle both cases (list or dict) for the response structure
if isinstance(response_json, list): # If the response is a list
result = response_json[0].get('generated_text', 'No text generated')
elif isinstance(response_json, dict): # If it's a dictionary
result = response_json.get('generated_text', 'No text generated')
else:
result = 'Unexpected response format'
# Print the final generated test cases
print("Generated Test Cases:", result)
# # If the response is a list, access the first item (which is the dictionary)
# if isinstance(response_json, list):
# result = response_json[0].get('generated_text', 'No text generated')
# else:
# result = response_json.get('generated_text', 'No text generated')
#
# print("Generated Text:", result)
#
# if response.status_code != 200:
# raise Exception(f"API call failed: {response.text}")
#
# result = response.json()
# generated_text = result.get("generated_text", "")
#
# return generated_text
# Step 4: Run the generator
if __name__ == "__main__":
# Example user story
user_story = """
As a user, I want to reset my password so that I can regain access to my account if I forget my password.
"""
#print("Generating test cases for the following user story:\n")
#print(user_story)
try:
test_cases = generate_test_cases(user_story)
#print("\n--- Generated Test Cases ---\n")
#print(test_cases)
except Exception as e:
print(f"Failed to generate test cases: {e}")
|