Update app.py

app.py

@@ -1,14 +1,22 @@
 import streamlit as st
 import ast
 import json
-from transformers import AutoTokenizer
-import transformers
-import torch
+import openai
+from llama_index.llms.openai import OpenAI
+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
+
 # Streamlit UI
 st.title("Auto Test Case Generation using LLM")
 
@@ -38,10 +46,25 @@ if uploaded_files:
             return functions
         
         functions = extract_functions_from_file(file_path)
-
         list_of_functions = list(functions.keys())
         st.write(list_of_functions)
 
+        def res(prompt):
+
+            response = openai.chat.completions.create(
+                model=model,
+                messages=[
+                    {"role":"system",
+                     "content":"You are a helpful coding assistant. Your task is to generate test cases. If the function can't be found, politely refuse"
+                    },
+                    {"role": "user",
+                     "content": prompt,
+                    }
+                ]
+            )
+
+            return response.choices[0].message.content
+
         # Initialize session state for chat messages
         if "messages" not in st.session_state:
             st.session_state.messages = []
@@ -64,18 +87,8 @@ if uploaded_files:
             else:
                 snippet = functions[func]
 
-                # model = "codellama/CodeLlama-7b-Instruct-hf"
-                model = "TinyLlama/TinyLlama_v1.1"
-                
-                tokenizer = AutoTokenizer.from_pretrained(model)
-                pipeline = transformers.pipeline(
-                   "text-generation",
-                    model=model,
-                    torch_dtype=torch.float16,
-                    device_map="auto",
-                )
-
-                print('Done')
+                # Generation
+                model = "gpt-3.5-turbo"
 
                 # Generation
                 # resp = ollama.generate(model='codellama', 
@@ -84,17 +97,18 @@ if uploaded_files:
                 #                     \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.
-                    """,
-                    max_new_tokens=1024
-                )
-                print('Response ', resp)
-                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'])
+
+                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.
+                    """
+
+                print(prompt)
+              
+                resp = res(prompt)
+                st.session_state.messages.append({"role": "assistant", "content": f"{resp}"})
+                st.markdown(resp)
+                # st.session_state.messages.append({"role": "assistant", "content": f"{resp['response']}"})
+                # st.markdown(resp['response'])