File size: 3,370 Bytes
93ac83f 1c0aaf7 93ac83f | 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 90 91 92 | import os
import gradio as gr
from dotenv import load_dotenv
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.llms import HuggingFaceHub
from langchain_core.documents import Document
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain_community.tools.tavily_search.tool import TavilySearchResults
# Load environment variables
load_dotenv()
HUGGINGFACEHUB_API_TOKEN = os.getenv("HF_TOKEN")
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
# Set API keys
#os.environ["HUGGINGFACEHUB_API_TOKEN"] = HUGGINGFACEHUB_API_TOKEN
#os.environ["TAVILY_API_KEY"] = TAVILY_API_KEY
# Prompt Template
prompt_template = PromptTemplate(
input_variables=["context", "user_story"],
template="""You are a QA expert. Based on the context below and the given user story, write a detailed list of test cases.
Context:
{context}
User Story:
{user_story}
Test Cases:"""
)
# Load knowledge from RAG (local file)
def load_rag_knowledge():
with open("rag_knowledge_base.txt", "r", encoding="utf-8") as file:
content = file.read()
docs = [Document(page_content=content)]
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vector_store = FAISS.from_documents(docs, embeddings)
return vector_store.similarity_search("test case generation", k=1)
# Tavily search
def tavily_search(query):
search = TavilySearchResults(k=1)
results = search.run(query)
return results[0]['content'] if results else "No relevant results from Tavily."
# LLM call with combined context
def call_llm_with_context(context, user_story):
llm = HuggingFaceHub(
repo_id="mistralai/Mistral-7B-Instruct-v0.3",
model_kwargs={"temperature": 0.7, "max_new_tokens": 500}
)
chain = LLMChain(llm=llm, prompt=prompt_template)
return chain.run({"context": context, "user_story": user_story})
# Generate test cases pipeline
def generate_test_cases(user_story):
rag_docs = load_rag_knowledge()
rag_text = "\n".join([doc.page_content for doc in rag_docs])
tavily_text = tavily_search(user_story)
full_context = f"{rag_text}\n\n{tavily_text}"
test_cases = call_llm_with_context(full_context, user_story)
return rag_text.strip(), tavily_text.strip(), test_cases.strip()
# Gradio handler
def handle_generate(user_story):
rag, tavily, result = generate_test_cases(user_story)
return rag, tavily, result
# ----------------- Gradio UI -----------------
with gr.Blocks() as demo:
gr.Markdown("# 🧪 TechTales TestCaseGenerator using RAG + Tavily + Mistral + LangChain - Developed by Pankaj Kumar")
gr.Markdown("🚀 Enter a user story below to generate test cases using your knowledge base and Tavily search.")
user_input = gr.Textbox(label="📝 Enter User Story", lines=4, placeholder="As a user, I want to...")
btn = gr.Button("🔍 Generate Test Cases")
rag_output = gr.Textbox(label="📚 Knowledge from RAG File", lines=8)
tavily_output = gr.Textbox(label="🌐 Knowledge from Tavily Search", lines=8)
result_output = gr.Textbox(label="✅ Final Test Cases", lines=12)
btn.click(
handle_generate,
inputs=[user_input],
outputs=[rag_output, tavily_output, result_output]
)
demo.launch()
|