Code Assist cleanup and docs (#993)

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');
-      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,
       );
     });
   });

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 {