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'])