Update app.py

app.py

@@ -4,6 +4,8 @@ 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
 import nest_asyncio
 
 nest_asyncio.apply()
@@ -14,9 +16,10 @@ nest_asyncio.apply()
 
 
 # OpenAI credentials
-key = os.getenv('OPENAI_API_KEY')
-openai.api_key = key 
-os.environ["OPENAI_API_KEY"] = key
+# key = os.getenv('OPENAI_API_KEY')
+key = os.getenv('CLAUDE_API_KEY')
+# openai.api_key = key 
+# os.environ["OPENAI_API_KEY"] = key
 
 # Streamlit UI
 st.title("Auto Test Case Generation using LLM")
@@ -50,18 +53,27 @@ if uploaded_files:
         list_of_functions = list(functions.keys())
         st.write(list_of_functions)
 
-        def res(prompt):
+        def res(prompt, model):
 
-            response = openai.chat.completions.create(
-                model=model,
-                messages=[
-                    {"role": "user",
-                     "content": prompt,
-                    }
-                ]
-            )
+            # response = openai.chat.completions.create(
+            #     model=model,
+            #     messages=[
+            #         {"role": "user",
+            #          "content": prompt,
+            #         }
+            #     ]
+            # )
 
-            return response.choices[0].message.content
+            # return response.choices[0].message.content
+
+            response = [
+                ChatMessage(
+                    role="system", content="You are a helpful assistant"
+                ),
+                ChatMessage(role="user", content=prompt),
+            ]
+            resp = Anthropic(model=model).chat(messages)
+            return resp
 
         # Initialize session state for chat messages
         if "messages" not in st.session_state:
@@ -86,7 +98,8 @@ if uploaded_files:
                 snippet = functions[func]
 
                 # Generation
-                model = "gpt-3.5-turbo"
+                # model = "gpt-3.5-turbo"
+                model = "claude-3-haiku-20240307"
 
                 # Generation
                 # resp = ollama.generate(model='codellama', 
@@ -102,9 +115,9 @@ if uploaded_files:
                        \n\nThere should be no duplicate test cases. Avoid generating repeated statements.
                     """
 
-                print(prompt)
+                # print(prompt)
               
-                resp = res(prompt)
+                resp = res(prompt, model)
                 st.session_state.messages.append({"role": "assistant", "content": f"{resp}"})
                 st.markdown(resp)
                 # st.session_state.messages.append({"role": "assistant", "content": f"{resp['response']}"})