From a2a3797bd681a0272d254cdffcb22c6d889c2efe Mon Sep 17 00:00:00 2001
From: Krish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 8 Aug 2023 11:40:28 -0700
Subject: [PATCH] add support for claude, azure, cohere, replicate

---
 datachad/backend/models.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/datachad/backend/models.py b/datachad/backend/models.py
index 2faa592..9bf62ea 100644
--- a/datachad/backend/models.py
+++ b/datachad/backend/models.py
@@ -3,6 +3,8 @@
 
 import streamlit as st
 import tiktoken
+import litellm
+from litellm import completion
 from langchain.base_language import BaseLanguageModel
 from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 from langchain.chat_models import ChatOpenAI
@@ -81,6 +83,13 @@ def get_model(options: dict, credentials: dict) -> BaseLanguageModel:
                 verbose=True,
                 callbacks=[StreamingStdOutCallbackHandler()],
             )
+        case *litellm.model_list:
+            model = ChatOpenAI(
+                model_name=options["model"].name,
+                temperature=options["temperature"],
+                openai_api_key=credentials["openai_api_key"],
+                client=completion
+            )
         # Added models need to be cased here
         case _default:
             msg = f"Model {options['model'].name} not supported!"