Update app.py

app.py

@@ -6,6 +6,7 @@ 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()
@@ -21,8 +22,13 @@ nest_asyncio.apply()
 # os.environ["OPENAI_API_KEY"] = key
 
 # Anthropic credentials
-key = os.getenv('CLAUDE_API_KEY')
-os.environ["ANTHROPIC_API_KEY"] = key
+# key = os.getenv('CLAUDE_API_KEY')
+# os.environ["ANTHROPIC_API_KEY"] = key
+
+# Mistral
+key = os.getenv('CODESTRAL_KEY')
+os.environ["MISTRAL_API_KEY"] = key
+
 
 # Streamlit UI
 st.title("Auto Test Case Generation using LLM")
@@ -56,7 +62,7 @@ if uploaded_files:
         list_of_functions = list(functions.keys())
         st.write(list_of_functions)
 
-        def res(prompt, model):
+        def res(prompt, model=None):
 
             # response = openai.chat.completions.create(
             #     model=model,
@@ -73,7 +79,8 @@ if uploaded_files:
                 ChatMessage(role="system", content="You are a sincere and helpful coding assistant"),
                 ChatMessage(role="user", content=prompt),
             ]
-            resp = Anthropic(model=model).chat(response)
+            # resp = Anthropic(model=model).chat(response)
+            resp = MistralAI().chat(messages)
             return resp
 
         # Initialize session state for chat messages
@@ -100,10 +107,9 @@ if uploaded_files:
 
                 # Generation
                 # model = "gpt-3.5-turbo"
-                # claude-3-sonnet-20240229
                 # model = "claude-3-haiku-20240307"
                 # model = "claude-3-sonnet-20240229"
-                model = "claude-3-opus-20240229"
+                # model = "claude-3-opus-20240229"
                 
 
                 # Generation
@@ -127,7 +133,7 @@ if uploaded_files:
 
                 # print(prompt)
               
-                resp = res(prompt, model)
+                resp = res(prompt, model = None)
                 st.session_state.messages.append({"role": "assistant", "content": f"{resp}"})
                 st.markdown(resp)
                 # st.session_state.messages.append({"role": "assistant", "content": f"{resp['response']}"})