Support custom endpoints and models for LLMChat (#88)

Author: mahaloz

dailalib/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "3.17.0"
+__version__ = "3.18.0"
 
 import os
 # stop LiteLLM from querying at all to the remote server

dailalib/api/litellm/litellm_api.py

@@ -13,6 +13,7 @@
 
 active_model = None
 active_prompt_style = None
+active_custom_endpoint = None
 
 _l = logging.getLogger(__name__)
 
@@ -64,9 +65,10 @@ def __init__(
         self.prompts_by_name = {p.name: p for p in prompts}
 
         # update the globals (for threading hacks)
-        global active_model, active_prompt_style
-        active_model = self.model
+        global active_model, active_prompt_style, active_custom_endpoint
+        active_model = self.model if not self.custom_model else self.custom_model
         active_prompt_style = self.prompt_style
+        active_custom_endpoint = self.custom_endpoint
 
     def load_or_create_config(self, new_config=None) -> bool:
         if new_config:
@@ -195,7 +197,6 @@ def query_model(
         # get the answer
         try:
             answer = response.choices[0].message.content
-            if self.custom_endpoint: print(answer)
         except (KeyError, IndexError) as e:
             answer = None
 
@@ -261,10 +262,19 @@ def _set_model(self, model):
         global active_model
         active_model = model
 
+    def _set_custom_endpoint(self, custom_endpoint):
+        self.custom_endpoint = custom_endpoint
+        global active_custom_endpoint
+        active_custom_endpoint = custom_endpoint
+
     def get_model(self):
         # TODO: this hack needs to be refactored later
-        global active_model
-        return str(active_model)
+        global active_model, active_custom_endpoint
+        return str(active_model) if not active_custom_endpoint else str(active_custom_endpoint)
+
+    def get_custom_endpoint(self):
+        global active_custom_endpoint
+        return active_custom_endpoint
 
     #
     # LLM Settings

dailalib/llm_chat/llm_chat_ui.py

@@ -28,6 +28,7 @@ class LLMChatClient(QWidget):
     def __init__(self, ai_api: "LiteLLMAIAPI", parent=None, context: Context = None):
         super(LLMChatClient, self).__init__(parent)
         self.model = ai_api.get_model()
+        self.custom_endpoint = ai_api.get_custom_endpoint()
         self.ai_api = ai_api
         self.context = context
         self.setWindowTitle('LLM Chat')
@@ -65,6 +66,7 @@ def __init__(self, ai_api: "LiteLLMAIAPI", parent=None, context: Context = None)
 
         # Chat history
         self.chat_history = []
+        self.thread = None
 
         # model check
         if not self.model:
@@ -152,7 +154,9 @@ def send_message(self, add_text=True, role="user"):
         self.send_button.setDisabled(True)
 
         # Start a thread to get the response
-        self.thread = LLMThread(self.chat_history, self.model)
+        self.thread = LLMThread(
+            self.chat_history, self.model, custom_endpoint=self.custom_endpoint, api_key=self.ai_api.api_key
+        )
         self.thread.response_received.connect(lambda msg: self.receive_message(msg))
         self.thread.start()
 
@@ -183,10 +187,12 @@ def closeEvent(self, event):
 class LLMThread(QThread):
     response_received = pyqtSignal(str)
 
-    def __init__(self, chat_history, model_name):
+    def __init__(self, chat_history, model_name, custom_endpoint=None, api_key=None):
         super().__init__()
         self.chat_history = chat_history.copy()
         self.model_name = model_name
+        self.custom_endpoint = custom_endpoint
+        self.api_key = api_key
 
     def run(self):
         import litellm
@@ -196,8 +202,9 @@ def run(self):
         response = litellm.completion(
             model=self.model_name,
             messages=self.chat_history,
-            timeout=60,
-
+            timeout=60 if not self.custom_endpoint else 300,
+            api_base=self.custom_endpoint if self.custom_endpoint else None,  # Use custom endpoint if set
+            api_key=self.api_key if not self.custom_endpoint else "dummy" # In most of cases custom endpoint doesn't need the api_key
         )
         litellm.modify_params = False