feat: added local LLM context length option

src/agent/llm/LocalBrowserLLMProvider.test.ts

@@ -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

@@ -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}`;