From 2bdcbf20fb9a08446ca9dc35741e107bdbcbe0c3 Mon Sep 17 00:00:00 2001 From: Willie Zutz Date: Sat, 15 Feb 2025 16:03:24 -0700 Subject: [PATCH 1/5] User customizable context window for ollama models. --- src/routes/images.ts | 6 +++ src/routes/search.ts | 6 +++ src/routes/suggestions.ts | 5 ++ src/routes/videos.ts | 6 +++ src/websocket/connectionManager.ts | 6 +++ ui/app/settings/page.tsx | 75 ++++++++++++++++++++++++++++++ ui/components/ChatWindow.tsx | 2 + ui/components/SearchImages.tsx | 5 +- ui/components/SearchVideos.tsx | 5 +- ui/lib/actions.ts | 4 ++ 10 files changed, 118 insertions(+), 2 deletions(-) diff --git a/src/routes/images.ts b/src/routes/images.ts index 5671657..2e8e912 100644 --- a/src/routes/images.ts +++ b/src/routes/images.ts @@ -5,6 +5,7 @@ import { getAvailableChatModelProviders } from '../lib/providers'; import { HumanMessage, AIMessage } from '@langchain/core/messages'; import logger from '../utils/logger'; import { ChatOpenAI } from '@langchain/openai'; +import { ChatOllama } from '@langchain/community/chat_models/ollama'; import { getCustomOpenaiApiKey, getCustomOpenaiApiUrl, @@ -16,6 +17,7 @@ const router = express.Router(); interface ChatModel { provider: string; model: string; + ollamaContextWindow?: number; } interface ImageSearchBody { @@ -61,6 +63,10 @@ router.post('/', async (req, res) => { ) { llm = chatModelProviders[chatModelProvider][chatModel] .model as unknown as BaseChatModel | undefined; + + if (llm instanceof ChatOllama) { + llm.numCtx = body.chatModel?.ollamaContextWindow || 2048; + } } if (!llm) { diff --git a/src/routes/search.ts b/src/routes/search.ts index 57d90a3..daefece 100644 --- a/src/routes/search.ts +++ b/src/routes/search.ts @@ -15,12 +15,14 @@ import { getCustomOpenaiApiUrl, getCustomOpenaiModelName, } from '../config'; +import { ChatOllama } from '@langchain/community/chat_models/ollama'; const router = express.Router(); interface chatModel { provider: string; model: string; + ollamaContextWindow?: number; customOpenAIKey?: string; customOpenAIBaseURL?: string; } @@ -78,6 +80,7 @@ router.post('/', async (req, res) => { const embeddingModel = body.embeddingModel?.model || Object.keys(embeddingModelProviders[embeddingModelProvider])[0]; + const ollamaContextWindow = body.chatModel?.ollamaContextWindow || 2048; let llm: BaseChatModel | undefined; let embeddings: Embeddings | undefined; @@ -99,6 +102,9 @@ router.post('/', async (req, res) => { ) { llm = chatModelProviders[chatModelProvider][chatModel] .model as unknown as BaseChatModel | undefined; + if (llm instanceof ChatOllama) { + llm.numCtx = ollamaContextWindow; + } } if ( diff --git a/src/routes/suggestions.ts b/src/routes/suggestions.ts index 7dd1739..c7a1409 100644 --- a/src/routes/suggestions.ts +++ b/src/routes/suggestions.ts @@ -10,12 +10,14 @@ import { getCustomOpenaiApiUrl, getCustomOpenaiModelName, } from '../config'; +import { ChatOllama } from '@langchain/community/chat_models/ollama'; const router = express.Router(); interface ChatModel { provider: string; model: string; + ollamaContextWindow?: number; } interface SuggestionsBody { @@ -60,6 +62,9 @@ router.post('/', async (req, res) => { ) { llm = chatModelProviders[chatModelProvider][chatModel] .model as unknown as BaseChatModel | undefined; + if (llm instanceof ChatOllama) { + llm.numCtx = body.chatModel?.ollamaContextWindow || 2048; + } } if (!llm) { diff --git a/src/routes/videos.ts b/src/routes/videos.ts index b631f26..debe3cd 100644 --- a/src/routes/videos.ts +++ b/src/routes/videos.ts @@ -10,12 +10,14 @@ import { getCustomOpenaiApiUrl, getCustomOpenaiModelName, } from '../config'; +import { ChatOllama } from '@langchain/community/chat_models/ollama'; const router = express.Router(); interface ChatModel { provider: string; model: string; + ollamaContextWindow?: number; } interface VideoSearchBody { @@ -61,6 +63,10 @@ router.post('/', async (req, res) => { ) { llm = chatModelProviders[chatModelProvider][chatModel] .model as unknown as BaseChatModel | undefined; + + if (llm instanceof ChatOllama) { + llm.numCtx = body.chatModel?.ollamaContextWindow || 2048; + } } if (!llm) { diff --git a/src/websocket/connectionManager.ts b/src/websocket/connectionManager.ts index bb8f242..979b8a0 100644 --- a/src/websocket/connectionManager.ts +++ b/src/websocket/connectionManager.ts @@ -14,6 +14,7 @@ import { getCustomOpenaiApiUrl, getCustomOpenaiModelName, } from '../config'; +import { ChatOllama } from '@langchain/community/chat_models/ollama'; export const handleConnection = async ( ws: WebSocket, @@ -42,6 +43,8 @@ export const handleConnection = async ( searchParams.get('embeddingModel') || Object.keys(embeddingModelProviders[embeddingModelProvider])[0]; + const ollamaContextWindow = searchParams.get('ollamaContextWindow'); + let llm: BaseChatModel | undefined; let embeddings: Embeddings | undefined; @@ -52,6 +55,9 @@ export const handleConnection = async ( ) { llm = chatModelProviders[chatModelProvider][chatModel] .model as unknown as BaseChatModel | undefined; + if (llm instanceof ChatOllama) { + llm.numCtx = ollamaContextWindow ? parseInt(ollamaContextWindow) : 2048; + } } else if (chatModelProvider == 'custom_openai') { const customOpenaiApiKey = getCustomOpenaiApiKey(); const customOpenaiApiUrl = getCustomOpenaiApiUrl(); diff --git a/ui/app/settings/page.tsx b/ui/app/settings/page.tsx index 371d091..37a2f6b 100644 --- a/ui/app/settings/page.tsx +++ b/ui/app/settings/page.tsx @@ -23,6 +23,7 @@ interface SettingsType { customOpenaiApiKey: string; customOpenaiApiUrl: string; customOpenaiModelName: string; + ollamaContextWindow: number; } interface InputProps extends React.InputHTMLAttributes { @@ -112,6 +113,11 @@ const Page = () => { const [automaticImageSearch, setAutomaticImageSearch] = useState(false); const [automaticVideoSearch, setAutomaticVideoSearch] = useState(false); const [savingStates, setSavingStates] = useState>({}); + const [contextWindowSize, setContextWindowSize] = useState(2048); + const [isCustomContextWindow, setIsCustomContextWindow] = useState(false); + const predefinedContextSizes = [ + 1024, 2048, 3072, 4096, 8192, 16384, 32768, 65536, 131072 + ]; useEffect(() => { const fetchConfig = async () => { @@ -123,6 +129,7 @@ const Page = () => { }); const data = (await res.json()) as SettingsType; + setConfig(data); const chatModelProvidersKeys = Object.keys(data.chatModelProviders || {}); @@ -171,6 +178,9 @@ const Page = () => { setAutomaticVideoSearch( localStorage.getItem('autoVideoSearch') === 'true', ); + const storedContextWindow = parseInt(localStorage.getItem('ollamaContextWindow') ?? '2048'); + setContextWindowSize(storedContextWindow); + setIsCustomContextWindow(!predefinedContextSizes.includes(storedContextWindow)); setIsLoading(false); }; @@ -331,6 +341,8 @@ const Page = () => { localStorage.setItem('embeddingModelProvider', value); } else if (key === 'embeddingModel') { localStorage.setItem('embeddingModel', value); + } else if (key === 'ollamaContextWindow') { + localStorage.setItem('ollamaContextWindow', value.toString()); } } catch (err) { console.error('Failed to save:', err); @@ -548,6 +560,69 @@ const Page = () => { ]; })()} /> + {selectedChatModelProvider === 'ollama' && ( +
+

+ Chat Context Window Size +

+ { + // Allow any value to be typed + const value = parseInt(e.target.value) || contextWindowSize; + setContextWindowSize(value); + }} + onSave={(value) => { + // Validate only when saving + const numValue = Math.max(512, parseInt(value) || 2048); + setContextWindowSize(numValue); + setConfig((prev) => ({ + ...prev!, + ollamaContextWindow: numValue, + })); + saveConfig('ollamaContextWindow', numValue); + }} + /> +
+ )} +

+ {isCustomContextWindow + ? "Adjust the context window size for Ollama models (minimum 512 tokens)" + : "Adjust the context window size for Ollama models"} +

+ + )} )} diff --git a/ui/components/ChatWindow.tsx b/ui/components/ChatWindow.tsx index 1940f42..366048f 100644 --- a/ui/components/ChatWindow.tsx +++ b/ui/components/ChatWindow.tsx @@ -197,6 +197,8 @@ const useSocket = ( 'openAIBaseURL', localStorage.getItem('openAIBaseURL')!, ); + } else { + searchParams.append('ollamaContextWindow', localStorage.getItem('ollamaContextWindow') || '2048') } searchParams.append('embeddingModel', embeddingModel!); diff --git a/ui/components/SearchImages.tsx b/ui/components/SearchImages.tsx index 383f780..21778b8 100644 --- a/ui/components/SearchImages.tsx +++ b/ui/components/SearchImages.tsx @@ -33,9 +33,9 @@ const SearchImages = ({ const chatModelProvider = localStorage.getItem('chatModelProvider'); const chatModel = localStorage.getItem('chatModel'); - const customOpenAIBaseURL = localStorage.getItem('openAIBaseURL'); const customOpenAIKey = localStorage.getItem('openAIApiKey'); + const ollamaContextWindow = localStorage.getItem('ollamaContextWindow') || '2048'; const res = await fetch( `${process.env.NEXT_PUBLIC_API_URL}/images`, @@ -54,6 +54,9 @@ const SearchImages = ({ customOpenAIBaseURL: customOpenAIBaseURL, customOpenAIKey: customOpenAIKey, }), + ...(chatModelProvider === 'ollama' && { + ollamaContextWindow: parseInt(ollamaContextWindow), + }), }, }), }, diff --git a/ui/components/SearchVideos.tsx b/ui/components/SearchVideos.tsx index c284dc2..4a67d47 100644 --- a/ui/components/SearchVideos.tsx +++ b/ui/components/SearchVideos.tsx @@ -48,9 +48,9 @@ const Searchvideos = ({ const chatModelProvider = localStorage.getItem('chatModelProvider'); const chatModel = localStorage.getItem('chatModel'); - const customOpenAIBaseURL = localStorage.getItem('openAIBaseURL'); const customOpenAIKey = localStorage.getItem('openAIApiKey'); + const ollamaContextWindow = localStorage.getItem('ollamaContextWindow') || '2048'; const res = await fetch( `${process.env.NEXT_PUBLIC_API_URL}/videos`, @@ -69,6 +69,9 @@ const Searchvideos = ({ customOpenAIBaseURL: customOpenAIBaseURL, customOpenAIKey: customOpenAIKey, }), + ...(chatModelProvider === 'ollama' && { + ollamaContextWindow: parseInt(ollamaContextWindow), + }), }, }), }, diff --git a/ui/lib/actions.ts b/ui/lib/actions.ts index a4409b0..4191733 100644 --- a/ui/lib/actions.ts +++ b/ui/lib/actions.ts @@ -6,6 +6,7 @@ export const getSuggestions = async (chatHisory: Message[]) => { const customOpenAIKey = localStorage.getItem('openAIApiKey'); const customOpenAIBaseURL = localStorage.getItem('openAIBaseURL'); + const ollamaContextWindow = localStorage.getItem('ollamaContextWindow') || '2048'; const res = await fetch(`${process.env.NEXT_PUBLIC_API_URL}/suggestions`, { method: 'POST', @@ -21,6 +22,9 @@ export const getSuggestions = async (chatHisory: Message[]) => { customOpenAIKey, customOpenAIBaseURL, }), + ...(chatModelProvider === 'ollama' && { + ollamaContextWindow: parseInt(ollamaContextWindow), + }), }, }), }); From 18b6f5b6749243a81338923823ce175be46e4758 Mon Sep 17 00:00:00 2001 From: Willie Zutz Date: Sat, 15 Feb 2025 16:07:19 -0700 Subject: [PATCH 2/5] Updated formatting --- ui/app/settings/page.tsx | 37 +++++++++++++++++++++++----------- ui/components/ChatWindow.tsx | 5 ++++- ui/components/SearchImages.tsx | 3 ++- ui/components/SearchVideos.tsx | 3 ++- ui/lib/actions.ts | 3 ++- 5 files changed, 35 insertions(+), 16 deletions(-) diff --git a/ui/app/settings/page.tsx b/ui/app/settings/page.tsx index 37a2f6b..26cdbd6 100644 --- a/ui/app/settings/page.tsx +++ b/ui/app/settings/page.tsx @@ -116,7 +116,7 @@ const Page = () => { const [contextWindowSize, setContextWindowSize] = useState(2048); const [isCustomContextWindow, setIsCustomContextWindow] = useState(false); const predefinedContextSizes = [ - 1024, 2048, 3072, 4096, 8192, 16384, 32768, 65536, 131072 + 1024, 2048, 3072, 4096, 8192, 16384, 32768, 65536, 131072, ]; useEffect(() => { @@ -178,9 +178,13 @@ const Page = () => { setAutomaticVideoSearch( localStorage.getItem('autoVideoSearch') === 'true', ); - const storedContextWindow = parseInt(localStorage.getItem('ollamaContextWindow') ?? '2048'); + const storedContextWindow = parseInt( + localStorage.getItem('ollamaContextWindow') ?? '2048', + ); setContextWindowSize(storedContextWindow); - setIsCustomContextWindow(!predefinedContextSizes.includes(storedContextWindow)); + setIsCustomContextWindow( + !predefinedContextSizes.includes(storedContextWindow), + ); setIsLoading(false); }; @@ -566,7 +570,11 @@ const Page = () => { Chat Context Window Size

handleCompactChange(e.target.checked)} + className="form-checkbox h-4 w-4 text-blue-600 transition duration-150 ease-in-out" + /> +
+ +
+

+ Compact Mode +

+

+ Generate more concise responses +

+
+
+ + From b84e4e4ce6f1dfcd76d4d18d898a9deec55fe1d7 Mon Sep 17 00:00:00 2001 From: Willie Zutz Date: Sun, 16 Feb 2025 15:08:30 -0700 Subject: [PATCH 4/5] Added an icon to indicate that compact mode is enabled. --- ui/components/MessageInputActions/Optimization.tsx | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/ui/components/MessageInputActions/Optimization.tsx b/ui/components/MessageInputActions/Optimization.tsx index 649c555..9f0fef7 100644 --- a/ui/components/MessageInputActions/Optimization.tsx +++ b/ui/components/MessageInputActions/Optimization.tsx @@ -70,6 +70,12 @@ const Optimization = ({ className="p-2 text-black/50 dark:text-white/50 rounded-xl hover:bg-light-secondary dark:hover:bg-dark-secondary active:scale-95 transition duration-200 hover:text-black dark:hover:text-white" >
+ {isCompact && ( + + )} { OptimizationModes.find((mode) => mode.key === optimizationMode) ?.icon From 690ef428613058fec0669de184c8bb51a64726b1 Mon Sep 17 00:00:00 2001 From: Willie Zutz Date: Mon, 17 Feb 2025 01:22:34 -0700 Subject: [PATCH 5/5] Fixes a bug with rewriting where history wouldn't get removed. --- ui/components/ChatWindow.tsx | 35 +++++++++++++++++------------------ 1 file changed, 17 insertions(+), 18 deletions(-) diff --git a/ui/components/ChatWindow.tsx b/ui/components/ChatWindow.tsx index 6462761..f87292d 100644 --- a/ui/components/ChatWindow.tsx +++ b/ui/components/ChatWindow.tsx @@ -472,10 +472,10 @@ const ChatWindow = ({ id }: { id?: string }) => { } }, [isMessagesLoaded, isWSReady]); - const sendMessage = async ( +const sendMessage = async ( message: string, messageId?: string, - options?: { isCompact?: boolean }, + options?: { isCompact?: boolean; rewriteIndex?: number }, ) => { if (loading) return; if (!ws || ws.readyState !== WebSocket.OPEN) { @@ -489,6 +489,17 @@ const ChatWindow = ({ id }: { id?: string }) => { let sources: Document[] | undefined = undefined; let recievedMessage = ''; let added = false; + let messageChatHistory = chatHistory; + + if (options?.rewriteIndex !== undefined) { + const rewriteIndex = options.rewriteIndex; + setMessages((prev) => { + return [...prev.slice(0, messages.length > 2 ? rewriteIndex - 1 : 0)] + }); + + messageChatHistory = chatHistory.slice(0, messages.length > 2 ? rewriteIndex - 1 : 0) + setChatHistory(messageChatHistory); + } messageId = messageId ?? crypto.randomBytes(7).toString('hex'); let messageData = { @@ -501,10 +512,9 @@ const ChatWindow = ({ id }: { id?: string }) => { files: fileIds, focusMode: focusMode, optimizationMode: optimizationMode, - history: [...chatHistory, ['human', message]], + history: [...messageChatHistory, ['human', message]], isCompact: options?.isCompact ?? isCompact, }; - //console.log('sending:', messageData, JSON.stringify(messageData)); ws.send(JSON.stringify(messageData)); setMessages((prevMessages) => [ @@ -622,20 +632,9 @@ const ChatWindow = ({ id }: { id?: string }) => { }; const rewrite = (messageId: string) => { - const index = messages.findIndex((msg) => msg.messageId === messageId); - - if (index === -1) return; - - const message = messages[index - 1]; - - setMessages((prev) => { - return [...prev.slice(0, messages.length > 2 ? index - 1 : 0)]; - }); - setChatHistory((prev) => { - return [...prev.slice(0, messages.length > 2 ? index - 1 : 0)]; - }); - - sendMessage(message.content, message.messageId, { isCompact }); + const messageIndex = messages.findIndex((msg) => msg.messageId === messageId); + if(messageIndex == -1) return; + sendMessage(messages[messageIndex - 1].content, messageId, { isCompact, rewriteIndex: messageIndex }); }; useEffect(() => {