Code Assist cleanup and docs (#993)

This commit is contained in:
Tommaso Sciortino 2025-06-12 18:00:17 -07:00 committed by GitHub
parent 9a11567f73
commit 431ee839a0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 122 additions and 80 deletions

View File

@ -30,7 +30,24 @@ Or ask it to perform a task using its tools:
The Gemini CLI requires you to authenticate with Google's AI services. You'll need to configure **one** of the following authentication methods:
1. **Gemini API key:**
1. **Gemini Code Assist:**
- To enable this mode you only need set the GEMINI_CODE_ASSIST environment variable to true.
- Enterprise users must also provide a GOOGLE_CLOUD_PROJECT environment variable specifying their Google Cloud project.
In the following methods, replace `GOOGLE_CLOUD_PROJECT` with the relevant values for your project:
- You can temporarily set the environment variable in your current shell session using the following command:
```bash
export GEMINI_CODE_ASSIST="true"
export GOOGLE_CLOUD_PROJECT="YOUR_PROJECT_ID" // Enterprise users only.
```
- For repeated use, you can add the environment variable to your `.env` file (located in the project directory or user home directory) or your shell's configuration file (like `~/.bashrc`, `~/.zshrc`, or `~/.profile`). For example, the following command adds the environment variable to a `~/.bashrc` file:
```bash
echo 'export GEMINI_CODE_ASSIST="true"' >> ~/.bashrc
echo 'export GOOGLE_CLOUD_PROJECT="YOUR_PROJECT_ID"' >> ~/.bashrc # Enterprise users only.
source ~/.bashrc
```
2. **Gemini API key:**
- Obtain your API key from Google AI Studio: [https://aistudio.google.com/app/apikey](https://aistudio.google.com/app/apikey)
- Set the `GEMINI_API_KEY` environment variable. In the following methods, replace `YOUR_GEMINI_API_KEY` with the API key you obtained from Google AI Studio:
@ -40,11 +57,11 @@ The Gemini CLI requires you to authenticate with Google's AI services. You'll ne
```
- For repeated use, you can add the environment variable to your `.env` file (located in the project directory or user home directory) or your shell's configuration file (like `~/.bashrc`, `~/.zshrc`, or `~/.profile`). For example, the following command adds the environment variable to a `~/.bashrc` file:
```bash
echo 'export GEMINI_API_KEY="YOUR_GEMINI_API_KEY"' >> ~/.bashrc # Or your preferred shell config file
source ~/.bashrc # Reload the config
echo 'export GEMINI_API_KEY="YOUR_GEMINI_API_KEY"' >> ~/.bashrc
source ~/.bashrc
```
2. **Google API Key (Vertex AI Express Mode):**
3. **Google API Key (Vertex AI Express Mode):**
- You can use a general Google Cloud API key if it has been enabled for the Gemini API or Vertex AI.
- Set the `GOOGLE_API_KEY` and `GOOGLE_GENAI_USE_VERTEXAI` environment variables. In the following methods, replace `YOUR_GEMINI_API_KEY` with your Google Cloud API key:
@ -55,12 +72,12 @@ The Gemini CLI requires you to authenticate with Google's AI services. You'll ne
```
- For repeated use, you can add the environment variables to your `.env` file (located in the project directory or user home directory) or your shell's configuration file (like `~/.bashrc`, `~/.zshrc`, or `~/.profile`). For example, the following commands adds the environment variables to a `~/.bashrc` file:
```bash
echo 'export GEMINI_API_KEY="YOUR_GEMINI_API_KEY"' >> ~/.bashrc # Or your preferred shell config file
echo 'export GOOGLE_GENAI_USE_VERTEXAI=true' >> ~/.bashrc # Or your preferred shell config file
source ~/.bashrc # Reload the config
echo 'export GEMINI_API_KEY="YOUR_GEMINI_API_KEY"' >> ~/.bashrc
echo 'export GOOGLE_GENAI_USE_VERTEXAI=true' >> ~/.bashrc
source ~/.bashrc
```
3. **Vertex AI (Project and Location):**
4. **Vertex AI (Project and Location):**
- Ensure you have a Google Cloud project and have enabled the Vertex AI API.
- Set up Application Default Credentials (ADC), using the following command:
```bash
@ -76,10 +93,10 @@ The Gemini CLI requires you to authenticate with Google's AI services. You'll ne
```
- For repeated use, you can add the environment variables to your `.env` file (located in the project directory or user home directory) or your shell's configuration file (like `~/.bashrc`, `~/.zshrc`, or `~/.profile`). For example, the following commands adds the environment variables to a `~/.bashrc` file:
```bash
echo 'export GOOGLE_CLOUD_PROJECT="YOUR_PROJECT_ID"' >> ~/.bashrc # Or your preferred shell config file
echo 'export GOOGLE_CLOUD_LOCATION="YOUR_PROJECT_LOCATION"' >> ~/.bashrc # Or your preferred shell config file
echo 'export GOOGLE_GENAI_USE_VERTEXAI=true' >> ~/.bashrc # Or your preferred shell config file
source ~/.bashrc # Reload the config
echo 'export GOOGLE_CLOUD_PROJECT="YOUR_PROJECT_ID"' >> ~/.bashrc
echo 'export GOOGLE_CLOUD_LOCATION="YOUR_PROJECT_LOCATION"' >> ~/.bashrc
echo 'export GOOGLE_GENAI_USE_VERTEXAI=true' >> ~/.bashrc
source ~/.bashrc
```
### Next Steps

View File

@ -215,6 +215,9 @@ The CLI automatically loads environment variables from an `.env` file. The loadi
- Set to any value to disable all color output in the CLI.
- **`CLI_TITLE`**:
- Set to a string to customize the title of the CLI.
- **`CODE_ASSIST_ENDPOINT`**:
- Specifies the endpoint for the code assist server.
- This is useful for development and testing.
## 3. Command-Line Arguments

View File

@ -234,6 +234,7 @@ async function createContentGeneratorConfig(
const googleCloudProject = process.env.GOOGLE_CLOUD_PROJECT;
const googleCloudLocation = process.env.GOOGLE_CLOUD_LOCATION;
const hasCodeAssist = process.env.GEMINI_CODE_ASSIST === 'true';
const hasGeminiApiKey = !!geminiApiKey;
const hasGoogleApiKey = !!googleApiKey;
const hasVertexProjectLocationConfig =
@ -244,12 +245,18 @@ async function createContentGeneratorConfig(
'Both GEMINI_API_KEY and GOOGLE_API_KEY are set. Using GOOGLE_API_KEY.',
);
}
if (!hasGeminiApiKey && !hasGoogleApiKey && !hasVertexProjectLocationConfig) {
if (
!hasCodeAssist &&
!hasGeminiApiKey &&
!hasGoogleApiKey &&
!hasVertexProjectLocationConfig
) {
logger.error(
'No valid API authentication configuration found. Please set ONE of the following combinations in your environment variables or .env file:\n' +
'1. GEMINI_API_KEY (for Gemini API access).\n' +
'2. GOOGLE_API_KEY (for Gemini API or Vertex AI Express Mode access).\n' +
'3. GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION (for Vertex AI access).\n\n' +
'1. GEMINI_CODE_ASSIST=true (for Code Assist access).\n' +
'2. GEMINI_API_KEY (for Gemini API access).\n' +
'3. GOOGLE_API_KEY (for Gemini API or Vertex AI Express Mode access).\n' +
'4. GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION (for Vertex AI access).\n\n' +
'For Gemini API keys, visit: https://ai.google.dev/gemini-api/docs/api-key\n' +
'For Vertex AI authentication, visit: https://cloud.google.com/vertex-ai/docs/start/authentication\n' +
'The GOOGLE_GENAI_USE_VERTEXAI environment variable can also be set to true/false to influence service selection when ambiguity exists.',
@ -261,7 +268,7 @@ async function createContentGeneratorConfig(
model: argv.model || DEFAULT_GEMINI_MODEL,
apiKey: googleApiKey || geminiApiKey || '',
vertexai: hasGeminiApiKey ? false : undefined,
codeAssist: !!process.env.GEMINI_CODE_ASSIST,
codeAssist: hasCodeAssist,
};
if (config.apiKey) {

View File

@ -7,7 +7,7 @@
import { ContentGenerator } from '../core/contentGenerator.js';
import { getOauthClient } from './oauth2.js';
import { setupUser } from './setup.js';
import { CcpaServer } from './ccpaServer.js';
import { CodeAssistServer } from './server.js';
export async function createCodeAssistContentGenerator(): Promise<ContentGenerator> {
const oauth2Client = await getOauthClient();
@ -15,5 +15,5 @@ export async function createCodeAssistContentGenerator(): Promise<ContentGenerat
oauth2Client,
process.env.GOOGLE_CLOUD_PROJECT,
);
return new CcpaServer(oauth2Client, projectId);
return new CodeAssistServer(oauth2Client, projectId);
}

View File

@ -5,7 +5,11 @@
*/
import { describe, it, expect } from 'vitest';
import { toCcpaRequest, fromCcpaResponse, CcpaResponse } from './converter.js';
import {
toCodeAssistRequest,
fromCodeAsistResponse,
CodeAssistResponse,
} from './converter.js';
import {
GenerateContentParameters,
GenerateContentResponse,
@ -14,14 +18,14 @@ import {
} from '@google/genai';
describe('converter', () => {
describe('toCcpaRequest', () => {
describe('toCodeAssistRequest', () => {
it('should convert a simple request with project', () => {
const genaiReq: GenerateContentParameters = {
model: 'gemini-pro',
contents: [{ role: 'user', parts: [{ text: 'Hello' }] }],
};
const ccpaReq = toCcpaRequest(genaiReq, 'my-project');
expect(ccpaReq).toEqual({
const codeAssistReq = toCodeAssistRequest(genaiReq, 'my-project');
expect(codeAssistReq).toEqual({
model: 'gemini-pro',
project: 'my-project',
request: {
@ -42,8 +46,8 @@ describe('converter', () => {
model: 'gemini-pro',
contents: [{ role: 'user', parts: [{ text: 'Hello' }] }],
};
const ccpaReq = toCcpaRequest(genaiReq);
expect(ccpaReq).toEqual({
const codeAssistReq = toCodeAssistRequest(genaiReq);
expect(codeAssistReq).toEqual({
model: 'gemini-pro',
project: undefined,
request: {
@ -64,8 +68,8 @@ describe('converter', () => {
model: 'gemini-pro',
contents: 'Hello',
};
const ccpaReq = toCcpaRequest(genaiReq);
expect(ccpaReq.request.contents).toEqual([
const codeAssistReq = toCodeAssistRequest(genaiReq);
expect(codeAssistReq.request.contents).toEqual([
{ role: 'user', parts: [{ text: 'Hello' }] },
]);
});
@ -75,8 +79,8 @@ describe('converter', () => {
model: 'gemini-pro',
contents: [{ text: 'Hello' }, { text: 'World' }],
};
const ccpaReq = toCcpaRequest(genaiReq);
expect(ccpaReq.request.contents).toEqual([
const codeAssistReq = toCodeAssistRequest(genaiReq);
expect(codeAssistReq.request.contents).toEqual([
{ role: 'user', parts: [{ text: 'Hello' }] },
{ role: 'user', parts: [{ text: 'World' }] },
]);
@ -90,8 +94,8 @@ describe('converter', () => {
systemInstruction: 'You are a helpful assistant.',
},
};
const ccpaReq = toCcpaRequest(genaiReq);
expect(ccpaReq.request.systemInstruction).toEqual({
const codeAssistReq = toCodeAssistRequest(genaiReq);
expect(codeAssistReq.request.systemInstruction).toEqual({
role: 'user',
parts: [{ text: 'You are a helpful assistant.' }],
});
@ -106,8 +110,8 @@ describe('converter', () => {
topK: 40,
},
};
const ccpaReq = toCcpaRequest(genaiReq);
expect(ccpaReq.request.generationConfig).toEqual({
const codeAssistReq = toCodeAssistRequest(genaiReq);
expect(codeAssistReq.request.generationConfig).toEqual({
temperature: 0.8,
topK: 40,
});
@ -132,8 +136,8 @@ describe('converter', () => {
responseMimeType: 'application/json',
},
};
const ccpaReq = toCcpaRequest(genaiReq);
expect(ccpaReq.request.generationConfig).toEqual({
const codeAssistReq = toCodeAssistRequest(genaiReq);
expect(codeAssistReq.request.generationConfig).toEqual({
temperature: 0.1,
topP: 0.2,
topK: 3,
@ -150,9 +154,9 @@ describe('converter', () => {
});
});
describe('fromCcpaResponse', () => {
describe('fromCodeAssistResponse', () => {
it('should convert a simple response', () => {
const ccpaRes: CcpaResponse = {
const codeAssistRes: CodeAssistResponse = {
response: {
candidates: [
{
@ -167,13 +171,13 @@ describe('converter', () => {
],
},
};
const genaiRes = fromCcpaResponse(ccpaRes);
const genaiRes = fromCodeAsistResponse(codeAssistRes);
expect(genaiRes).toBeInstanceOf(GenerateContentResponse);
expect(genaiRes.candidates).toEqual(ccpaRes.response.candidates);
expect(genaiRes.candidates).toEqual(codeAssistRes.response.candidates);
});
it('should handle prompt feedback and usage metadata', () => {
const ccpaRes: CcpaResponse = {
const codeAssistRes: CodeAssistResponse = {
response: {
candidates: [],
promptFeedback: {
@ -187,13 +191,17 @@ describe('converter', () => {
},
},
};
const genaiRes = fromCcpaResponse(ccpaRes);
expect(genaiRes.promptFeedback).toEqual(ccpaRes.response.promptFeedback);
expect(genaiRes.usageMetadata).toEqual(ccpaRes.response.usageMetadata);
const genaiRes = fromCodeAsistResponse(codeAssistRes);
expect(genaiRes.promptFeedback).toEqual(
codeAssistRes.response.promptFeedback,
);
expect(genaiRes.usageMetadata).toEqual(
codeAssistRes.response.usageMetadata,
);
});
it('should handle automatic function calling history', () => {
const ccpaRes: CcpaResponse = {
const codeAssistRes: CodeAssistResponse = {
response: {
candidates: [],
automaticFunctionCallingHistory: [
@ -213,9 +221,9 @@ describe('converter', () => {
],
},
};
const genaiRes = fromCcpaResponse(ccpaRes);
const genaiRes = fromCodeAsistResponse(codeAssistRes);
expect(genaiRes.automaticFunctionCallingHistory).toEqual(
ccpaRes.response.automaticFunctionCallingHistory,
codeAssistRes.response.automaticFunctionCallingHistory,
);
});
});

View File

@ -27,13 +27,13 @@ import {
ToolConfig,
} from '@google/genai';
export interface CcpaRequest {
export interface CodeAssistRequest {
model: string;
project?: string;
request: CcpaGenerateContentRequest;
request: CodeAssistGenerateContentRequest;
}
interface CcpaGenerateContentRequest {
interface CodeAssistGenerateContentRequest {
contents: Content[];
systemInstruction?: Content;
cachedContent?: string;
@ -41,10 +41,10 @@ interface CcpaGenerateContentRequest {
toolConfig?: ToolConfig;
labels?: Record<string, string>;
safetySettings?: SafetySetting[];
generationConfig?: CcpaGenerationConfig;
generationConfig?: CodeAssistGenerationConfig;
}
interface CcpaGenerationConfig {
interface CodeAssistGenerationConfig {
temperature?: number;
topP?: number;
topK?: number;
@ -67,7 +67,7 @@ interface CcpaGenerationConfig {
thinkingConfig?: ThinkingConfig;
}
export interface CcpaResponse {
export interface CodeAssistResponse {
response: VertexResponse;
}
@ -78,18 +78,20 @@ interface VertexResponse {
usageMetadata?: GenerateContentResponseUsageMetadata;
}
export function toCcpaRequest(
export function toCodeAssistRequest(
req: GenerateContentParameters,
project?: string,
): CcpaRequest {
): CodeAssistRequest {
return {
model: req.model,
project,
request: toCcpaGenerateContentRequest(req),
request: toCodeAssistGenerateContentRequest(req),
};
}
export function fromCcpaResponse(res: CcpaResponse): GenerateContentResponse {
export function fromCodeAsistResponse(
res: CodeAssistResponse,
): GenerateContentResponse {
const inres = res.response;
const out = new GenerateContentResponse();
out.candidates = inres.candidates;
@ -99,9 +101,9 @@ export function fromCcpaResponse(res: CcpaResponse): GenerateContentResponse {
return out;
}
function toCcpaGenerateContentRequest(
function toCodeAssistGenerateContentRequest(
req: GenerateContentParameters,
): CcpaGenerateContentRequest {
): CodeAssistGenerateContentRequest {
return {
contents: toContents(req.contents),
systemInstruction: maybeToContent(req.config?.systemInstruction),
@ -110,7 +112,7 @@ function toCcpaGenerateContentRequest(
toolConfig: req.config?.toolConfig,
labels: req.config?.labels,
safetySettings: req.config?.safetySettings,
generationConfig: toCcpaGenerationConfig(req.config),
generationConfig: toCodeAssistGenerationConfig(req.config),
};
}
@ -168,9 +170,9 @@ function toPart(part: PartUnion): Part {
return part;
}
function toCcpaGenerationConfig(
function toCodeAssistGenerationConfig(
config?: GenerateContentConfig,
): CcpaGenerationConfig | undefined {
): CodeAssistGenerationConfig | undefined {
if (!config) {
return undefined;
}

View File

@ -21,15 +21,20 @@ import {
} from '@google/genai';
import * as readline from 'readline';
import { ContentGenerator } from '../core/contentGenerator.js';
import { CcpaResponse, toCcpaRequest, fromCcpaResponse } from './converter.js';
import {
CodeAssistResponse,
toCodeAssistRequest,
fromCodeAsistResponse,
} from './converter.js';
import { PassThrough } from 'node:stream';
// TODO: Use production endpoint once it supports our methods.
export const CCPA_ENDPOINT =
export const CODE_ASSIST_ENDPOINT =
process.env.CODE_ASSIST_ENDPOINT ??
'https://staging-cloudcode-pa.sandbox.googleapis.com';
export const CCPA_API_VERSION = 'v1internal';
export const CODE_ASSIST_API_VERSION = 'v1internal';
export class CcpaServer implements ContentGenerator {
export class CodeAssistServer implements ContentGenerator {
constructor(
readonly auth: OAuth2Client,
readonly projectId?: string,
@ -38,13 +43,13 @@ export class CcpaServer implements ContentGenerator {
async generateContentStream(
req: GenerateContentParameters,
): Promise<AsyncGenerator<GenerateContentResponse>> {
const resps = await this.streamEndpoint<CcpaResponse>(
const resps = await this.streamEndpoint<CodeAssistResponse>(
'streamGenerateContent',
toCcpaRequest(req, this.projectId),
toCodeAssistRequest(req, this.projectId),
);
return (async function* (): AsyncGenerator<GenerateContentResponse> {
for await (const resp of resps) {
yield fromCcpaResponse(resp);
yield fromCodeAsistResponse(resp);
}
})();
}
@ -52,11 +57,11 @@ export class CcpaServer implements ContentGenerator {
async generateContent(
req: GenerateContentParameters,
): Promise<GenerateContentResponse> {
const resp = await this.callEndpoint<CcpaResponse>(
const resp = await this.callEndpoint<CodeAssistResponse>(
'generateContent',
toCcpaRequest(req, this.projectId),
toCodeAssistRequest(req, this.projectId),
);
return fromCcpaResponse(resp);
return fromCodeAsistResponse(resp);
}
async onboardUser(
@ -89,7 +94,7 @@ export class CcpaServer implements ContentGenerator {
async callEndpoint<T>(method: string, req: object): Promise<T> {
const res = await this.auth.request({
url: `${CCPA_ENDPOINT}/${CCPA_API_VERSION}:${method}`,
url: `${CODE_ASSIST_ENDPOINT}/${CODE_ASSIST_API_VERSION}:${method}`,
method: 'POST',
headers: {
'Content-Type': 'application/json',
@ -106,7 +111,7 @@ export class CcpaServer implements ContentGenerator {
req: object,
): Promise<AsyncGenerator<T>> {
const res = await this.auth.request({
url: `${CCPA_ENDPOINT}/${CCPA_API_VERSION}:${method}`,
url: `${CODE_ASSIST_ENDPOINT}/${CODE_ASSIST_API_VERSION}:${method}`,
method: 'POST',
params: {
alt: 'sse',

View File

@ -5,7 +5,7 @@
*/
import { ClientMetadata, OnboardUserRequest } from './types.js';
import { CcpaServer } from './ccpaServer.js';
import { CodeAssistServer } from './server.js';
import { OAuth2Client } from 'google-auth-library';
import { GaxiosError } from 'gaxios';
import { clearCachedCredentials } from './oauth2.js';
@ -19,7 +19,7 @@ export async function setupUser(
oAuth2Client: OAuth2Client,
projectId?: string,
): Promise<string> {
const ccpaServer: CcpaServer = new CcpaServer(oAuth2Client, projectId);
const caServer = new CodeAssistServer(oAuth2Client, projectId);
const clientMetadata: ClientMetadata = {
ideType: 'IDE_UNSPECIFIED',
platform: 'PLATFORM_UNSPECIFIED',
@ -30,7 +30,7 @@ export async function setupUser(
}
// TODO: Support Free Tier user without projectId.
const loadRes = await ccpaServer.loadCodeAssist({
const loadRes = await caServer.loadCodeAssist({
cloudaicompanionProject: process.env.GOOGLE_CLOUD_PROJECT,
metadata: clientMetadata,
});
@ -42,10 +42,10 @@ export async function setupUser(
};
try {
// Poll onboardUser until long running operation is complete.
let lroRes = await ccpaServer.onboardUser(onboardReq);
let lroRes = await caServer.onboardUser(onboardReq);
while (!lroRes.done) {
await new Promise((f) => setTimeout(f, 5000));
lroRes = await ccpaServer.onboardUser(onboardReq);
lroRes = await caServer.onboardUser(onboardReq);
}
return lroRes.response?.cloudaicompanionProject?.id || '';

View File

@ -53,7 +53,7 @@ export interface LoadCodeAssistResponse {
}
/**
* GeminiUserTier reflects the structure received from the CCPA when calling LoadCodeAssist.
* GeminiUserTier reflects the structure received from the CodeAssist when calling LoadCodeAssist.
*/
export interface GeminiUserTier {
id: UserTierId;
@ -109,7 +109,7 @@ export enum UserTierId {
}
/**
* PrivacyNotice reflects the structure received from the CCPA in regards to a tier
* PrivacyNotice reflects the structure received from the CodeAssist in regards to a tier
* privacy notice.
*/
export interface PrivacyNotice {