Update chat_engine.py

chat_engine.py

@@ -69,23 +69,36 @@ def extract_text_from_pdf(pdf_path):
 
 
 def get_api_type(api_type):
-    if api_type == 'openai':
-        # default is gpt-3.5-turbo, can also be gpt-4-0314
-        return OpenAI(model='gpt-4o', temperature=0.9) # for QA, temp is low
-    elif api_type == 'claude':
+    if model_name == 'openai-gpt-4o':
+        return OpenAI(model='gpt-4o', temperature=0.7)
+    elif model_name == 'openai-gpt-4-turbo':
+        return OpenAI(model='gpt-4-turbo', temperature=0.7'})
+    elif model_name == 'openai-gpt-3.5-turbo':
+        return OpenAI(model='gpt-3.5-turbo', temperature=0.7'})
+    elif model_name == 'claude-sonnet-3.5':
         return Anthropic(model="claude-3-5-sonnet-20240620")
-    elif api_type == 'llama':
-        return Anyscale(model='meta-llama/Llama-2-70b-chat-hf')
-    elif api_type == 'mistral':
-        return Anyscale(model='mistralai/Mixtral-8x7B-Instruct-v0.1', max_tokens=10000)
+    elif model_name == 'claude-opus-3':
+        return Anthropic(model="claude-3-opus-20240229")
+    elif model_name == 'claude-sonnet-3':
+        return Anthropic(model="claude-3-sonnet-20240229")
+    elif model_name == 'claude-haiku-3':
+        return Anthropic(model="claude-3-haiku-20240307")
+    elif model_name == 'llama-3-70B':
+        return Anyscale(model='meta-llama/Meta-Llama-3-70B-Instruct')
+    elif model_name == 'llama-3-8B':
+        return Anyscale(model='meta-llama/Meta-Llama-3-70B-Instruct')
+    elif model_name == 'mistral-8x7B':
+        return Anyscale(model='mistralai/Mixtral-8x7B-Instruct-v0.1')
+    elif model_name == 'mistral-8x22B':
+        return Anyscale(model='mistralai/Mixtral-8x7B-Instruct-v0.1')
     else:
         raise NotImplementedError
   
 
-def get_chat_engine(files, progress=gr.Progress()):
+def get_chat_engine(files, api_type, progress=gr.Progress()):
 
     progress(0, desc="Uploading Documents...")
-    llm = get_api_type('openai')
+    llm = get_api_type(api_type)
     Settings.llm = llm
     embed_model = MistralAIEmbedding(model_name='mistral-embed', api_key=mistral_api_key)
     Settings.embed_model = embed_model