From 8d2ee595338b4155aaf67a6f54cb99fb7bf6bc1e Mon Sep 17 00:00:00 2001 From: binyangzhu000-sudo <224954946+binyangzhu000-sudo@users.noreply.github.com> Date: Fri, 10 Jul 2026 14:16:51 +0800 Subject: [PATCH] Add Atlas Cloud custom LLM shortcut --- README.md | 10 +++- docs/cli/authentication.md | 11 ++++ packages/cli/src/config/auth.test.ts | 26 +++++++++ packages/cli/src/config/auth.ts | 8 ++- packages/cli/src/gemini.tsx | 2 +- packages/cli/src/ui/components/Footer.tsx | 19 +++++-- packages/cli/src/ui/hooks/useAuthCommand.ts | 5 +- packages/core/src/config/config.ts | 5 +- packages/core/src/custom_llm/config.test.ts | 60 +++++++++++++++++++++ packages/core/src/custom_llm/config.ts | 41 ++++++++++++++ packages/core/src/custom_llm/index.ts | 20 +++---- packages/core/src/index.ts | 1 + 12 files changed, 184 insertions(+), 24 deletions(-) create mode 100644 packages/core/src/custom_llm/config.test.ts create mode 100644 packages/core/src/custom_llm/config.ts diff --git a/README.md b/README.md index 184db529..cd147157 100644 --- a/README.md +++ b/README.md @@ -81,6 +81,15 @@ export CUSTOM_LLM_TOP_P=1 # Top P (default: 1) When these variables are set, Easy LLM CLI will use your custom LLM instead of the default Gemini model. +For Atlas Cloud, you can use the OpenAI-compatible endpoint with fewer variables: + +```bash +export ATLASCLOUD_API_KEY="your-atlascloud-api-key" +# Optional: export CUSTOM_LLM_MODEL_NAME="deepseek-ai/deepseek-v4-pro" +``` + +By default, Atlas Cloud uses `https://api.atlascloud.ai/v1` with `qwen/qwen3.5-flash`. + ## Examples @@ -191,4 +200,3 @@ Use MCP servers to integrate your local system tools with your enterprise collab ```text > Organize my PDF invoices by month of expenditure. ``` - diff --git a/docs/cli/authentication.md b/docs/cli/authentication.md index f797e454..31d83e4e 100644 --- a/docs/cli/authentication.md +++ b/docs/cli/authentication.md @@ -75,6 +75,17 @@ The Gemini CLI requires you to authenticate with Google's AI services. On initia - It automatically uses the credentials of the logged-in user in the Cloud Shell environment. - This is the default authentication method when running in Cloud Shell and no other method is configured. +5. **OpenAI-compatible custom LLM API:** + - Set `USE_CUSTOM_LLM=true`, `CUSTOM_LLM_ENDPOINT`, and `CUSTOM_LLM_MODEL_NAME` for any OpenAI-compatible provider. + - Atlas Cloud can be enabled with only `ATLASCLOUD_API_KEY`; the CLI will use `https://api.atlascloud.ai/v1` and the `qwen/qwen3.5-flash` model by default. + - You can still override the Atlas Cloud default model with `CUSTOM_LLM_MODEL_NAME` when needed. + + ```bash + export ATLASCLOUD_API_KEY="YOUR_ATLASCLOUD_API_KEY" + # Optional override: + export CUSTOM_LLM_MODEL_NAME="deepseek-ai/deepseek-v4-pro" + ``` + ### Persisting Environment Variables with `.env` Files You can create a **`.gemini/.env`** file in your project directory or in your home directory. Creating a plain **`.env`** file also works, but `.gemini/.env` is recommended to keep Gemini variables isolated from other tools. diff --git a/packages/cli/src/config/auth.test.ts b/packages/cli/src/config/auth.test.ts index 96defb1e..9e89df64 100644 --- a/packages/cli/src/config/auth.test.ts +++ b/packages/cli/src/config/auth.test.ts @@ -67,6 +67,32 @@ describe('validateAuthMethod', () => { }); }); + describe('CUSTOM_LLM_API', () => { + it('should return null if ATLASCLOUD_API_KEY is set', () => { + process.env.ATLASCLOUD_API_KEY = 'test-atlas-key'; + expect(validateAuthMethod(AuthType.CUSTOM_LLM_API)).toBeNull(); + }); + + it('should return null if CUSTOM_LLM_ENDPOINT and CUSTOM_LLM_MODEL_NAME are set', () => { + process.env.CUSTOM_LLM_ENDPOINT = 'https://example.test/v1'; + process.env.CUSTOM_LLM_MODEL_NAME = 'example-model'; + expect(validateAuthMethod(AuthType.CUSTOM_LLM_API)).toBeNull(); + }); + + it('should return an endpoint error if custom LLM variables are missing', () => { + expect(validateAuthMethod(AuthType.CUSTOM_LLM_API)).toBe( + 'CUSTOM_LLM_ENDPOINT environment variable not found. Add that to your .env and try again, no reload needed! Alternatively set ATLASCLOUD_API_KEY to use the Atlas Cloud OpenAI-compatible endpoint.', + ); + }); + + it('should return a model error if CUSTOM_LLM_MODEL_NAME is missing', () => { + process.env.CUSTOM_LLM_ENDPOINT = 'https://example.test/v1'; + expect(validateAuthMethod(AuthType.CUSTOM_LLM_API)).toBe( + 'CUSTOM_LLM_MODEL_NAME environment variable not found. Add that to your .env and try again, no reload needed! Alternatively set ATLASCLOUD_API_KEY to use the Atlas Cloud default model.', + ); + }); + }); + it('should return an error message for an invalid auth method', () => { expect(validateAuthMethod('invalid-method')).toBe( 'Invalid auth method selected.', diff --git a/packages/cli/src/config/auth.ts b/packages/cli/src/config/auth.ts index 122c0977..a52d5024 100644 --- a/packages/cli/src/config/auth.ts +++ b/packages/cli/src/config/auth.ts @@ -39,12 +39,16 @@ export const validateAuthMethod = (authMethod: string): string | null => { } if (authMethod === AuthType.CUSTOM_LLM_API) { + if (process.env.ATLASCLOUD_API_KEY) { + return null; + } + if (!process.env.CUSTOM_LLM_ENDPOINT) { - return 'CUSTOM_LLM_ENDPOINT environment variable not found. Add that to your .env and try again, no reload needed!'; + return 'CUSTOM_LLM_ENDPOINT environment variable not found. Add that to your .env and try again, no reload needed! Alternatively set ATLASCLOUD_API_KEY to use the Atlas Cloud OpenAI-compatible endpoint.'; } if (!process.env.CUSTOM_LLM_MODEL_NAME) { - return 'CUSTOM_LLM_MODEL_NAME environment variable not found. Add that to your .env and try again, no reload needed!'; + return 'CUSTOM_LLM_MODEL_NAME environment variable not found. Add that to your .env and try again, no reload needed! Alternatively set ATLASCLOUD_API_KEY to use the Atlas Cloud default model.'; } return null; } diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index 1ea0cd06..8bf90e6a 100644 --- a/packages/cli/src/gemini.tsx +++ b/packages/cli/src/gemini.tsx @@ -122,7 +122,7 @@ export async function main() { } } - if (process.env.USE_CUSTOM_LLM) { + if (process.env.USE_CUSTOM_LLM || process.env.ATLASCLOUD_API_KEY) { settings.setValue( SettingScope.User, 'selectedAuthType', diff --git a/packages/cli/src/ui/components/Footer.tsx b/packages/cli/src/ui/components/Footer.tsx index 9b7f0f52..85d95160 100644 --- a/packages/cli/src/ui/components/Footer.tsx +++ b/packages/cli/src/ui/components/Footer.tsx @@ -7,7 +7,12 @@ import React from 'react'; import { Box, Text } from 'ink'; import { Colors } from '../colors.js'; -import { shortenPath, tildeifyPath, tokenLimit } from '@google/gemini-cli-core'; +import { + resolveCustomLLMRuntimeConfig, + shortenPath, + tildeifyPath, + tokenLimit, +} from '@google/gemini-cli-core'; import { ConsoleSummaryDisplay } from './ConsoleSummaryDisplay.js'; import process from 'node:process'; import { MemoryUsageDisplay } from './MemoryUsageDisplay.js'; @@ -41,10 +46,14 @@ export const Footer: React.FC = ({ const percentage = promptTokenCount / limit; let modelInfo = model; - if (process.env.USE_CUSTOM_LLM && process.env.CUSTOM_LLM_MODEL_NAME) { - modelInfo = - process.env.CUSTOM_LLM_MODEL_NAME + - `(${process.env.CUSTOM_LLM_PROVIDER})`; + const customConfig = resolveCustomLLMRuntimeConfig(); + if ( + (process.env.USE_CUSTOM_LLM || process.env.ATLASCLOUD_API_KEY) && + customConfig.modelName + ) { + modelInfo = customConfig.provider + ? `${customConfig.modelName}(${customConfig.provider})` + : customConfig.modelName; } return ( diff --git a/packages/cli/src/ui/hooks/useAuthCommand.ts b/packages/cli/src/ui/hooks/useAuthCommand.ts index cc5f9f78..6d88d4c9 100644 --- a/packages/cli/src/ui/hooks/useAuthCommand.ts +++ b/packages/cli/src/ui/hooks/useAuthCommand.ts @@ -14,7 +14,10 @@ import { } from '@google/gemini-cli-core'; function isInitAuth(settings: LoadedSettings) { - if (process.env.USE_CUSTOM_LLM && process.env.USE_CUSTOM_LLM !== 'false') { + if ( + (process.env.USE_CUSTOM_LLM && process.env.USE_CUSTOM_LLM !== 'false') || + process.env.ATLASCLOUD_API_KEY + ) { return false; } if (settings.merged.selectedAuthType !== AuthType.CUSTOM_LLM_API) { diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 2f738c14..2479370d 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -44,6 +44,7 @@ import { DEFAULT_GEMINI_FLASH_MODEL, } from './models.js'; import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js'; +import { resolveCustomLLMRuntimeConfig } from '../custom_llm/config.js'; export enum ApprovalMode { DEFAULT = 'default', @@ -278,8 +279,8 @@ export class Config { } getModel(): string { - if (process.env.USE_CUSTOM_LLM) { - return process.env.CUSTOM_LLM_MODEL_NAME || ''; + if (process.env.USE_CUSTOM_LLM || process.env.ATLASCLOUD_API_KEY) { + return resolveCustomLLMRuntimeConfig().modelName; } return this.contentGeneratorConfig?.model || this.model; } diff --git a/packages/core/src/custom_llm/config.test.ts b/packages/core/src/custom_llm/config.test.ts new file mode 100644 index 00000000..15c27e09 --- /dev/null +++ b/packages/core/src/custom_llm/config.test.ts @@ -0,0 +1,60 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it } from 'vitest'; +import { + ATLASCLOUD_DEFAULT_MODEL, + ATLASCLOUD_OPENAI_BASE_URL, + resolveCustomLLMRuntimeConfig, +} from './config.js'; + +describe('resolveCustomLLMRuntimeConfig', () => { + it('uses existing CUSTOM_LLM_* variables unchanged', () => { + const config = resolveCustomLLMRuntimeConfig({ + CUSTOM_LLM_API_KEY: 'custom-key', + CUSTOM_LLM_ENDPOINT: 'https://example.test/v1', + CUSTOM_LLM_MODEL_NAME: 'custom-model', + CUSTOM_LLM_PROVIDER: 'custom-provider', + CUSTOM_LLM_TEMPERATURE: '0.3', + CUSTOM_LLM_MAX_TOKENS: '4096', + CUSTOM_LLM_TOP_P: '0.8', + }); + + expect(config).toEqual({ + apiKey: 'custom-key', + baseURL: 'https://example.test/v1', + modelName: 'custom-model', + provider: 'custom-provider', + temperature: 0.3, + maxTokens: 4096, + topP: 0.8, + }); + }); + + it('uses Atlas Cloud defaults when ATLASCLOUD_API_KEY is set', () => { + const config = resolveCustomLLMRuntimeConfig({ + ATLASCLOUD_API_KEY: 'atlas-key', + }); + + expect(config.apiKey).toBe('atlas-key'); + expect(config.baseURL).toBe(ATLASCLOUD_OPENAI_BASE_URL); + expect(config.modelName).toBe(ATLASCLOUD_DEFAULT_MODEL); + expect(config.provider).toBe('atlascloud'); + }); + + it('lets CUSTOM_LLM_* override Atlas Cloud defaults', () => { + const config = resolveCustomLLMRuntimeConfig({ + ATLASCLOUD_API_KEY: 'atlas-key', + CUSTOM_LLM_MODEL_NAME: 'deepseek-ai/deepseek-v4-pro', + CUSTOM_LLM_PROVIDER: 'atlascloud-eu', + }); + + expect(config.apiKey).toBe('atlas-key'); + expect(config.baseURL).toBe(ATLASCLOUD_OPENAI_BASE_URL); + expect(config.modelName).toBe('deepseek-ai/deepseek-v4-pro'); + expect(config.provider).toBe('atlascloud-eu'); + }); +}); diff --git a/packages/core/src/custom_llm/config.ts b/packages/core/src/custom_llm/config.ts new file mode 100644 index 00000000..d1a3a28f --- /dev/null +++ b/packages/core/src/custom_llm/config.ts @@ -0,0 +1,41 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +export const ATLASCLOUD_OPENAI_BASE_URL = 'https://api.atlascloud.ai/v1'; +export const ATLASCLOUD_DEFAULT_MODEL = 'qwen/qwen3.5-flash'; +export const ATLASCLOUD_PROVIDER_NAME = 'atlascloud'; + +export interface CustomLLMRuntimeConfig { + apiKey: string; + baseURL: string; + modelName: string; + provider: string; + temperature: number; + maxTokens: number; + topP: number; +} + +export function resolveCustomLLMRuntimeConfig( + env: NodeJS.ProcessEnv = process.env, +): CustomLLMRuntimeConfig { + const hasAtlasCloudKey = Boolean(env.ATLASCLOUD_API_KEY); + + return { + apiKey: env.CUSTOM_LLM_API_KEY || env.ATLASCLOUD_API_KEY || '', + baseURL: + env.CUSTOM_LLM_ENDPOINT || + (hasAtlasCloudKey ? ATLASCLOUD_OPENAI_BASE_URL : ''), + modelName: + env.CUSTOM_LLM_MODEL_NAME || + (hasAtlasCloudKey ? ATLASCLOUD_DEFAULT_MODEL : ''), + provider: + env.CUSTOM_LLM_PROVIDER || + (hasAtlasCloudKey ? ATLASCLOUD_PROVIDER_NAME : ''), + temperature: Number(env.CUSTOM_LLM_TEMPERATURE || 0), + maxTokens: Number(env.CUSTOM_LLM_MAX_TOKENS || 8192), + topP: Number(env.CUSTOM_LLM_TOP_P || 1), + }; +} diff --git a/packages/core/src/custom_llm/index.ts b/packages/core/src/custom_llm/index.ts index f36a5186..6f22beb0 100644 --- a/packages/core/src/custom_llm/index.ts +++ b/packages/core/src/custom_llm/index.ts @@ -17,26 +17,22 @@ import { ContentGenerator } from '../core/contentGenerator.js'; import { CustomLLMContentGeneratorConfig, ToolCallMap } from './types.js'; import { extractToolFunctions } from './util.js'; import { ModelConverter } from './converter.js'; +import { resolveCustomLLMRuntimeConfig } from './config.js'; export class CustomLLMContentGenerator implements ContentGenerator { private model: OpenAI; - private apiKey: string = process.env.CUSTOM_LLM_API_KEY || ''; - private baseURL: string = process.env.CUSTOM_LLM_ENDPOINT || ''; - private modelName: string = process.env.CUSTOM_LLM_MODEL_NAME || ''; - private temperature: number = Number(process.env.CUSTOM_LLM_TEMPERATURE || 0); - private maxTokens: number = Number(process.env.CUSTOM_LLM_MAX_TOKENS || 8192); - private topP: number = Number(process.env.CUSTOM_LLM_TOP_P || 1); + private runtimeConfig = resolveCustomLLMRuntimeConfig(); private config: CustomLLMContentGeneratorConfig = { - model: this.modelName, - temperature: this.temperature, - max_tokens: this.maxTokens, - top_p: this.topP, + model: this.runtimeConfig.modelName, + temperature: this.runtimeConfig.temperature, + max_tokens: this.runtimeConfig.maxTokens, + top_p: this.runtimeConfig.topP, }; constructor() { this.model = new OpenAI({ - apiKey: this.apiKey, - baseURL: this.baseURL, + apiKey: this.runtimeConfig.apiKey, + baseURL: this.runtimeConfig.baseURL, }); } diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index aff37f50..5cb6df85 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -18,6 +18,7 @@ export * from './core/turn.js'; export * from './core/geminiRequest.js'; export * from './core/coreToolScheduler.js'; export * from './core/nonInteractiveToolExecutor.js'; +export * from './custom_llm/config.js'; export * from './code_assist/codeAssist.js'; export * from './code_assist/oauth2.js';