Adding LM Studio Support

sample.config.toml

@@ -21,8 +21,11 @@ NETLIFY = "<YOUR_NETLIFY_API_KEY>"
 BING = "https://api.bing.microsoft.com/v7.0/search"
 GOOGLE = "https://www.googleapis.com/customsearch/v1"
 OLLAMA = "http://127.0.0.1:11434"
+
+LM_STUDIO = "http://localhost:1234/v1"
 OPENAI = "https://api.openai.com/v1"
 
+
 [LOGGING]
 LOG_REST_API = "true"
 LOG_PROMPTS = "false"

src/config.py

@@ -59,6 +59,9 @@ def get_google_search_api_endpoint(self):
 
     def get_ollama_api_endpoint(self):
         return self.config["API_ENDPOINTS"]["OLLAMA"]
+
+    def get_lmstudio_api_endpoint(self):
+        return self.config["API_ENDPOINTS"]["LM_STUDIO"]
 
     def get_claude_api_key(self):
         return self.config["API_KEYS"]["CLAUDE"]
@@ -131,6 +134,10 @@ def set_google_search_api_endpoint(self, endpoint):
     def set_ollama_api_endpoint(self, endpoint):
         self.config["API_ENDPOINTS"]["OLLAMA"] = endpoint
         self.save_config()
+
+    def set_lmstudio_api_endpoint(self, endpoint):
+        self.config["API_ENDPOINTS"]["LM_STUDIO"] = endpoint
+        self.save_config()
 
     def set_claude_api_key(self, key):
         self.config["API_KEYS"]["CLAUDE"] = key

src/llm/llm.py

@@ -10,6 +10,7 @@
 from .gemini_client import Gemini
 from .mistral_client import MistralAi
 from .groq_client import Groq
+from .lm_studio_client import LMStudio
 
 from src.state import AgentState
 
@@ -57,7 +58,11 @@ def __init__(self, model_id: str = None):
                 ("Mixtral", "mixtral-8x7b-32768"),
                 ("GEMMA 7B", "gemma-7b-it"),
             ],
-            "OLLAMA": []
+            "OLLAMA": [],
+            "LM_STUDIO": [
+                ("LM Studio", "local-model"),    
+            ],
+
         }
         if ollama.client:
             self.models["OLLAMA"] = [(model["name"], model["name"]) for model in ollama.models]
@@ -96,7 +101,8 @@ def inference(self, prompt: str, project_name: str) -> str:
             "OPENAI": OpenAi(),
             "GOOGLE": Gemini(),
             "MISTRAL": MistralAi(),
-            "GROQ": Groq()
+            "GROQ": Groq(),
+            "LM_STUDIO": LMStudio()
         }
 
         try:

src/llm/lm_studio_client.py

@@ -0,0 +1,30 @@
+from src.logger import Logger
+from src.config import Config
+from openai import OpenAI
+
+
+log = Logger()
+
+class LMStudio:
+    def __init__(self):
+        try:
+            self.api_endpoint = Config().get_lmstudio_api_endpoint()
+            self.client = OpenAI(base_url=self.api_endpoint, api_key="not-needed")
+            log.info("LM Studio available")
+        except:
+            self.api_endpoint = None
+            self.client = None
+            log.warning("LM Studio not available")
+            log.warning("Make sure to set the LM Studio API endpoint in the config")
+
+    def inference(self, model_id: str, prompt: str) -> str:
+        chat_completion = self.client.chat.completions.create(
+            messages=[
+                {
+                    "role": "user",
+                    "content": prompt.strip(),
+                }
+            ],
+            model=model_id, # unused 
+        )
+        return chat_completion.choices[0].message.content