replacing openai compeltion with litellm

requirements.txt

@@ -1,5 +1,5 @@
 streamlit==1.37.1
-openai==1.40.6
+litellm==1.47.0
 discord-webhook==1.3.1
 python-dotenv==1.0.1
 colorama==0.4.6

src/utils/constants.py

@@ -1,27 +1,28 @@
-"""Module doc string"""
-
-from .config import OPENAI_API_KEY
-
-
-class ConstantVariables:
-    """Module doc string"""
-
-    model_list_tuple = (
-        "gpt-4o",
-        "gpt-4o-mini",
-        "gpt-4-turbo",
-        "gpt-4",
-        "gpt-3.5-turbo",
-    )
-    default_model = "gpt-4o-mini"
-
-    max_tokens = 180
-    min_token = 20
-    step = 20
-    default = round(((max_tokens + min_token) / 2) / step) * step
-    default_token = max(min_token, min(max_tokens, default))
-
-    if OPENAI_API_KEY != "NO_KEY":
-        api_key = OPENAI_API_KEY
-    else:
-        api_key = ""
+"""Module doc string"""
+
+from .config import OPENAI_API_KEY
+
+
+class ConstantVariables:
+    """Module doc string"""
+
+    model_list_tuple = (
+        "gpt-4o",
+        "gpt-4o-mini",
+        "gpt-4-turbo",
+        "gpt-3.5-turbo",
+        "o1-preview",
+        "o1-mini"
+    )
+    default_model = "gpt-4o-mini"
+
+    max_tokens = 180
+    min_token = 20
+    step = 20
+    default = round(((max_tokens + min_token) / 2) / step) * step
+    default_token = max(min_token, min(max_tokens, default))
+
+    if OPENAI_API_KEY != "NO_KEY":
+        api_key = OPENAI_API_KEY
+    else:
+        api_key = ""

src/utils/openai_utils.py

@@ -1,70 +1,73 @@
-"""Module doc string"""
-
-import openai
-import streamlit as st
-from openai import OpenAI
-
-from .logs import log_execution_time, logger
-
-
-class OpenAIFunctions:
-    """Module doc string"""
-
-    @log_execution_time
-    @staticmethod
-    def invoke_model():
-        """_summary_"""
-        logger.debug("OpenAI invoked")
-        client = OpenAI(api_key=st.session_state.openai_api_key)
-        with st.chat_message("assistant"):
-            stream = client.chat.completions.create(
-                model=st.session_state["openai_model"],
-                messages=[
-                    {"role": m["role"], "content": m["content"]}
-                    for m in st.session_state.messages
-                ],
-                max_tokens=st.session_state["openai_maxtokens"],
-                stream=True,
-                stream_options={"include_usage": True},
-            )
-
-            def stream_data():
-                for chunk in stream:
-                    if chunk.choices != []:
-                        word = chunk.choices[0].delta.content
-                        if word is not None:
-                            yield word
-                    if chunk.usage is not None:
-                        yield {
-                            "completion_tokens": chunk.usage.completion_tokens,
-                            "prompt_tokens": chunk.usage.prompt_tokens,
-                            "total_tokens": chunk.usage.total_tokens,
-                        }
-
-            return st.write_stream(stream_data)
-
-    @log_execution_time
-    @staticmethod
-    def check_openai_api_key():
-        """_summary_"""
-        logger.info("Checking OpenAI Key")
-        try:
-            client = OpenAI(api_key=st.session_state.openai_api_key)
-            client.models.list()
-            logger.debug("OpenAI key Working")
-            return True
-        except openai.AuthenticationError as auth_error:
-            with st.chat_message("assistant"):
-                st.error(str(auth_error))
-            logger.error("AuthenticationError: %s", auth_error)
-            return False
-        except openai.OpenAIError as openai_error:
-            with st.chat_message("assistant"):
-                st.error(str(openai_error))
-            logger.error("OpenAIError: %s", openai_error)
-            return False
-        except Exception as general_error:
-            with st.chat_message("assistant"):
-                st.error(str(general_error))
-            logger.error("Unexpected error: %s", general_error)
-            return False
+"""Module doc string"""
+
+import openai
+import streamlit as st
+from litellm import completion
+from openai import OpenAI
+
+from .logs import log_execution_time, logger
+
+
+class OpenAIFunctions:
+    """Module doc string"""
+
+    @log_execution_time
+    @staticmethod
+    def invoke_model():
+        """_summary_"""
+        logger.debug("OpenAI invoked")
+        with st.chat_message("assistant"):
+            messages = [
+                {"role": m["role"], "content": m["content"]}
+                for m in st.session_state.messages
+            ]
+
+            stream = completion(
+                api_key=st.session_state.openai_api_key,
+                model=st.session_state["openai_model"],
+                messages=messages,
+                max_tokens=st.session_state["openai_maxtokens"],
+                stream=True,
+                stream_options={"include_usage": True},
+            )
+
+            def stream_data():
+                for chunk in stream:
+                    if chunk.choices != []:
+                        word = chunk.choices[0].delta.content
+                        if word is not None:
+                            yield word
+                    if hasattr(chunk, "usage"):
+                        yield {
+                            "completion_tokens": chunk.usage.completion_tokens,
+                            "prompt_tokens": chunk.usage.prompt_tokens,
+                            "total_tokens": chunk.usage.total_tokens,
+                        }
+
+            return st.write_stream(stream_data)
+
+    @log_execution_time
+    @staticmethod
+    def check_openai_api_key():
+        """_summary_"""
+        logger.info("Checking OpenAI Key")
+        try:
+            client = OpenAI(api_key=st.session_state.openai_api_key)
+            client.models.list()
+            logger.debug("OpenAI key Working")
+            return True
+        except openai.AuthenticationError as auth_error:
+            with st.chat_message("assistant"):
+                st.error(str(auth_error))
+            logger.error("AuthenticationError: %s", auth_error)
+            return False
+        except openai.OpenAIError as openai_error:
+            with st.chat_message("assistant"):
+                st.error(str(openai_error))
+            logger.error("OpenAIError: %s", openai_error)
+            return False
+        except Exception as general_error:
+            with st.chat_message("assistant"):
+                st.error(str(general_error))
+            logger.error("Unexpected error: %s", general_error)
+            return False