3dtours/KGStudio
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: added local LLM context length option

Browse Source
This commit is contained in:
Xiaohan-Tian
2026-05-14 19:44:25 -07:00
parent dbf0c20542
commit a48966d761
13 changed files with 467 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
public/config.json
+3
View File
@@ -25,6 +25,9 @@
"api_key": "",
"base_url": "",
"model": ""
},
"local_browser": {
"context_length": 32768
},
"soundfont": {
"base_url": "https://cdn.jsdelivr.net/npm/soundfont-for-samplers/FluidR3_GM/"
src/agent/llm/LocalBrowserLLMProvider.test.ts
+106
View File
@@ -0,0 +1,106 @@
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
const {
configGetMock,
ensureRuntimeSupportedMock,
notifyLoadProgressMock,
notifyLoadStartMock,
notifyCacheReadyMock,
notifyLoadErrorMock,
loadModelReaderWithCacheMock,
} = vi.hoisted(() => ({
configGetMock: vi.fn(),
ensureRuntimeSupportedMock: vi.fn(async () => undefined),
notifyLoadProgressMock: vi.fn(),
notifyLoadStartMock: vi.fn(),
notifyCacheReadyMock: vi.fn(),
notifyLoadErrorMock: vi.fn(),
loadModelReaderWithCacheMock: vi.fn(),
}));
vi.mock('../../core/config/ConfigManager', () => ({
ConfigManager: {
instance: () => ({
get: configGetMock,
}),
},
}));
vi.mock('../../util/localLLMModelManager', () => ({
LocalLLMModelManager: {
ensureRuntimeSupported: ensureRuntimeSupportedMock,
notifyLoadProgress: notifyLoadProgressMock,
notifyLoadStart: notifyLoadStartMock,
notifyCacheReady: notifyCacheReadyMock,
notifyLoadError: notifyLoadErrorMock,
},
}));
vi.mock('../../util/localLLMModelCache', () => ({
LocalLLMModelCache: {
loadModelReaderWithCache: loadModelReaderWithCacheMock,
},
}));
import { LocalBrowserLLMProvider } from './LocalBrowserLLMProvider';
describe('LocalBrowserLLMProvider', () => {
afterEach(() => {
vi.restoreAllMocks();
});
beforeEach(() => {
configGetMock.mockReset();
ensureRuntimeSupportedMock.mockClear();
notifyLoadProgressMock.mockClear();
notifyLoadStartMock.mockClear();
notifyCacheReadyMock.mockClear();
notifyLoadErrorMock.mockClear();
loadModelReaderWithCacheMock.mockReset();
});
async function runProviderAndCaptureOptions(configValue: unknown): Promise<Record<string, unknown>> {
const createFromOptionsMock = vi.fn(async (_fileset: unknown, options: Record<string, unknown>) => ({
generateResponse: (_prompt: string, callback: (partial: string, done: boolean) => void) => {
callback('hello', true);
},
sizeInTokens: (text: string) => text.length,
}));
configGetMock.mockReturnValue(configValue);
loadModelReaderWithCacheMock.mockResolvedValue({
reader: new Uint8Array([1, 2, 3]),
totalBytes: 3,
fromCache: true,
cacheWritePromise: null,
});
vi.spyOn(LocalBrowserLLMProvider.prototype as never, 'getMediaPipeModule' as never).mockResolvedValue({
FilesetResolver: {
forGenAiTasks: vi.fn(async () => ({})),
},
LlmInference: {
createFromOptions: createFromOptionsMock,
},
});
const provider = new LocalBrowserLLMProvider();
const chunks: unknown[] = [];
for await (const chunk of provider.generateStream([])) {
chunks.push(chunk);
}
expect(chunks.length).toBeGreaterThan(0);
return createFromOptionsMock.mock.calls[0][1] as Record<string, unknown>;
}
it('uses the configured maxTokens value', async () => {
const options = await runProviderAndCaptureOptions(65536);
expect(options.maxTokens).toBe(65536);
});
it('falls back to 32768 when config is invalid', async () => {
const options = await runProviderAndCaptureOptions(99999);
expect(options.maxTokens).toBe(32768);
});
});
src/agent/llm/LocalBrowserLLMProvider.ts
+22 -2
View File
@@ -9,9 +9,15 @@ import {
parseToolCalls,
stripToolProtocol,
} from './gemmaToolProtocol';
import { LOCAL_LLM_MODEL_FILENAME, LOCAL_LLM_MODEL_URL } from '../../util/localLLMConfig';
import {
LOCAL_LLM_DEFAULT_CONTEXT_LENGTH,
LOCAL_LLM_MODEL_FILENAME,
LOCAL_LLM_MODEL_URL,
normalizeLocalLLMContextLength,
} from '../../util/localLLMConfig';
import { LocalLLMModelCache } from '../../util/localLLMModelCache';
import type { LLMProvider } from './LLMProvider';
import { ConfigManager } from '../../core/config/ConfigManager';
type MediaPipeGenAI = {
FilesetResolver: {
@@ -60,6 +66,9 @@ export class LocalBrowserLLMProvider implements LLMProvider {
return this.inference;
}
const maxTokens = this.getConfiguredContextLength();
console.log(`[localLLM] Initializing with max context length: ${maxTokens} tokens`);
const [{ FilesetResolver, LlmInference }, modelLoad] = await Promise.all([
this.getMediaPipeModule(),
LocalLLMModelCache.loadModelReaderWithCache(
@@ -86,7 +95,7 @@ export class LocalBrowserLLMProvider implements LLMProvider {
modelAssetBuffer: modelLoad.reader,
},
numResponses: 1,
maxTokens: 32768,
maxTokens,
topK: 64,
temperature: 1.0,
});
@@ -112,6 +121,17 @@ export class LocalBrowserLLMProvider implements LLMProvider {
return importMediaPipe();
}
private getConfiguredContextLength(): number {
try {
const configManager = ConfigManager.instance();
return normalizeLocalLLMContextLength(
configManager.get('general.local_browser.context_length'),
);
} catch {
return LOCAL_LLM_DEFAULT_CONTEXT_LENGTH;
}
}
private applyTemplate(message: { role: 'user' | 'model'; text: string }): string {
const template = PROMPT_TEMPLATE[message.role];
return `${template.pre}${message.text}${template.post}`;
src/components/ChatBox.tsx
+1
View File
@@ -200,6 +200,7 @@ const ChatBox: React.FC<ChatBoxProps> = ({ isVisible }) => {
const unsubscribe = configManager.addChangeListener((changedKeys) => {
if (
changedKeys.includes('general.llm_provider') ||
changedKeys.includes('general.local_browser.context_length') ||
changedKeys.some(k => k.startsWith('general.openai.')) ||
changedKeys.some(k => k.startsWith('general.openai_compatible.'))
) {
src/components/settings/sections/GeneralSettings.test.tsx
+120
View File
@@ -0,0 +1,120 @@
import React from 'react';
import { beforeEach, describe, expect, it, vi } from 'vitest';
import { fireEvent, render, screen, waitFor } from '@testing-library/react';
import GeneralSettings from './GeneralSettings';
const configState = new Map<string, unknown>([
['general.llm_provider', 'local_browser'],
['general.persist_api_keys_non_localhost', false],
['general.openai.api_key', ''],
['general.openai.model', 'gpt-5.4-mini'],
['general.openai.flex', false],
['general.gemini.api_key', ''],
['general.gemini.model', 'gemini-2.5-flash'],
['general.claude.api_key', ''],
['general.claude.model', 'claude-sonnet-4.6'],
['general.claude_openrouter.api_key', ''],
['general.claude_openrouter.base_url', 'https://openrouter.ai/api/v1'],
['general.claude_openrouter.model', 'anthropic/claude-sonnet-4.6'],
['general.openai_compatible.api_key', ''],
['general.openai_compatible.base_url', ''],
['general.openai_compatible.model', ''],
['general.local_browser.context_length', 65536],
['general.soundfont.base_url', 'https://cdn.jsdelivr.net/npm/soundfont-for-samplers/FluidR3_GM/'],
['general.kgone.enabled', false],
['general.kgone.base_url', 'http://127.0.0.1:8000'],
]);
const configManagerMock = {
getIsInitialized: vi.fn(() => true),
initialize: vi.fn().mockResolvedValue(undefined),
get: vi.fn((key: string) => configState.get(key)),
set: vi.fn(async (key: string, value: unknown) => {
configState.set(key, value);
}),
isKGOneServerManaged: vi.fn(() => false),
isSoundfontServerManaged: vi.fn(() => false),
};
vi.mock('../../../core/config/ConfigManager', () => ({
ConfigManager: {
instance: () => configManagerMock,
},
}));
vi.mock('../../../util/localLLMModelManager', () => ({
LocalLLMModelManager: {
getState: () => ({
isCached: false,
isChecking: false,
isDownloading: false,
isDeleting: false,
progressPercent: 0,
progressText: '',
error: '',
runtimeSupport: {
supported: true,
webgpuExposed: true,
crossOriginIsolated: true,
sharedArrayBufferAvailable: true,
secureContext: true,
reason: null,
},
}),
subscribe: (listener: (state: unknown) => void) => {
listener({
isCached: false,
isChecking: false,
isDownloading: false,
isDeleting: false,
progressPercent: 0,
progressText: '',
error: '',
runtimeSupport: {
supported: true,
webgpuExposed: true,
crossOriginIsolated: true,
sharedArrayBufferAvailable: true,
secureContext: true,
reason: null,
},
});
return () => {};
},
deleteCachedModel: vi.fn().mockResolvedValue(undefined),
},
}));
describe('GeneralSettings', () => {
beforeEach(() => {
configState.set('general.local_browser.context_length', 65536);
configManagerMock.get.mockClear();
configManagerMock.set.mockClear();
});
it('renders the local context length selector and VRAM hint', async () => {
render(<GeneralSettings />);
expect(await screen.findByText('Gemma 4 E4B Local Runtime')).toBeTruthy();
expect(screen.getByLabelText('Context Length')).toBeTruthy();
expect(screen.getByText(/require more VRAM/i)).toBeTruthy();
});
it('initializes the local context length from config', async () => {
render(<GeneralSettings />);
const select = await screen.findByLabelText('Context Length');
expect((select as HTMLSelectElement).value).toBe('65536');
});
it('persists local context length changes', async () => {
render(<GeneralSettings />);
const select = await screen.findByLabelText('Context Length');
fireEvent.change(select, { target: { value: '131072' } });
await waitFor(() => {
expect(configManagerMock.set).toHaveBeenCalledWith('general.local_browser.context_length', 131072);
});
});
});
src/components/settings/sections/GeneralSettings.tsx
+44 -1
View File
@@ -1,7 +1,15 @@
import React, { useState, useEffect, useCallback, useMemo } from 'react';
import { ConfigManager } from '../../../core/config/ConfigManager';
import { LocalLLMModelManager, type LocalLLMModelState } from '../../../util/localLLMModelManager';
import { LOCAL_LLM_DISPLAY_NAME, LOCAL_LLM_PROVIDER_KEY } from '../../../util/localLLMConfig';
import {
formatLocalLLMContextLength,
LOCAL_LLM_CONTEXT_LENGTH_OPTIONS,
LOCAL_LLM_DEFAULT_CONTEXT_LENGTH,
LOCAL_LLM_DISPLAY_NAME,
LOCAL_LLM_PROVIDER_KEY,
normalizeLocalLLMContextLength,
type LocalLLMContextLength,
} from '../../../util/localLLMConfig';
const GeneralSettings: React.FC = () => {
const [llmProvider, setLlmProvider] = useState<string>(LOCAL_LLM_PROVIDER_KEY);
@@ -24,6 +32,7 @@ const GeneralSettings: React.FC = () => {
const [kgoneBaseUrl, setKgoneBaseUrl] = useState<string>('');
const [kgoneServerManaged, setKgoneServerManaged] = useState<boolean>(false);
const [soundfontServerManaged, setSoundfontServerManaged] = useState<boolean>(false);
const [localContextLength, setLocalContextLength] = useState<LocalLLMContextLength>(LOCAL_LLM_DEFAULT_CONTEXT_LENGTH);
const [localModelState, setLocalModelState] = useState<LocalLLMModelState>(LocalLLMModelManager.getState());
const configManager = ConfigManager.instance();
@@ -64,6 +73,7 @@ const GeneralSettings: React.FC = () => {
setCompatibleKey((configManager.get('general.openai_compatible.api_key') as string) || '');
setCompatibleBaseUrl((configManager.get('general.openai_compatible.base_url') as string) || '');
setCompatibleModel((configManager.get('general.openai_compatible.model') as string) || '');
setLocalContextLength(normalizeLocalLLMContextLength(configManager.get('general.local_browser.context_length')));
setSoundfontBaseUrl((configManager.get('general.soundfont.base_url') as string) || '');
setKgoneEnabled((configManager.get('general.kgone.enabled') as boolean) ?? false);
setKgoneBaseUrl((configManager.get('general.kgone.base_url') as string) || '');
@@ -210,6 +220,18 @@ const GeneralSettings: React.FC = () => {
}
};
const handleLocalContextLengthChange = async (value: string) => {
const parsed = Number(value);
const normalized = normalizeLocalLLMContextLength(parsed);
setLocalContextLength(normalized);
try {
await configManager.set('general.local_browser.context_length', normalized);
console.log('Local browser context length changed to:', normalized);
} catch (error) {
console.error('Failed to save local browser context length:', error);
}
};
// NOTE: Gemini and Claude are not supported yet due to CORS issues.
return (
<div className="settings-section">
@@ -279,6 +301,27 @@ const GeneralSettings: React.FC = () => {
</div>
</div>
<div className="settings-item">
<label className="settings-label" htmlFor="local-llm-context-length">
Context Length
</label>
<select
id="local-llm-context-length"
className="settings-select"
value={localContextLength}
onChange={(e) => void handleLocalContextLengthChange(e.target.value)}
>
{LOCAL_LLM_CONTEXT_LENGTH_OPTIONS.map(option => (
<option key={option} value={option}>
{formatLocalLLMContextLength(option)}
</option>
))}
</select>
<div className="settings-help" style={{ fontSize: '12px', color: '#888', marginTop: '4px' }}>
Larger context lengths require more VRAM and may also reduce performance as conversations become longer.
</div>
</div>
{!localModelState.isCached && !localModelState.isDownloading && localModelState.runtimeSupport.supported && (
<div className="settings-help" style={{ fontSize: '12px', color: '#888', marginTop: '4px', marginBottom: '8px' }}>
The local model downloads automatically the next time you chat with `Local LLM (Browser)`.
src/constants/coreConstants.ts
+1 -1
View File
@@ -104,7 +104,7 @@ export const OPFS_CONSTANTS = {
export const CONFIG_UPGRADER_CONSTANTS = {
VERSION_KEY: '__config_version',
CURRENT_VERSION: 2,
CURRENT_VERSION: 3,
};
export const URL_CONSTANTS = {
src/core/config-upgrader/KGConfigUpgrader.ts
+5
View File
@@ -2,6 +2,7 @@ import { KGConfigStorage } from '../io/KGConfigStorage';
import { CONFIG_UPGRADER_CONSTANTS } from '../../constants/coreConstants';
import { upgradeConfigToV1 } from './upgradeConfigToV1';
import { upgradeConfigToV2 } from './upgradeConfigToV2';
import { upgradeConfigToV3 } from './upgradeConfigToV3';
/**
* KGConfigUpgrader — Orchestrates app-level migrations (e.g., storage backend changes).
@@ -38,6 +39,10 @@ export class KGConfigUpgrader {
await upgradeConfigToV2();
break;
}
case 3: {
await upgradeConfigToV3();
break;
}
default: {
throw new Error(`No config upgrader found for version ${nextVersion}`);
}
src/core/config-upgrader/upgradeConfigToV3.test.ts
+88
View File
@@ -0,0 +1,88 @@
import { beforeEach, describe, expect, it, vi } from 'vitest';
const configStore = new Map<string, { name: string; data: Record<string, unknown>; lastModified: number }>();
vi.mock('../io/KGConfigStorage', () => ({
KGConfigStorage: {
getInstance: () => ({
getRaw: vi.fn(async (name: string) => configStore.get(name)?.data ?? null),
saveRaw: vi.fn(async (name: string, data: Record<string, unknown>) => {
configStore.set(name, { name, data, lastModified: Date.now() });
}),
}),
},
}));
import { upgradeConfigToV3 } from './upgradeConfigToV3';
describe('upgradeConfigToV3', () => {
beforeEach(() => {
configStore.clear();
});
it('adds the default local browser context length when missing', async () => {
configStore.set('userConfig', {
name: 'userConfig',
data: {
general: {
llm_provider: 'local_browser',
},
},
lastModified: Date.now(),
});
await upgradeConfigToV3();
expect(
((configStore.get('userConfig')?.data.general as Record<string, unknown>).local_browser as Record<string, unknown>).context_length,
).toBe(32768);
});
it.each([32768, 65536, 131072])('preserves existing context length %s', async (existingValue) => {
configStore.set('userConfig', {
name: 'userConfig',
data: {
general: {
local_browser: {
context_length: existingValue,
},
},
},
lastModified: Date.now(),
});
await upgradeConfigToV3();
expect(
((configStore.get('userConfig')?.data.general as Record<string, unknown>).local_browser as Record<string, unknown>).context_length,
).toBe(existingValue);
});
it('is a no-op when general is missing', async () => {
configStore.set('userConfig', {
name: 'userConfig',
data: {},
lastModified: Date.now(),
});
await upgradeConfigToV3();
expect(configStore.get('userConfig')?.data).toEqual({});
});
it('is a no-op when general is malformed', async () => {
configStore.set('userConfig', {
name: 'userConfig',
data: {
general: 'invalid',
},
lastModified: Date.now(),
});
await upgradeConfigToV3();
expect(configStore.get('userConfig')?.data).toEqual({
general: 'invalid',
});
});
});
src/core/config-upgrader/upgradeConfigToV3.ts
+35
View File
@@ -0,0 +1,35 @@
import { KGConfigStorage } from '../io/KGConfigStorage';
import { LOCAL_LLM_DEFAULT_CONTEXT_LENGTH } from '../../util/localLLMConfig';
const CONFIG_KEY = 'userConfig';
export async function upgradeConfigToV3(): Promise<void> {
const storage = KGConfigStorage.getInstance();
const rawConfig = await storage.getRaw(CONFIG_KEY);
if (!rawConfig || typeof rawConfig !== 'object') {
return;
}
const config = rawConfig as Record<string, unknown>;
const general = config.general;
if (!general || typeof general !== 'object') {
return;
}
const generalRecord = general as Record<string, unknown>;
const localBrowser = generalRecord.local_browser;
if (!localBrowser || typeof localBrowser !== 'object') {
generalRecord.local_browser = { context_length: LOCAL_LLM_DEFAULT_CONTEXT_LENGTH };
await storage.saveRaw(CONFIG_KEY, config);
return;
}
const localBrowserRecord = localBrowser as Record<string, unknown>;
if ('context_length' in localBrowserRecord) {
return;
}
localBrowserRecord.context_length = LOCAL_LLM_DEFAULT_CONTEXT_LENGTH;
await storage.saveRaw(CONFIG_KEY, config);
}
src/core/config/ConfigManager.ts
+6
View File
@@ -8,6 +8,9 @@ interface AppConfig {
language: string;
llm_provider: 'local_browser' | 'openai' | 'gemini' | 'claude' | 'claude_openrouter' | 'openai_compatible';
persist_api_keys_non_localhost: boolean;
local_browser: {
context_length: 32768 | 65536 | 131072;
};
openai: {
api_key: string;
flex: boolean;
@@ -208,6 +211,9 @@ export class ConfigManager {
base_url: '',
model: ''
},
local_browser: {
context_length: 32768
},
soundfont: {
base_url: 'https://cdn.jsdelivr.net/npm/soundfont-for-samplers/FluidR3_GM/'
},
src/util/localLLMConfig.test.ts
+19
View File
@@ -0,0 +1,19 @@
import { describe, expect, it } from 'vitest';
import {
formatLocalLLMContextLength,
LOCAL_LLM_DEFAULT_CONTEXT_LENGTH,
normalizeLocalLLMContextLength,
} from './localLLMConfig';
describe('localLLMConfig', () => {
it('defaults invalid context lengths to 32768', () => {
expect(normalizeLocalLLMContextLength(undefined)).toBe(LOCAL_LLM_DEFAULT_CONTEXT_LENGTH);
expect(normalizeLocalLLMContextLength(12345)).toBe(LOCAL_LLM_DEFAULT_CONTEXT_LENGTH);
});
it('formats context lengths using k suffixes', () => {
expect(formatLocalLLMContextLength(32768)).toBe('32k');
expect(formatLocalLLMContextLength(65536)).toBe('64k');
expect(formatLocalLLMContextLength(131072)).toBe('128k');
});
});
src/util/localLLMConfig.ts
+17
View File
@@ -6,6 +6,10 @@ export const LOCAL_LLM_DISPLAY_NAME = 'Gemma 4 E4B';
export const LOCAL_LLM_LEGACY_FILENAMES = [
'gemma-3n-E4B-it-int4-Web.litertlm',
];
export const LOCAL_LLM_CONTEXT_LENGTH_OPTIONS = [32768, 65536, 131072] as const;
export const LOCAL_LLM_DEFAULT_CONTEXT_LENGTH = 32768;
export type LocalLLMContextLength = typeof LOCAL_LLM_CONTEXT_LENGTH_OPTIONS[number];
export interface LocalLLMRuntimeSupport {
supported: boolean;
@@ -40,3 +44,16 @@ export function detectLocalLLMRuntimeSupport(): LocalLLMRuntimeSupport {
reason,
};
}
export function isLocalLLMContextLength(value: unknown): value is LocalLLMContextLength {
return typeof value === 'number'
&& (LOCAL_LLM_CONTEXT_LENGTH_OPTIONS as readonly number[]).includes(value);
}
export function normalizeLocalLLMContextLength(value: unknown): LocalLLMContextLength {
return isLocalLLMContextLength(value) ? value : LOCAL_LLM_DEFAULT_CONTEXT_LENGTH;
}
export function formatLocalLLMContextLength(value: LocalLLMContextLength): string {
return `${Math.round(value / 1024)}k`;
}