feat: add keep alive strategy for ollama api

This commit is contained in:
mok-liee 2024-10-04 15:14:17 +02:00
parent 1680a1786e
commit 6d6c7122b9
5 changed files with 67 additions and 2 deletions

View file

@ -9,4 +9,7 @@ ANTHROPIC = "" # Anthropic API key - sk-ant-1234567890abcdef1234567890abcdef
[API_ENDPOINTS] [API_ENDPOINTS]
SEARXNG = "http://localhost:32768" # SearxNG API URL SEARXNG = "http://localhost:32768" # SearxNG API URL
OLLAMA = "" # Ollama API URL - http://host.docker.internal:11434 OLLAMA = "" # Ollama API URL - http://host.docker.internal:11434
[OLLAMA]
KEEP_ALIVE = "5m"

View file

@ -14,6 +14,9 @@ interface Config {
GROQ: string; GROQ: string;
ANTHROPIC: string; ANTHROPIC: string;
}; };
OLLAMA: {
KEEP_ALIVE: string;
};
API_ENDPOINTS: { API_ENDPOINTS: {
SEARXNG: string; SEARXNG: string;
OLLAMA: string; OLLAMA: string;
@ -45,6 +48,8 @@ export const getSearxngApiEndpoint = () =>
export const getOllamaApiEndpoint = () => loadConfig().API_ENDPOINTS.OLLAMA; export const getOllamaApiEndpoint = () => loadConfig().API_ENDPOINTS.OLLAMA;
export const getOllamaKeepAliveStrategy = () => loadConfig().OLLAMA.KEEP_ALIVE;
export const updateConfig = (config: RecursivePartial<Config>) => { export const updateConfig = (config: RecursivePartial<Config>) => {
const currentConfig = loadConfig(); const currentConfig = loadConfig();

View file

@ -1,5 +1,5 @@
import { OllamaEmbeddings } from '@langchain/community/embeddings/ollama'; import { OllamaEmbeddings } from '@langchain/community/embeddings/ollama';
import { getOllamaApiEndpoint } from '../../config'; import { getOllamaApiEndpoint, getOllamaKeepAliveStrategy } from '../../config';
import logger from '../../utils/logger'; import logger from '../../utils/logger';
import { ChatOllama } from '@langchain/community/chat_models/ollama'; import { ChatOllama } from '@langchain/community/chat_models/ollama';
@ -24,6 +24,7 @@ export const loadOllamaChatModels = async () => {
baseUrl: ollamaEndpoint, baseUrl: ollamaEndpoint,
model: model.model, model: model.model,
temperature: 0.7, temperature: 0.7,
keepAlive: getOllamaKeepAliveStrategy(),
}), }),
}; };
@ -57,6 +58,7 @@ export const loadOllamaEmbeddingsModels = async () => {
model: new OllamaEmbeddings({ model: new OllamaEmbeddings({
baseUrl: ollamaEndpoint, baseUrl: ollamaEndpoint,
model: model.model, model: model.model,
keepAlive: getOllamaKeepAliveStrategy(),
}), }),
}; };

View file

@ -9,6 +9,7 @@ import {
getAnthropicApiKey, getAnthropicApiKey,
getOpenaiApiKey, getOpenaiApiKey,
updateConfig, updateConfig,
getOllamaKeepAliveStrategy,
} from '../config'; } from '../config';
import logger from '../utils/logger'; import logger from '../utils/logger';
@ -50,6 +51,7 @@ router.get('/', async (_, res) => {
config['openaiApiKey'] = getOpenaiApiKey(); config['openaiApiKey'] = getOpenaiApiKey();
config['ollamaApiUrl'] = getOllamaApiEndpoint(); config['ollamaApiUrl'] = getOllamaApiEndpoint();
config['ollamaKeepAliveStrategy'] = getOllamaKeepAliveStrategy();
config['anthropicApiKey'] = getAnthropicApiKey(); config['anthropicApiKey'] = getAnthropicApiKey();
config['groqApiKey'] = getGroqApiKey(); config['groqApiKey'] = getGroqApiKey();
@ -72,6 +74,9 @@ router.post('/', async (req, res) => {
API_ENDPOINTS: { API_ENDPOINTS: {
OLLAMA: config.ollamaApiUrl, OLLAMA: config.ollamaApiUrl,
}, },
OLLAMA: {
KEEP_ALIVE: config.ollamaKeepAliveStrategy,
},
}; };
updateConfig(updatedConfig); updateConfig(updatedConfig);

View file

@ -54,6 +54,7 @@ interface SettingsType {
embeddingModelProviders: { embeddingModelProviders: {
[key: string]: [Record<string, any>]; [key: string]: [Record<string, any>];
}; };
ollamaKeepAliveStrategy: string;
openaiApiKey: string; openaiApiKey: string;
groqApiKey: string; groqApiKey: string;
anthropicApiKey: string; anthropicApiKey: string;
@ -78,6 +79,10 @@ const SettingsDialog = ({
const [selectedChatModel, setSelectedChatModel] = useState<string | null>( const [selectedChatModel, setSelectedChatModel] = useState<string | null>(
null, null,
); );
const [selectedOllamaKeepAliveStrategy, setSelectedOllamaKeepAliveStrategy] =
useState<string | null>(null);
const [selectedEmbeddingModelProvider, setSelectedEmbeddingModelProvider] = const [selectedEmbeddingModelProvider, setSelectedEmbeddingModelProvider] =
useState<string | null>(null); useState<string | null>(null);
const [selectedEmbeddingModel, setSelectedEmbeddingModel] = useState< const [selectedEmbeddingModel, setSelectedEmbeddingModel] = useState<
@ -124,6 +129,10 @@ const SettingsDialog = ({
(data.chatModelProviders && (data.chatModelProviders &&
data.chatModelProviders[chatModelProvider]?.[0].name) || data.chatModelProviders[chatModelProvider]?.[0].name) ||
''; '';
const ollamaKeepAliveStrategy =
localStorage.getItem('ollamaKeepAliveStrategy') ||
data.ollamaKeepAliveStrategy ||
'';
const embeddingModelProvider = const embeddingModelProvider =
localStorage.getItem('embeddingModelProvider') || localStorage.getItem('embeddingModelProvider') ||
defaultEmbeddingModelProvider || defaultEmbeddingModelProvider ||
@ -136,6 +145,7 @@ const SettingsDialog = ({
setSelectedChatModelProvider(chatModelProvider); setSelectedChatModelProvider(chatModelProvider);
setSelectedChatModel(chatModel); setSelectedChatModel(chatModel);
setSelectedOllamaKeepAliveStrategy(ollamaKeepAliveStrategy);
setSelectedEmbeddingModelProvider(embeddingModelProvider); setSelectedEmbeddingModelProvider(embeddingModelProvider);
setSelectedEmbeddingModel(embeddingModel); setSelectedEmbeddingModel(embeddingModel);
setCustomOpenAIApiKey(localStorage.getItem('openAIApiKey') || ''); setCustomOpenAIApiKey(localStorage.getItem('openAIApiKey') || '');
@ -164,6 +174,10 @@ const SettingsDialog = ({
localStorage.setItem('chatModelProvider', selectedChatModelProvider!); localStorage.setItem('chatModelProvider', selectedChatModelProvider!);
localStorage.setItem('chatModel', selectedChatModel!); localStorage.setItem('chatModel', selectedChatModel!);
localStorage.setItem(
'ollamaKeepAliveStrategy',
selectedOllamaKeepAliveStrategy!,
);
localStorage.setItem( localStorage.setItem(
'embeddingModelProvider', 'embeddingModelProvider',
selectedEmbeddingModelProvider!, selectedEmbeddingModelProvider!,
@ -293,6 +307,42 @@ const SettingsDialog = ({
/> />
</div> </div>
)} )}
{selectedChatModelProvider &&
selectedChatModelProvider === 'ollama' && (
<div className="flex flex-col space-y-1">
<p className="text-black/70 dark:text-white/70 text-sm">
KeepAlive Strategy
</p>
<Select
value={selectedOllamaKeepAliveStrategy ?? undefined}
onChange={(e) => {
setSelectedOllamaKeepAliveStrategy(
e.target.value,
);
setConfig({
...config,
ollamaKeepAliveStrategy: e.target.value,
});
}}
options={[
{
value: '5m',
label: '5 Minutes',
},
{
value: '60m',
label: '1 Hour',
},
{
value: '-1m',
label: 'Forever',
},
]}
/>
</div>
)}
{selectedChatModelProvider && {selectedChatModelProvider &&
selectedChatModelProvider === 'custom_openai' && ( selectedChatModelProvider === 'custom_openai' && (
<> <>