Update app.py

app.py

@@ -1,191 +1,328 @@
-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 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 nest_asyncio
+# nest_asyncio.apply()
 
-# import ollama
-# from llama_index.llms.ollama import Ollama
-# from llama_index.core.llms import ChatMessage
+# # 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
+# # 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
+# # 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
+# # Mistral
+# key = os.getenv('MISTRAL_API_KEY')
+# os.environ["MISTRAL_API_KEY"] = key
 
 
 
-# Streamlit UI
-st.title("Auto Test Case Generation using LLM")
+# # 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)
+# 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...")
+# 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)
+#         # 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}"
+#         st.success("Fetching list of functions...")
+#         file_path = f"./data/{uploaded_file.name}"
         
-        def extract_functions_from_file(file_path, file_extension):
+#         def extract_functions_from_file(file_path, file_extension):
 
-            if file_extension == '.py':
-                with open(file_path, "r") as file:
-                    file_content = file.read()
+#             if file_extension == '.py':
+#                 with open(file_path, "r") as file:
+#                     file_content = file.read()
                 
-                parsed_content = ast.parse(file_content)
-                methods = {}
+#                 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
+#                 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()
+#             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
+#                 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*\{')
+#                 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 = []
+#                 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'])
+#                 if 'main' in methods.keys():
+#                     del(methods['main'])
 
-            return methods
+#             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
+#         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)
+        print(file_extension)
+        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_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 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.
+        """
+        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 = []
 
-         # 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())
+            test_cases = self.generate_test_cases(func)
+            st.session_state.messages.append({"role": "assistant", "content": f"{test_cases}"})
+            st.markdown(test_cases)
 
-            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'])
-                
+if __name__ == "__main__":
+    test_case_generator = TestCaseGenerator()
+    test_case_generator.run()