# import os # import re # import streamlit as st # import ast # import json # import openai # from llama_index.llms.openai import OpenAI # from llama_index.core.llms import ChatMessage # from llama_index.llms.anthropic import Anthropic # from llama_index.llms.mistralai import MistralAI # import nest_asyncio # nest_asyncio.apply() # # import ollama # # from llama_index.llms.ollama import Ollama # # from llama_index.core.llms import ChatMessage # # OpenAI credentials # # key = os.getenv('OPENAI_API_KEY') # # openai.api_key = key # # os.environ["OPENAI_API_KEY"] = key # # Anthropic credentials # # key = os.getenv('CLAUDE_API_KEY') # # os.environ["ANTHROPIC_API_KEY"] = key # # Mistral # key = os.getenv('MISTRAL_API_KEY') # os.environ["MISTRAL_API_KEY"] = key # # Streamlit UI # st.title("Auto Test Case Generation using LLM") # uploaded_files = st.file_uploader("Upload a python or Java file", type=[".py","java"], accept_multiple_files=True) # if uploaded_files: # for uploaded_file in uploaded_files: # with open(f"./data/{uploaded_file.name}", 'wb') as f: # f.write(uploaded_file.getbuffer()) # st.success("File uploaded...") # # Check file type # _, file_extension = os.path.splitext(uploaded_file.name) # print(file_extension) # st.success("Fetching list of functions...") # file_path = f"./data/{uploaded_file.name}" # def extract_functions_from_file(file_path, file_extension): # if file_extension == '.py': # with open(file_path, "r") as file: # file_content = file.read() # parsed_content = ast.parse(file_content) # methods = {} # for node in ast.walk(parsed_content): # if isinstance(node, ast.FunctionDef): # func_name = node.name # func_body = ast.get_source_segment(file_content, node) # methods[func_name] = func_body # elif file_extension == '.java': # with open(file_path, 'r') as file: # lines = file.readlines() # methods = {} # inside_method = False # method_name = None # method_body = [] # brace_count = 0 # method_signature_pattern = re.compile(r'((?:public|protected|private|static|\s)*)\s+[\w<>\[\]]+\s+(\w+)\s*\([^)]*\)\s*\{') # for line in lines: # if not inside_method: # match = method_signature_pattern.search(line) # if match: # modifiers, method_name = match.groups() # inside_method = True # method_body.append(line) # brace_count = line.count('{') - line.count('}') # else: # method_body.append(line) # brace_count += line.count('{') - line.count('}') # if brace_count == 0: # inside_method = False # methods[method_name] = ''.join(method_body) # method_body = [] # if 'main' in methods.keys(): # del(methods['main']) # return methods # functions = extract_functions_from_file(file_path, file_extension) # list_of_functions = list(functions.keys()) # st.write(list_of_functions) # def res(prompt, model=None): # # response = openai.chat.completions.create( # # model=model, # # messages=[ # # {"role": "user", # # "content": prompt, # # } # # ] # # ) # # return response.choices[0].message.content # response = [ # ChatMessage(role="system", content="You are a sincere and helpful coding assistant"), # ChatMessage(role="user", content=prompt), # ] # # resp = Anthropic(model=model).chat(response) # resp = MistralAI(model).chat(response) # return resp # # Initialize session state for chat messages # if "messages" not in st.session_state: # st.session_state.messages = [] # # Display chat messages from history on app rerun # for message in st.session_state.messages: # with st.chat_message(message["role"]): # st.markdown(message["content"]) # # Accept user input # if func := st.chat_input("Enter the function name for generating test cases:"): # st.session_state.messages.append({"role": "assistant", "content": f"Generating test cases for {func}"}) # st.success(f"Generating test cases for {func}") # func = ''.join(func.split()) # if func not in list_of_functions: # st.write("Incorrect function name") # else: # snippet = functions[func] # # Generation # # model = "gpt-3.5-turbo" # # model = "claude-3-haiku-20240307" # # model = "claude-3-sonnet-20240229" # # model = "claude-3-opus-20240229" # model = "codestral-latest" # # Generation # # resp = ollama.generate(model='codellama', # # prompt=f""" Your task is to generate unit test cases for this function : {snippet}\ # # \n\n Politely refuse if the function is not suitable for generating test cases. # # \n\n Generate atleast 5 unit test case. Include couple of edge cases as well. # # \n\n There should be no duplicate test cases. # # \n\n Avoid generating repeated statements. # # """) # prompt=f""" Your task is to generate unit test cases for this function : \n\n{snippet}\ # \n\n Generate between 3 to 8 unique unit test cases. Include couple of edge cases as well. # \n\n All the test cases should have the mandatory assert statement. # \n\n Every test case should be defined as a method inside the class. # \n\n All the test cases should have textual description. # \n\n Politely refuse if the function is not suitable for generating test cases. # \n\n There should be no duplicate and incomplete test case. # \n\n Avoid generating repeated statements. # \n\n Recheck your response before generating. # \n\n Do not share the last Test Case. # """ # # print(prompt) # resp = res(prompt = prompt, model = model) # # Post Processing # post_prompt = f"""Except the last test case, display everything that is present in this end to end: \n\n{resp}\ # \n\n Do not add anything extra. Just copy and paste everything except the last test case. # \n\n Do not mention the count of total number of test cases in the response. # \n\n Do not mention this sentence - "I have excluded the last test case as per your request" # """ # post_resp = res(prompt = post_prompt, model = model) # st.session_state.messages.append({"role": "assistant", "content": f"{post_resp}"}) # st.markdown(post_resp) # # st.session_state.messages.append({"role": "assistant", "content": f"{resp['response']}"}) # # st.markdown(resp['response']) import os import re import ast import streamlit as st from llama_index.llms.openai import OpenAI from llama_index.core.llms import ChatMessage from llama_index.llms.anthropic import Anthropic from llama_index.llms.mistralai import MistralAI import nest_asyncio class TestCaseGenerator: def __init__(self): nest_asyncio.apply() self.key = os.getenv('MISTRAL_API_KEY') os.environ["MISTRAL_API_KEY"] = self.key self.model = "codestral-latest" self.functions = {} self.list_of_functions = [] def setup_streamlit_ui(self): st.title("Auto Test Case Generation using LLM") uploaded_files = st.file_uploader("Upload a python or Java file", type=["py", "java"], accept_multiple_files=True) if uploaded_files: for uploaded_file in uploaded_files: self.process_uploaded_file(uploaded_file) def process_uploaded_file(self, uploaded_file): with open(f"./data/{uploaded_file.name}", 'wb') as f: f.write(uploaded_file.getbuffer()) st.success("File uploaded...") _, file_extension = os.path.splitext(uploaded_file.name) st.success("Fetching list of functions...") file_path = f"./data/{uploaded_file.name}" self.extract_functions_from_file(file_path, file_extension) st.write(self.list_of_functions) def extract_functions_from_file(self, file_path, file_extension): if file_extension == '.py': self.extract_python_functions(file_path) elif file_extension == '.java': self.extract_java_functions(file_path) if 'main' in self.functions.keys(): del(self.functions['main']) self.list_of_functions = list(self.functions.keys()) def extract_python_functions(self, file_path): with open(file_path, "r") as file: file_content = file.read() parsed_content = ast.parse(file_content) for node in ast.walk(parsed_content): if isinstance(node, ast.FunctionDef): func_name = node.name func_body = ast.get_source_segment(file_content, node) self.functions[func_name] = func_body def extract_java_functions(self, file_path): with open(file_path, 'r') as file: lines = file.readlines() inside_method = False method_name = None method_body = [] brace_count = 0 method_signature_pattern = re.compile(r'((?:public|protected|private|static|\s)*)\s+[\w<>\[\]]+\s+(\w+)\s*\([^)]*\)\s*\{') for line in lines: if not inside_method: match = method_signature_pattern.search(line) if match: modifiers, method_name = match.groups() inside_method = True method_body.append(line) brace_count = line.count('{') - line.count('}') else: method_body.append(line) brace_count += line.count('{') - line.count('}') if brace_count == 0: inside_method = False self.functions[method_name] = ''.join(method_body) method_body = [] # def generate_response(self, prompt): # response = [ # ChatMessage(role="system", content="You are a sincere and helpful coding assistant"), # ChatMessage(role="user", content=prompt), # ] # resp = MistralAI(self.model).chat(response) # return resp def generate_response(self, prompt): response = [ ChatMessage(role="system", content="You are tasked with generating unit test cases, including descriptions and assert statements, for a given function. Follow these instructions carefully:"), ChatMessage(role="user", content=prompt), ] resp = MistralAI(self.model).chat(response) return resp def generate_test_cases(self, func): if func not in self.list_of_functions: st.write("Incorrect function name") return snippet = self.functions[func] prompt=f""" Your task is to generate unit test cases for this function : \n\n{snippet}\ \n\n Generate between 3 to 8 unique unit test cases. Include couple of edge cases as well. \n\n All the test cases should have the mandatory assert statement. \n\n All the test cases should have textual description. \n\n Politely refuse if the function is not suitable for generating test cases. \n\n There should be no duplicate and incomplete test case. \n\n Avoid generating repeated statements. \n\n Recheck your response before generating. \n\n Do not share the last Test Case. """ resp = self.generate_response(prompt) post_prompt = f"""Except the last test case, display everything that is present in this end to end: \n\n{resp}\ \n\n - Do not add anything extra. Just copy and paste everything except the last test case. \n\n - Do not mention the count of total number of test cases in the response. \n\n - Do not mention this sentence - "I have excluded the last test case as per your request" """ post_resp = self.generate_response(post_prompt) return post_resp def run(self): self.setup_streamlit_ui() if "messages" not in st.session_state: st.session_state.messages = [] for message in st.session_state.messages: with st.chat_message(message["role"]): st.markdown(message["content"]) if func := st.chat_input("Enter the function name for generating test cases:"): st.session_state.messages.append({"role": "assistant", "content": f"Generating test cases for {func}"}) st.success(f"Generating test cases for {func}") func = ''.join(func.split()) test_cases = self.generate_test_cases(func) st.session_state.messages.append({"role": "assistant", "content": f"{test_cases}"}) st.markdown(test_cases) if __name__ == "__main__": test_case_generator = TestCaseGenerator() test_case_generator.run()