Code Assist cleanup and docs (#993)

README.md

@@ -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
+        echo 'export GEMINI_API_KEY="YOUR_GEMINI_API_KEY"' >> ~/.bashrc
-        source ~/.bashrc # Reload the config
+        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 GEMINI_API_KEY="YOUR_GEMINI_API_KEY"' >> ~/.bashrc
-        echo 'export GOOGLE_GENAI_USE_VERTEXAI=true' >> ~/.bashrc # Or your preferred shell config file
+        echo 'export GOOGLE_GENAI_USE_VERTEXAI=true' >> ~/.bashrc
-        source ~/.bashrc # Reload the config
+        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_PROJECT="YOUR_PROJECT_ID"' >> ~/.bashrc
-        echo 'export GOOGLE_CLOUD_LOCATION="YOUR_PROJECT_LOCATION"' >> ~/.bashrc # Or your preferred shell config file
+        echo 'export GOOGLE_CLOUD_LOCATION="YOUR_PROJECT_LOCATION"' >> ~/.bashrc
-        echo 'export GOOGLE_GENAI_USE_VERTEXAI=true' >> ~/.bashrc # Or your preferred shell config file
+        echo 'export GOOGLE_GENAI_USE_VERTEXAI=true' >> ~/.bashrc
-        source ~/.bashrc # Reload the config
+        source ~/.bashrc
         ```
 
 ### Next Steps

docs/cli/configuration.md

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

packages/cli/src/config/config.ts

@@ -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' +
+        '1. GEMINI_CODE_ASSIST=true (for Code Assist access).\n' +
-        '2. GOOGLE_API_KEY (for Gemini API or Vertex AI Express Mode access).\n' +
+        '2. GEMINI_API_KEY (for Gemini API access).\n' +
-        '3. GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION (for Vertex AI access).\n\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) {

packages/core/src/code_assist/codeAssist.ts

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

packages/core/src/code_assist/converter.test.ts

@@ -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');
+      const codeAssistReq = toCodeAssistRequest(genaiReq, 'my-project');
-      expect(ccpaReq).toEqual({
+      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);
+      const codeAssistReq = toCodeAssistRequest(genaiReq);
-      expect(ccpaReq).toEqual({
+      expect(codeAssistReq).toEqual({
         model: 'gemini-pro',
         project: undefined,
         request: {
@@ -64,8 +68,8 @@ describe('converter', () => {
         model: 'gemini-pro',
         contents: 'Hello',
       };
-      const ccpaReq = toCcpaRequest(genaiReq);
+      const codeAssistReq = toCodeAssistRequest(genaiReq);
-      expect(ccpaReq.request.contents).toEqual([
+      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);
+      const codeAssistReq = toCodeAssistRequest(genaiReq);
-      expect(ccpaReq.request.contents).toEqual([
+      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);
+      const codeAssistReq = toCodeAssistRequest(genaiReq);
-      expect(ccpaReq.request.systemInstruction).toEqual({
+      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);
+      const codeAssistReq = toCodeAssistRequest(genaiReq);
-      expect(ccpaReq.request.generationConfig).toEqual({
+      expect(codeAssistReq.request.generationConfig).toEqual({
         temperature: 0.8,
         topK: 40,
       });
@@ -132,8 +136,8 @@ describe('converter', () => {
           responseMimeType: 'application/json',
         },
       };
-      const ccpaReq = toCcpaRequest(genaiReq);
+      const codeAssistReq = toCodeAssistRequest(genaiReq);
-      expect(ccpaReq.request.generationConfig).toEqual({
+      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);
+      const genaiRes = fromCodeAsistResponse(codeAssistRes);
-      expect(genaiRes.promptFeedback).toEqual(ccpaRes.response.promptFeedback);
+      expect(genaiRes.promptFeedback).toEqual(
-      expect(genaiRes.usageMetadata).toEqual(ccpaRes.response.usageMetadata);
+        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,
       );
     });
   });

packages/core/src/code_assist/converter.ts

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

packages/core/src/code_assist/server.ts

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

packages/core/src/code_assist/setup.ts

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

packages/core/src/code_assist/types.ts

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