| 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_dotenv() |
| HUGGINGFACEHUB_API_TOKEN = os.getenv("HF_TOKEN") |
| TAVILY_API_KEY = os.getenv("TAVILY_API_KEY") |
|
|
| |
| |
| |
|
|
| |
| 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:""" |
| ) |
|
|
| |
| 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) |
|
|
| |
| def tavily_search(query): |
| search = TavilySearchResults(k=1) |
| results = search.run(query) |
| return results[0]['content'] if results else "No relevant results from Tavily." |
|
|
| |
| 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}) |
|
|
| |
| 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() |
|
|
| |
| def handle_generate(user_story): |
| rag, tavily, result = generate_test_cases(user_story) |
| return rag, tavily, result |
|
|
| |
| 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() |
|
|
