From 6edac6938c0f63d13d9b5531626c46230f4bb026 Mon Sep 17 00:00:00 2001 From: haddadrm <121486289+haddadrm@users.noreply.github.com> Date: Sun, 26 Jan 2025 18:18:35 +0400 Subject: [PATCH] feat: Add LM Studio Support and Thinking Model Panel LM Studio Integration: - Added LM Studio provider with OpenAI-compatible API support - Dynamic model discovery via /v1/models endpoint - Support for both chat and embeddings models - Docker-compatible networking configuration Thinking Model Panel: - Added collapsible UI panel for model's chain of thought - Parses responses with tags to separate reasoning - Maintains backward compatibility with regular responses - Styled consistently with app theme for light/dark modes - Preserves all existing message functionality (sources, markdown, etc.) These improvements enhance the app's compatibility with local LLMs and provide better visibility into model reasoning processes while maintaining existing functionality. --- src/config.ts | 28 ++++++++++- src/lib/providers/index.ts | 3 ++ src/lib/providers/lmstudio.ts | 89 +++++++++++++++++++++++++++++++++++ src/routes/config.ts | 3 ++ ui/components/MessageBox.tsx | 86 ++++++++++++++++++++++++++------- 5 files changed, 191 insertions(+), 18 deletions(-) create mode 100644 src/lib/providers/lmstudio.ts diff --git a/src/config.ts b/src/config.ts index 001c259..c0556b3 100644 --- a/src/config.ts +++ b/src/config.ts @@ -16,9 +16,10 @@ interface Config { ANTHROPIC: string; GEMINI: string; }; - API_ENDPOINTS: { - SEARXNG: string; + API_ENDPOINTS: { OLLAMA: string; + LMSTUDIO: string; + SEARXNG: string; }; } @@ -51,6 +52,8 @@ export const getSearxngApiEndpoint = () => export const getOllamaApiEndpoint = () => loadConfig().API_ENDPOINTS.OLLAMA; +export const getLMStudioApiEndpoint = () => loadConfig().API_ENDPOINTS.LMSTUDIO; + export const updateConfig = (config: RecursivePartial) => { const currentConfig = loadConfig(); @@ -72,6 +75,27 @@ export const updateConfig = (config: RecursivePartial) => { } } + /* +export const updateConfig = (config: RecursivePartial) => { + const currentConfig = loadConfig(); + + // Merge existing config with new values + const mergedConfig: RecursivePartial = { + GENERAL: { + ...currentConfig.GENERAL, + ...config.GENERAL, + }, + API_KEYS: { + ...currentConfig.API_KEYS, + ...config.API_KEYS, + }, + API_ENDPOINTS: { + ...currentConfig.API_ENDPOINTS, + ...config.API_ENDPOINTS, + }, + }; +*/ + fs.writeFileSync( path.join(__dirname, `../${configFileName}`), toml.stringify(config), diff --git a/src/lib/providers/index.ts b/src/lib/providers/index.ts index 98846e7..1da17a7 100644 --- a/src/lib/providers/index.ts +++ b/src/lib/providers/index.ts @@ -4,6 +4,7 @@ import { loadOpenAIChatModels, loadOpenAIEmbeddingsModels } from './openai'; import { loadAnthropicChatModels } from './anthropic'; import { loadTransformersEmbeddingsModels } from './transformers'; import { loadGeminiChatModels, loadGeminiEmbeddingsModels } from './gemini'; +import { loadLMStudioChatModels, loadLMStudioEmbeddingsModels } from './lmstudio'; const chatModelProviders = { openai: loadOpenAIChatModels, @@ -11,6 +12,7 @@ const chatModelProviders = { ollama: loadOllamaChatModels, anthropic: loadAnthropicChatModels, gemini: loadGeminiChatModels, + lm_studio: loadLMStudioChatModels, }; const embeddingModelProviders = { @@ -18,6 +20,7 @@ const embeddingModelProviders = { local: loadTransformersEmbeddingsModels, ollama: loadOllamaEmbeddingsModels, gemini: loadGeminiEmbeddingsModels, + lm_studio: loadLMStudioEmbeddingsModels, }; export const getAvailableChatModelProviders = async () => { diff --git a/src/lib/providers/lmstudio.ts b/src/lib/providers/lmstudio.ts new file mode 100644 index 0000000..623a658 --- /dev/null +++ b/src/lib/providers/lmstudio.ts @@ -0,0 +1,89 @@ +import { OpenAIEmbeddings } from '@langchain/openai'; +import { ChatOpenAI } from '@langchain/openai'; +import { getKeepAlive, getLMStudioApiEndpoint } from '../../config'; +import logger from '../../utils/logger'; +import axios from 'axios'; + +interface LMStudioModel { + id: string; + // add other properties if LM Studio API provides them +} + +interface ChatModelConfig { + displayName: string; + model: ChatOpenAI; +} + +export const loadLMStudioChatModels = async (): Promise> => { + const lmStudioEndpoint = getLMStudioApiEndpoint(); + + if (!lmStudioEndpoint) { + logger.debug('LM Studio endpoint not configured, skipping'); + return {}; + } + + try { + const response = await axios.get<{ data: LMStudioModel[] }>(`${lmStudioEndpoint}/models`, { + headers: { + 'Content-Type': 'application/json', + }, + }); + + const lmStudioModels = response.data.data; + + const chatModels = lmStudioModels.reduce>((acc, model) => { + acc[model.id] = { + displayName: model.id, + model: new ChatOpenAI({ + openAIApiKey: 'lm-studio', + configuration: { + baseURL: lmStudioEndpoint, + }, + modelName: model.id, + temperature: 0.7, + }), + }; + return acc; + }, {}); + + return chatModels; + } catch (err) { + logger.error(`Error loading LM Studio models: ${err}`); + return {}; + } +}; + +export const loadLMStudioEmbeddingsModels = async () => { + const lmStudioEndpoint = getLMStudioApiEndpoint(); + + if (!lmStudioEndpoint) return {}; + + try { + const response = await axios.get(`${lmStudioEndpoint}/models`, { + headers: { + 'Content-Type': 'application/json', + }, + }); + + const lmStudioModels = response.data.data; + + const embeddingsModels = lmStudioModels.reduce((acc, model) => { + acc[model.id] = { + displayName: model.id, + model: new OpenAIEmbeddings({ + openAIApiKey: 'lm-studio', // Dummy key required by LangChain + configuration: { + baseURL: lmStudioEndpoint, + }, + modelName: model.id, + }), + }; + return acc; + }, {}); + + return embeddingsModels; + } catch (err) { + logger.error(`Error loading LM Studio embeddings model: ${err}`); + return {}; + } +}; \ No newline at end of file diff --git a/src/routes/config.ts b/src/routes/config.ts index 6ff80c6..afb19c6 100644 --- a/src/routes/config.ts +++ b/src/routes/config.ts @@ -6,6 +6,7 @@ import { import { getGroqApiKey, getOllamaApiEndpoint, + getLMStudioApiEndpoint, getAnthropicApiKey, getGeminiApiKey, getOpenaiApiKey, @@ -51,6 +52,7 @@ router.get('/', async (_, res) => { config['openaiApiKey'] = getOpenaiApiKey(); config['ollamaApiUrl'] = getOllamaApiEndpoint(); + config['lmStudioApiUrl'] = getLMStudioApiEndpoint(); config['anthropicApiKey'] = getAnthropicApiKey(); config['groqApiKey'] = getGroqApiKey(); config['geminiApiKey'] = getGeminiApiKey(); @@ -74,6 +76,7 @@ router.post('/', async (req, res) => { }, API_ENDPOINTS: { OLLAMA: config.ollamaApiUrl, + LMSTUDIO: config.lmStudioApiUrl, }, }; diff --git a/ui/components/MessageBox.tsx b/ui/components/MessageBox.tsx index f23127c..282b2ff 100644 --- a/ui/components/MessageBox.tsx +++ b/ui/components/MessageBox.tsx @@ -11,6 +11,8 @@ import { StopCircle, Layers3, Plus, + Brain, + ChevronDown, } from 'lucide-react'; import Markdown from 'markdown-to-jsx'; import Copy from './MessageActions/Copy'; @@ -41,26 +43,48 @@ const MessageBox = ({ }) => { const [parsedMessage, setParsedMessage] = useState(message.content); const [speechMessage, setSpeechMessage] = useState(message.content); + const [thinking, setThinking] = useState(''); + const [answer, setAnswer] = useState(''); + const [isThinkingExpanded, setIsThinkingExpanded] = useState(false); useEffect(() => { const regex = /\[(\d+)\]/g; - if ( - message.role === 'assistant' && - message?.sources && - message.sources.length > 0 - ) { - return setParsedMessage( - message.content.replace( - regex, - (_, number) => - `${number}`, - ), - ); - } + // First check for thinking content + const match = message.content.match(/(.*?)<\/think>(.*)/s); + if (match) { + const [_, thinkingContent, answerContent] = match; + setThinking(thinkingContent.trim()); + setAnswer(answerContent.trim()); - setSpeechMessage(message.content.replace(regex, '')); - setParsedMessage(message.content); + // Process the answer part for sources if needed + if (message.role === 'assistant' && message?.sources && message.sources.length > 0) { + setParsedMessage( + answerContent.trim().replace( + regex, + (_, number) => + `${number}`, + ), + ); + } else { + setParsedMessage(answerContent.trim()); + } + setSpeechMessage(answerContent.trim().replace(regex, '')); + } else { + // No thinking content - process as before + if (message.role === 'assistant' && message?.sources && message.sources.length > 0) { + setParsedMessage( + message.content.replace( + regex, + (_, number) => + `${number}`, + ), + ); + } else { + setParsedMessage(message.content); + } + setSpeechMessage(message.content.replace(regex, '')); + } }, [message.content, message.sources, message.role]); const { speechStatus, start, stop } = useSpeech({ text: speechMessage }); @@ -81,6 +105,37 @@ const MessageBox = ({ ref={dividerRef} className="flex flex-col space-y-6 w-full lg:w-9/12" > + {thinking && ( +
+ + + {isThinkingExpanded && ( +
+ + {thinking} + +
+ )} +
+ )} {message.sources && message.sources.length > 0 && (
@@ -199,4 +254,3 @@ const MessageBox = ({ ); }; -export default MessageBox;