From 431ee839a0ad3d1f97e0f437bacbab5108481b90 Mon Sep 17 00:00:00 2001 From: Tommaso Sciortino Date: Thu, 12 Jun 2025 18:00:17 -0700 Subject: Code Assist cleanup and docs (#993) --- packages/core/src/code_assist/ccpaServer.ts | 142 ----------------------- packages/core/src/code_assist/codeAssist.ts | 4 +- packages/core/src/code_assist/converter.test.ts | 62 +++++----- packages/core/src/code_assist/converter.ts | 32 +++--- packages/core/src/code_assist/server.ts | 147 ++++++++++++++++++++++++ packages/core/src/code_assist/setup.ts | 10 +- packages/core/src/code_assist/types.ts | 4 +- 7 files changed, 208 insertions(+), 193 deletions(-) delete mode 100644 packages/core/src/code_assist/ccpaServer.ts create mode 100644 packages/core/src/code_assist/server.ts (limited to 'packages/core') diff --git a/packages/core/src/code_assist/ccpaServer.ts b/packages/core/src/code_assist/ccpaServer.ts deleted file mode 100644 index 3ef8b084..00000000 --- a/packages/core/src/code_assist/ccpaServer.ts +++ /dev/null @@ -1,142 +0,0 @@ -/** - * @license - * Copyright 2025 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { OAuth2Client } from 'google-auth-library'; -import { - LoadCodeAssistResponse, - LoadCodeAssistRequest, - OnboardUserRequest, - LongrunningOperationResponse, -} from './types.js'; -import { - GenerateContentResponse, - GenerateContentParameters, - CountTokensParameters, - EmbedContentResponse, - CountTokensResponse, - EmbedContentParameters, -} from '@google/genai'; -import * as readline from 'readline'; -import { ContentGenerator } from '../core/contentGenerator.js'; -import { CcpaResponse, toCcpaRequest, fromCcpaResponse } from './converter.js'; -import { PassThrough } from 'node:stream'; - -// TODO: Use production endpoint once it supports our methods. -export const CCPA_ENDPOINT = - 'https://staging-cloudcode-pa.sandbox.googleapis.com'; -export const CCPA_API_VERSION = 'v1internal'; - -export class CcpaServer implements ContentGenerator { - constructor( - readonly auth: OAuth2Client, - readonly projectId?: string, - ) {} - - async generateContentStream( - req: GenerateContentParameters, - ): Promise> { - const resps = await this.streamEndpoint( - 'streamGenerateContent', - toCcpaRequest(req, this.projectId), - ); - return (async function* (): AsyncGenerator { - for await (const resp of resps) { - yield fromCcpaResponse(resp); - } - })(); - } - - async generateContent( - req: GenerateContentParameters, - ): Promise { - const resp = await this.callEndpoint( - 'generateContent', - toCcpaRequest(req, this.projectId), - ); - return fromCcpaResponse(resp); - } - - async onboardUser( - req: OnboardUserRequest, - ): Promise { - return await this.callEndpoint( - 'onboardUser', - req, - ); - } - - async loadCodeAssist( - req: LoadCodeAssistRequest, - ): Promise { - return await this.callEndpoint( - 'loadCodeAssist', - req, - ); - } - - async countTokens(_req: CountTokensParameters): Promise { - return { totalTokens: 0 }; - } - - async embedContent( - _req: EmbedContentParameters, - ): Promise { - throw Error(); - } - - async callEndpoint(method: string, req: object): Promise { - const res = await this.auth.request({ - url: `${CCPA_ENDPOINT}/${CCPA_API_VERSION}:${method}`, - method: 'POST', - headers: { - 'Content-Type': 'application/json', - 'X-Goog-User-Project': this.projectId || '', - }, - responseType: 'json', - body: JSON.stringify(req), - }); - return res.data as T; - } - - async streamEndpoint( - method: string, - req: object, - ): Promise> { - const res = await this.auth.request({ - url: `${CCPA_ENDPOINT}/${CCPA_API_VERSION}:${method}`, - method: 'POST', - params: { - alt: 'sse', - }, - headers: { 'Content-Type': 'application/json' }, - responseType: 'stream', - body: JSON.stringify(req), - }); - - return (async function* (): AsyncGenerator { - const rl = readline.createInterface({ - input: res.data as PassThrough, - crlfDelay: Infinity, // Recognizes '\r\n' and '\n' as line breaks - }); - - let bufferedLines: string[] = []; - for await (const line of rl) { - // blank lines are used to separate JSON objects in the stream - if (line === '') { - if (bufferedLines.length === 0) { - continue; // no data to yield - } - yield JSON.parse(bufferedLines.join('\n')) as T; - bufferedLines = []; // Reset the buffer after yielding - } else if (line.startsWith('data: ')) { - bufferedLines.push(line.slice(6).trim()); - } else { - throw new Error(`Unexpected line format in response: ${line}`); - } - } - })(); - } -} diff --git a/packages/core/src/code_assist/codeAssist.ts b/packages/core/src/code_assist/codeAssist.ts index 6467b416..5922cb41 100644 --- a/packages/core/src/code_assist/codeAssist.ts +++ b/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 { const oauth2Client = await getOauthClient(); @@ -15,5 +15,5 @@ export async function createCodeAssistContentGenerator(): Promise { - 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, ); }); }); diff --git a/packages/core/src/code_assist/converter.ts b/packages/core/src/code_assist/converter.ts index c7b0e7c7..495cbfae 100644 --- a/packages/core/src/code_assist/converter.ts +++ b/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; 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; } diff --git a/packages/core/src/code_assist/server.ts b/packages/core/src/code_assist/server.ts new file mode 100644 index 00000000..680b66ac --- /dev/null +++ b/packages/core/src/code_assist/server.ts @@ -0,0 +1,147 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { OAuth2Client } from 'google-auth-library'; +import { + LoadCodeAssistResponse, + LoadCodeAssistRequest, + OnboardUserRequest, + LongrunningOperationResponse, +} from './types.js'; +import { + GenerateContentResponse, + GenerateContentParameters, + CountTokensParameters, + EmbedContentResponse, + CountTokensResponse, + EmbedContentParameters, +} from '@google/genai'; +import * as readline from 'readline'; +import { ContentGenerator } from '../core/contentGenerator.js'; +import { + CodeAssistResponse, + toCodeAssistRequest, + fromCodeAsistResponse, +} from './converter.js'; +import { PassThrough } from 'node:stream'; + +// TODO: Use production endpoint once it supports our methods. +export const CODE_ASSIST_ENDPOINT = + process.env.CODE_ASSIST_ENDPOINT ?? + 'https://staging-cloudcode-pa.sandbox.googleapis.com'; +export const CODE_ASSIST_API_VERSION = 'v1internal'; + +export class CodeAssistServer implements ContentGenerator { + constructor( + readonly auth: OAuth2Client, + readonly projectId?: string, + ) {} + + async generateContentStream( + req: GenerateContentParameters, + ): Promise> { + const resps = await this.streamEndpoint( + 'streamGenerateContent', + toCodeAssistRequest(req, this.projectId), + ); + return (async function* (): AsyncGenerator { + for await (const resp of resps) { + yield fromCodeAsistResponse(resp); + } + })(); + } + + async generateContent( + req: GenerateContentParameters, + ): Promise { + const resp = await this.callEndpoint( + 'generateContent', + toCodeAssistRequest(req, this.projectId), + ); + return fromCodeAsistResponse(resp); + } + + async onboardUser( + req: OnboardUserRequest, + ): Promise { + return await this.callEndpoint( + 'onboardUser', + req, + ); + } + + async loadCodeAssist( + req: LoadCodeAssistRequest, + ): Promise { + return await this.callEndpoint( + 'loadCodeAssist', + req, + ); + } + + async countTokens(_req: CountTokensParameters): Promise { + return { totalTokens: 0 }; + } + + async embedContent( + _req: EmbedContentParameters, + ): Promise { + throw Error(); + } + + async callEndpoint(method: string, req: object): Promise { + const res = await this.auth.request({ + url: `${CODE_ASSIST_ENDPOINT}/${CODE_ASSIST_API_VERSION}:${method}`, + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Goog-User-Project': this.projectId || '', + }, + responseType: 'json', + body: JSON.stringify(req), + }); + return res.data as T; + } + + async streamEndpoint( + method: string, + req: object, + ): Promise> { + const res = await this.auth.request({ + url: `${CODE_ASSIST_ENDPOINT}/${CODE_ASSIST_API_VERSION}:${method}`, + method: 'POST', + params: { + alt: 'sse', + }, + headers: { 'Content-Type': 'application/json' }, + responseType: 'stream', + body: JSON.stringify(req), + }); + + return (async function* (): AsyncGenerator { + const rl = readline.createInterface({ + input: res.data as PassThrough, + crlfDelay: Infinity, // Recognizes '\r\n' and '\n' as line breaks + }); + + let bufferedLines: string[] = []; + for await (const line of rl) { + // blank lines are used to separate JSON objects in the stream + if (line === '') { + if (bufferedLines.length === 0) { + continue; // no data to yield + } + yield JSON.parse(bufferedLines.join('\n')) as T; + bufferedLines = []; // Reset the buffer after yielding + } else if (line.startsWith('data: ')) { + bufferedLines.push(line.slice(6).trim()); + } else { + throw new Error(`Unexpected line format in response: ${line}`); + } + } + })(); + } +} diff --git a/packages/core/src/code_assist/setup.ts b/packages/core/src/code_assist/setup.ts index 0efe5f0c..6ff61523 100644 --- a/packages/core/src/code_assist/setup.ts +++ b/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 { - 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 || ''; diff --git a/packages/core/src/code_assist/types.ts b/packages/core/src/code_assist/types.ts index 733780e9..0bcc2b79 100644 --- a/packages/core/src/code_assist/types.ts +++ b/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 { -- cgit v1.2.3