adding support ti run lm-studio

.gitignore

@@ -158,3 +158,5 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+/notebooks

src/app.py

@@ -24,7 +24,7 @@ def main():
         if (
             st.session_state.openai_api_key is not None
             and st.session_state.openai_api_key != ""
-        ):
+        ) or st.session_state.provider_select != "OpenAI":
             logger.info("OpenAI key Checking condition passed")
             if OpenAIFunctions.check_openai_api_key():
                 logger.info("Inference Started")

src/utils/config.py

@@ -1,12 +1,12 @@
-"""Module doc string"""
-
-import os
-
-from dotenv import find_dotenv, load_dotenv
-
-load_dotenv(find_dotenv(), override=True)
-
-LOGGER_LEVEL = os.getenv("LOGGER_LEVEL", "INFO")
-DISCORD_HOOK = os.getenv("DISCORD_HOOK", "NO_HOOK")
-ENVIRONMENT = os.getenv("ENVIRONMENT", "NOT_LOCAL")
-OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "NO_KEY")
+"""Module doc string"""
+
+import os
+
+from dotenv import find_dotenv, load_dotenv
+
+load_dotenv(find_dotenv(), override=True)
+
+LOGGER_LEVEL = os.getenv("LOGGER_LEVEL", "INFO")
+DISCORD_HOOK = os.getenv("DISCORD_HOOK", "NO_HOOK")
+ENVIRONMENT = os.getenv("ENVIRONMENT", "NOT_LOCAL")
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "NO_KEY")

src/utils/constants.py

@@ -12,13 +12,14 @@ class ConstantVariables:
         "gpt-4-turbo",
         "gpt-3.5-turbo",
         "o1-preview",
-        "o1-mini"
+        "o1-mini",
     )
+    provider = ("lm-studio", "OpenAI")
     default_model = "gpt-4o-mini"
-
-    max_tokens = 180
-    min_token = 20
-    step = 20
+    default_provider = "lm-studio"
+    max_tokens = 1024
+    min_token = 32
+    step = 32
     default = round(((max_tokens + min_token) / 2) / step) * step
     default_token = max(min_token, min(max_tokens, default))
 

src/utils/logs.py

@@ -1,98 +1,106 @@
-"""Module doc string"""
-
-import asyncio
-import logging
-import sys
-import time
-from functools import wraps
-
-from colorama import Back, Fore, Style, init
-
-from .config import LOGGER_LEVEL
-
-# Initialize colorama
-init(autoreset=True)
-
-logger = logging.getLogger(__name__)
-
-if not logger.hasHandlers():
-    logger.propagate = False
-    logger.setLevel(LOGGER_LEVEL)
-
-    # Define color codes for different log levels
-    log_colors = {
-        logging.DEBUG: Fore.CYAN,
-        logging.INFO: Fore.GREEN,
-        logging.WARNING: Fore.YELLOW,
-        logging.ERROR: Fore.RED,
-        logging.CRITICAL: Fore.RED + Back.WHITE + Style.BRIGHT,
-    }
-
-    class ColoredFormatter(logging.Formatter):
-        """Module doc string"""
-
-        def format(self, record):
-            """Module doc string"""
-
-            levelno = record.levelno
-            color = log_colors.get(levelno, "")
-
-            # Format the message
-            message = record.getMessage()
-
-            # Format the rest of the log details
-            details = self._fmt % {
-                "asctime": self.formatTime(record, self.datefmt),
-                "levelname": record.levelname,
-                "module": record.module,
-                "funcName": record.funcName,
-                "lineno": record.lineno,
-            }
-
-            # Combine details and colored message
-            return f"{Fore.WHITE}{details} :: {color}{message}{Style.RESET_ALL}"
-
-    normal_handler = logging.StreamHandler(sys.stdout)
-    normal_handler.setLevel(logging.DEBUG)
-    normal_handler.addFilter(lambda logRecord: logRecord.levelno < logging.WARNING)
-
-    error_handler = logging.StreamHandler(sys.stderr)
-    error_handler.setLevel(logging.WARNING)
-
-    formatter = ColoredFormatter(
-        "%(asctime)s :: %(levelname)s :: %(module)s :: %(funcName)s :: %(lineno)d"
-    )
-
-    normal_handler.setFormatter(formatter)
-    error_handler.setFormatter(formatter)
-
-    logger.addHandler(normal_handler)
-    logger.addHandler(error_handler)
-
-
-def log_execution_time(func):
-    """Module doc string"""
-
-    @wraps(func)
-    def sync_wrapper(*args, **kwargs):
-        start_time = time.time()
-        result = func(*args, **kwargs)
-        end_time = time.time()
-        execution_time = end_time - start_time
-        message_string = f"{func.__name__} executed in {execution_time:.4f} seconds"
-        logger.debug(message_string)
-        return result
-
-    @wraps(func)
-    async def async_wrapper(*args, **kwargs):
-        start_time = time.time()
-        result = await func(*args, **kwargs)
-        end_time = time.time()
-        execution_time = end_time - start_time
-        message_string = f"{func.__name__} executed in {execution_time:.4f} seconds"
-        logger.debug(message_string)
-        return result
-
-    if asyncio.iscoroutinefunction(func):
-        return async_wrapper
-    return sync_wrapper
+"""Module doc string"""
+
+import asyncio
+import logging
+import sys
+import time
+from functools import wraps
+
+from colorama import Back, Fore, Style, init
+
+from .config import LOGGER_LEVEL
+
+# Initialize colorama
+init(autoreset=True)
+
+logger = logging.getLogger(__name__)
+
+if not logger.hasHandlers():
+    logger.propagate = False
+    logger.setLevel(LOGGER_LEVEL)
+
+    # Define color codes for different log levels
+    log_colors = {
+        logging.DEBUG: Fore.CYAN,
+        logging.INFO: Fore.GREEN,
+        logging.WARNING: Fore.YELLOW,
+        logging.ERROR: Fore.RED,
+        logging.CRITICAL: Fore.RED + Back.WHITE + Style.BRIGHT,
+    }
+
+    class ColoredFormatter(logging.Formatter):
+        """Module doc string"""
+
+        def format(self, record):
+            """Module doc string"""
+
+            levelno = record.levelno
+            color = log_colors.get(levelno, "")
+
+            # Format the message
+            message = record.getMessage()
+
+            # Format the rest of the log details
+            details = self._fmt % {
+                "asctime": self.formatTime(record, self.datefmt),
+                "levelname": record.levelname,
+                "module": record.module,
+                "funcName": record.funcName,
+                "lineno": record.lineno,
+            }
+
+            # Combine details and colored message
+            return (
+                f"{Fore.WHITE}{details} :: {color}{message}{Style.RESET_ALL}"
+            )
+
+    normal_handler = logging.StreamHandler(sys.stdout)
+    normal_handler.setLevel(logging.DEBUG)
+    normal_handler.addFilter(
+        lambda logRecord: logRecord.levelno < logging.WARNING
+    )
+
+    error_handler = logging.StreamHandler(sys.stderr)
+    error_handler.setLevel(logging.WARNING)
+
+    formatter = ColoredFormatter(
+        "%(asctime)s :: %(levelname)s :: %(module)s :: %(funcName)s :: %(lineno)d"
+    )
+
+    normal_handler.setFormatter(formatter)
+    error_handler.setFormatter(formatter)
+
+    logger.addHandler(normal_handler)
+    logger.addHandler(error_handler)
+
+
+def log_execution_time(func):
+    """Module doc string"""
+
+    @wraps(func)
+    def sync_wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        execution_time = end_time - start_time
+        message_string = (
+            f"{func.__name__} executed in {execution_time:.4f} seconds"
+        )
+        logger.debug(message_string)
+        return result
+
+    @wraps(func)
+    async def async_wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = await func(*args, **kwargs)
+        end_time = time.time()
+        execution_time = end_time - start_time
+        message_string = (
+            f"{func.__name__} executed in {execution_time:.4f} seconds"
+        )
+        logger.debug(message_string)
+        return result
+
+    if asyncio.iscoroutinefunction(func):
+        return async_wrapper
+    return sync_wrapper

src/utils/openai_utils.py

@@ -15,21 +15,27 @@ class OpenAIFunctions:
     @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
             ]
+            comp_args = {}
+            if st.session_state.provider_select == "OpenAI":
+                comp_args["api_key"] = st.session_state.openai_api_key
+                comp_args["model"] = st.session_state["openai_model"]
+            elif st.session_state.provider_select == "lm-studio":
+                comp_args["base_url"] = "http://localhost:1234/v1"
+                comp_args["api_key"] = st.session_state.provider_select
+                comp_args["model"] = "gpt-4o-mini"
+            comp_args["messages"] = messages
+            comp_args["max_tokens"] = st.session_state["openai_maxtokens"]
+            comp_args["stream"] = True
+            comp_args["stream_options"] = {"include_usage": True}
 
-            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},
-            )
+            stream = completion(**comp_args)
 
             def stream_data():
                 for chunk in stream:
@@ -50,24 +56,28 @@ class OpenAIFunctions:
     @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")
+        if st.session_state.provider_select == "lm-studio":
+            logger.info("Local Provider is Sekected")
             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
+        else:
+            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

src/utils/streamlit_utils.py

@@ -25,36 +25,53 @@ class StreamlitFunctions:
     def streamlit_side_bar():
         """_summary_"""
         with st.sidebar:
-            st.text_input(
-                label="OpenAI API key",
-                value=ConstantVariables.api_key,
-                help="This will not be saved or stored.",
-                type="password",
-                key="api_key",
-            )
-
             st.selectbox(
-                "Select the GPT model",
-                ConstantVariables.model_list_tuple,
-                key="openai_model",
-            )
-            st.slider(
-                "Max Tokens",
-                min_value=ConstantVariables.min_token,
-                max_value=ConstantVariables.max_tokens,
-                step=ConstantVariables.step,
-                key="openai_maxtokens",
-            )
-            st.button(
-                "Start Chat",
-                on_click=StreamlitFunctions.start_app,
-                use_container_width=True,
-            )
-            st.button(
-                "Reset History",
-                on_click=StreamlitFunctions.reset_history,
-                use_container_width=True,
+                "Select Provider",
+                ConstantVariables.provider,
+                placeholder="Choose an option",
+                key="provider_select",
             )
+            if st.session_state.provider_select is not None:
+                if st.session_state.provider_select == "OpenAI":
+                    st.text_input(
+                        label="OpenAI API key",
+                        value=ConstantVariables.api_key,
+                        help="This will not be saved or stored.",
+                        type="password",
+                        key="api_key",
+                    )
+
+                    st.selectbox(
+                        "Select the GPT model",
+                        ConstantVariables.model_list_tuple,
+                        key="openai_model",
+                    )
+
+                elif st.session_state.provider_select == "lm-studio":
+                    st.header("NOTE")
+                    st.text(
+                        "lm-studio is configured to work on `http://localhost:1234/v1`"
+                    )
+
+                st.slider(
+                    "Max Tokens",
+                    min_value=ConstantVariables.min_token,
+                    max_value=ConstantVariables.max_tokens,
+                    step=ConstantVariables.step,
+                    key="openai_maxtokens",
+                )
+
+                st.button(
+                    "Start Chat",
+                    on_click=StreamlitFunctions.start_app,
+                    use_container_width=True,
+                )
+
+                st.button(
+                    "Reset History",
+                    on_click=StreamlitFunctions.reset_history,
+                    use_container_width=True,
+                )
 
     @staticmethod
     def streamlit_initialize_variables():
@@ -66,15 +83,23 @@ class StreamlitFunctions:
         if "openai_model" not in st.session_state:
             st.session_state["openai_model"] = ConstantVariables.default_model
 
+        if "provider_select" not in st.session_state:
+            st.session_state["provider_select"] = None
+
         if "openai_api_key" not in st.session_state:
             st.session_state["openai_api_key"] = None
 
         if "openai_maxtokens" not in st.session_state:
-            st.session_state["openai_maxtokens"] = ConstantVariables.default_token
+            st.session_state["openai_maxtokens"] = (
+                ConstantVariables.default_token
+            )
 
         if "start_app" not in st.session_state:
             st.session_state["start_app"] = False
 
+        if "api_key" not in st.session_state:
+            st.session_state["api_key"] = None
+
     @staticmethod
     def reset_history():
         """_summary_"""
@@ -102,7 +127,11 @@ class StreamlitFunctions:
         if prompt := st.chat_input("Type your Query"):
             with st.chat_message("user"):
                 st.markdown(prompt)
-            st.session_state.messages.append({"role": "user", "content": prompt})
+            st.session_state.messages.append(
+                {"role": "user", "content": prompt}
+            )
             response = OpenAIFunctions.invoke_model()
             logger.debug(response)
-            st.session_state.messages.append({"role": "assistant", "content": response[0]})
+            st.session_state.messages.append(
+                {"role": "assistant", "content": response[0]}
+            )