Merge pull request #4 from sidwan02/routing-via-local-gemini

feat: Add `LocalGeminiClient` and configurable provider for routing

Author: sidwan02

.gemini/settings.json

@@ -18,11 +18,12 @@
     },
     "useModelRouter": {
       "useGemmaRouting": {
-        "enabled": false,
-        "host": "http://localhost:11434",
+        "enabled": true,
+        "host": "http://localhost:8000",
         // "model": "gemma3n:e4b"
         // "model": "gemma3n:e2b"
-        "model": "gemma3:1b"
+        "model": "gemma3:1b",
+        "provider": "litert-lm"
         // "model": "gemma3:270m"
       }
     }

packages/core/src/config/config.ts

@@ -149,6 +149,7 @@ export interface UseModelRouterSettings {
     enabled?: boolean;
     model?: string;
     host?: string;
+    provider?: 'ollama' | 'litert-lm';
   };
 }
 
@@ -565,6 +566,7 @@ export class Config {
         host:
           params.useModelRouter?.useGemmaRouting?.host ??
           'http://localhost:11434',
+        provider: params.useModelRouter?.useGemmaRouting?.provider ?? 'ollama',
       },
     };
     this.disableModelRouterForAuth = params.disableModelRouterForAuth ?? [];

packages/core/src/core/localGeminiClient.ts

@@ -0,0 +1,87 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { GoogleGenAI } from '@google/genai';
+import type { Config } from '../config/config.js';
+import { debugLogger } from '../utils/debugLogger.js';
+import type { Content } from './ollamaChat.js';
+
+/**
+ * A client for making single, non-streaming calls to a local Gemini-compatible API
+ * and expecting a JSON response.
+ */
+export class LocalGeminiClient {
+  private readonly host: string;
+  private readonly model: string;
+  private readonly client: GoogleGenAI;
+
+  constructor(config: Config) {
+    const useGemmaRoutingSettings = config.getUseGemmaRoutingSettings();
+    this.host = useGemmaRoutingSettings?.host || 'http://localhost:8000';
+    this.model = useGemmaRoutingSettings?.model || 'Gemma3-1B-IT';
+
+    if (!this.model.toLowerCase().startsWith('gemma')) {
+      throw new Error(
+        `Invalid model name: ${this.model}. Model name must start with "Gemma" (case-insensitive).`,
+      );
+    }
+
+    this.client = new GoogleGenAI({
+      apiKey: 'no-api-key-needed',
+      httpOptions: {
+        baseUrl: this.host,
+      },
+    });
+  }
+
+  /**
+   * Sends a prompt to the local Gemini model and expects a JSON object in response.
+   * @param contents The history and current prompt.
+   * @param systemInstruction The system prompt.
+   * @returns A promise that resolves to the parsed JSON object.
+   */
+  async generateJson(
+    contents: Content[],
+    systemInstruction: string,
+  ): Promise<object> {
+    debugLogger.log(
+      `[LocalGeminiClient] Sending request to ${this.host} for model ${this.model}`,
+    );
+
+    const geminiContents = contents.map((c) => ({
+      role: c.role === 'model' ? 'model' : 'user',
+      parts: c.parts.map((p) => ({ text: p.text })),
+    }));
+
+    try {
+      const result = await this.client.models.generateContent({
+        model: this.model,
+        contents: geminiContents,
+        config: {
+          responseMimeType: 'application/json',
+          systemInstruction: systemInstruction
+            ? { parts: [{ text: systemInstruction }] }
+            : undefined,
+        },
+      });
+
+      const text = result.text;
+      if (!text) {
+        throw new Error(
+          'Invalid response from Local Gemini API: No text found',
+        );
+      }
+
+      return JSON.parse(text);
+    } catch (error) {
+      debugLogger.error(
+        `[LocalGeminiClient] Failed to generate content:`,
+        error,
+      );
+      throw error;
+    }
+  }
+}

packages/core/src/routing/strategies/classifierStrategy.test.ts

@@ -21,8 +21,13 @@ import { promptIdContext } from '../../utils/promptIdContext.js';
 import type { Content } from '@google/genai';
 import type { ResolvedModelConfig } from '../../services/modelConfigService.js';
 
+import type { OllamaClient as OllamaClientType } from '../../core/ollamaClient.js';
+import type { LocalGeminiClient as LocalGeminiClientType } from '../../core/localGeminiClient.js';
+
 vi.mock('../../core/baseLlmClient.js');
 vi.mock('../../utils/promptIdContext.js');
+vi.mock('../../core/ollamaClient.js');
+vi.mock('../../core/localGeminiClient.js');
 
 describe('ClassifierStrategy', () => {
   let strategy: ClassifierStrategy;
@@ -50,6 +55,7 @@ describe('ClassifierStrategy', () => {
         getResolvedConfig: vi.fn().mockReturnValue(mockResolvedConfig),
       },
       getPreviewFeatures: () => false,
+      getUseGemmaRoutingSettings: () => ({ enabled: false }),
     } as unknown as Config;
     mockBaseLlmClient = {
       generateJson: vi.fn(),
@@ -278,4 +284,88 @@ describe('ClassifierStrategy', () => {
     );
     consoleWarnSpy.mockRestore();
   });
+
+  describe('Local Routing', () => {
+    beforeEach(() => {
+      vi.resetModules();
+    });
+
+    it('should use OllamaClient when local routing is enabled with default provider', async () => {
+      const mockOllamaGenerateJson = vi.fn().mockResolvedValue({
+        reasoning: 'Local Ollama reasoning',
+        model_choice: 'flash',
+      });
+
+      // We need to mock the module before importing/using it in the test context if we want to capture the constructor
+      // However, since we are using vi.mock at the top level, we can just mock the implementation here.
+      const { OllamaClient } = await import('../../core/ollamaClient.js');
+      vi.mocked(OllamaClient).mockImplementation(
+        () =>
+          ({
+            generateJson: mockOllamaGenerateJson,
+          }) as unknown as OllamaClientType,
+      );
+
+      vi.spyOn(mockConfig, 'getUseGemmaRoutingSettings').mockReturnValue({
+        enabled: true,
+        provider: 'ollama',
+      });
+
+      const decision = await strategy.route(
+        mockContext,
+        mockConfig,
+        mockBaseLlmClient,
+      );
+
+      expect(OllamaClient).toHaveBeenCalled();
+      expect(mockOllamaGenerateJson).toHaveBeenCalled();
+      expect(decision).toEqual({
+        model: DEFAULT_GEMINI_FLASH_MODEL,
+        metadata: {
+          source: 'Classifier',
+          latencyMs: expect.any(Number),
+          reasoning: 'Local Ollama reasoning',
+        },
+      });
+    });
+
+    it('should use LocalGeminiClient when local routing is enabled with litert-lm provider', async () => {
+      const mockLocalGeminiGenerateJson = vi.fn().mockResolvedValue({
+        reasoning: 'Local Gemini reasoning',
+        model_choice: 'pro',
+      });
+
+      const { LocalGeminiClient } = await import(
+        '../../core/localGeminiClient.js'
+      );
+      vi.mocked(LocalGeminiClient).mockImplementation(
+        () =>
+          ({
+            generateJson: mockLocalGeminiGenerateJson,
+          }) as unknown as LocalGeminiClientType,
+      );
+
+      vi.spyOn(mockConfig, 'getUseGemmaRoutingSettings').mockReturnValue({
+        enabled: true,
+        provider: 'litert-lm',
+      });
+
+      const decision = await strategy.route(
+        mockContext,
+        mockConfig,
+        mockBaseLlmClient,
+      );
+
+      expect(LocalGeminiClient).toHaveBeenCalled();
+      expect(mockLocalGeminiGenerateJson).toHaveBeenCalled();
+      expect(decision).toEqual({
+        model: DEFAULT_GEMINI_MODEL,
+        metadata: {
+          source: 'Classifier',
+          latencyMs: expect.any(Number),
+          reasoning: 'Local Gemini reasoning',
+        },
+      });
+    });
+  });
 });

packages/core/src/routing/strategies/classifierStrategy.ts

@@ -5,6 +5,7 @@
  */
 
 import { OllamaClient } from '../../core/ollamaClient.js';
+import { LocalGeminiClient } from '../../core/localGeminiClient.js';
 import { z } from 'zod';
 import type { BaseLlmClient } from '../../core/baseLlmClient.js';
 import { promptIdContext } from '../../utils/promptIdContext.js';
@@ -190,8 +191,10 @@ export class ClassifierStrategy implements RoutingStrategy {
 
       let jsonResponse;
       if (config.getUseGemmaRoutingSettings()?.enabled) {
-        debugLogger.log(`[Routing] Using OllamaClient for classifier routing.`);
-        const ollamaClient = new OllamaClient(config);
+        const settings = config.getUseGemmaRoutingSettings();
+        const provider = settings?.provider ?? 'ollama';
+        debugLogger.log(`[Routing] Using ${provider} for classifier routing.`);
+
         const ollamaHistory = finalHistory
           .map(toOllamaContent)
           .filter((c): c is OllamaContent => c !== null);
@@ -202,12 +205,19 @@ export class ClassifierStrategy implements RoutingStrategy {
           ollamaHistory.push(ollamaRequest);
         }
 
-        // debugLogger.log(`[Routing] ollamaHistory:`, ollamaHistory);
-
-        jsonResponse = await ollamaClient.generateJson(
-          ollamaHistory,
-          CLASSIFIER_SYSTEM_PROMPT,
-        );
+        if (provider === 'litert-lm') {
+          const client = new LocalGeminiClient(config);
+          jsonResponse = await client.generateJson(
+            ollamaHistory,
+            CLASSIFIER_SYSTEM_PROMPT,
+          );
+        } else {
+          const ollamaClient = new OllamaClient(config);
+          jsonResponse = await ollamaClient.generateJson(
+            ollamaHistory,
+            CLASSIFIER_SYSTEM_PROMPT,
+          );
+        }
       } else {
         jsonResponse = await baseLlmClient.generateJson({
           modelConfigKey: { model: 'classifier' },

prompt.txt

@@ -484,96 +484,4 @@ comments.
 
 Use hyphens instead of underscores in flag names (e.g. `my-flag` instead of
 `my_flag`).
---- End of Context from: GEMINI.md ---
-
---- Context from: litellm/GEMINI.md ---
-# GEMINI.md
-
-This file provides guidance to Gemini when working with code in this repository.
-
-## Development Commands
-
-### Installation
-- `make install-dev` - Install core development dependencies
-- `make install-proxy-dev` - Install proxy development dependencies with full feature set
-- `make install-test-deps` - Install all test dependencies
-
-### Testing
-- `make test` - Run all tests
-- `make test-unit` - Run unit tests (tests/test_litellm) with 4 parallel workers
-- `make test-integration` - Run integration tests (excludes unit tests)
-- `pytest tests/` - Direct pytest execution
-
-### Code Quality
-- `make lint` - Run all linting (Ruff, MyPy, Black, circular imports, import safety)
-- `make format` - Apply Black code formatting
-- `make lint-ruff` - Run Ruff linting only
-- `make lint-mypy` - Run MyPy type checking only
-
-### Single Test Files
-- `poetry run pytest tests/path/to/test_file.py -v` - Run specific test file
-- `poetry run pytest tests/path/to/test_file.py::test_function -v` - Run specific test
-
-## Architecture Overview
-
-LiteLLM is a unified interface for 100+ LLM providers with two main components:
-
-### Core Library (`litellm/`)
-- **Main entry point**: `litellm/main.py` - Contains core completion() function
-- **Provider implementations**: `litellm/llms/` - Each provider has its own subdirectory
-- **Router system**: `litellm/router.py` + `litellm/router_utils/` - Load balancing and fallback logic
-- **Type definitions**: `litellm/types/` - Pydantic models and type hints
-- **Integrations**: `litellm/integrations/` - Third-party observability, caching, logging
-- **Caching**: `litellm/caching/` - Multiple cache backends (Redis, in-memory, S3, etc.)
-
-### Proxy Server (`litellm/proxy/`)
-- **Main server**: `proxy_server.py` - FastAPI application
-- **Authentication**: `auth/` - API key management, JWT, OAuth2
-- **Database**: `db/` - Prisma ORM with PostgreSQL/SQLite support
-- **Management endpoints**: `management_endpoints/` - Admin APIs for keys, teams, models
-- **Pass-through endpoints**: `pass_through_endpoints/` - Provider-specific API forwarding
-- **Guardrails**: `guardrails/` - Safety and content filtering hooks
-- **UI Dashboard**: Served from `_experimental/out/` (Next.js build)
-
-## Key Patterns
-
-### Provider Implementation
-- Providers inherit from base classes in `litellm/llms/base.py`
-- Each provider has transformation functions for input/output formatting
-- Support both sync and async operations
-- Handle streaming responses and function calling
-
-### Error Handling
-- Provider-specific exceptions mapped to OpenAI-compatible errors
-- Fallback logic handled by Router system
-- Comprehensive logging through `litellm/_logging.py`
-
-### Configuration
-- YAML config files for proxy server (see `proxy/example_config_yaml/`)
-- Environment variables for API keys and settings
-- Database schema managed via Prisma (`proxy/schema.prisma`)
-
-## Development Notes
-
-### Code Style
-- Uses Black formatter, Ruff linter, MyPy type checker
-- Pydantic v2 for data validation
-- Async/await patterns throughout
-- Type hints required for all public APIs
-
-### Testing Strategy
-- Unit tests in `tests/test_litellm/`
-- Integration tests for each provider in `tests/llm_translation/`
-- Proxy tests in `tests/proxy_unit_tests/`
-- Load tests in `tests/load_tests/`
-
-### Database Migrations
-- Prisma handles schema migrations
-- Migration files auto-generated with `prisma migrate dev`
-- Always test migrations against both PostgreSQL and SQLite
-
-### Enterprise Features
-- Enterprise-specific code in `enterprise/` directory
-- Optional features enabled via environment variables
-- Separate licensing and authentication for enterprise features
---- End of Context from: litellm/GEMINI.md ---
+--- End of Context from: GEMINI.md ---