Update app.py

app.py

@@ -1,11 +1,101 @@
 import streamlit as st
 import ast
 import json
-import ollama
-from llama_index.llms.ollama import Ollama
 from llama_index.core.llms import ChatMessage
+from transformers import AutoTokenizer
+import transformers
+import torch
+
+# import ollama
+# from llama_index.llms.ollama import Ollama
+# from llama_index.core.llms import ChatMessage
 
 # Streamlit UI
 st.title("Auto Test Case Generation using LLM")
 
-uploaded_files = st.file_uploader("Upload a python(.py) file", type=".py", accept_multiple_files=True)
+uploaded_files = st.file_uploader("Upload a python(.py) file", type=".py", 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...")
+
+        st.success("Fetching list of functions...")
+        file_path = f"./data/{uploaded_file.name}"
+        def extract_functions_from_file(file_path):
+            with open(file_path, "r") as file:
+                file_content = file.read()
+            
+            parsed_content = ast.parse(file_content)
+            functions = {}
+            
+            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)
+                    functions[func_name] = func_body
+            
+            return functions
+        
+        functions = extract_functions_from_file(file_path)
+
+        list_of_functions = list(functions.keys())
+        st.write(list_of_functions)
+
+        # 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]
+
+                model = "codellama/CodeLlama-7b-Instruct-hf"
+                
+                tokenizer = AutoTokenizer.from_pretrained(model)
+                pipeline = transformers.pipeline(
+                    model=model,
+                    torch_dtype=torch.float16,
+                    device_map="auto",
+                )
+
+                # Generation
+                # resp = ollama.generate(model='codellama', 
+                #                     prompt=f"""You are a helpful coding assistant. Your task is to generate unit test cases for this function : {snippet}\
+                #                     \n\nPolitely refuse if the function is not suitable for generating test cases.
+                #                     \n\nGenerate atleast 5 unit test case. Include couple of edge cases as well.
+                #                     \n\nThere should be no duplicate test cases. Avoid generating repeated statements.
+                #                     """)
+                resp = pipeline(
+                    f"""You are a helpful coding assistant. Your task is to generate unit test cases for this function : {snippet}\
+                    \n\nPolitely refuse if the function is not suitable for generating test cases.
+                    \n\nGenerate atleast 5 unit test case. Include couple of edge cases as well.
+                    \n\nThere should be no duplicate test cases. Avoid generating repeated statements.
+                    """,
+                    do_sample=True,
+                    top_k=10,
+                    temperature=0.1,
+                    top_p=0.95,
+                    num_return_sequences=1,
+                    eos_token_id=tokenizer.eos_token_id,
+                )
+                resp_list = [n['generated_text'] for n in resp]
+                response = " ".join(resp_list)
+                st.session_state.messages.append({"role": "assistant", "content": f"{resp['response']}"})
+                st.markdown(resp['response'])
+