Code Assist cleanup and docs (#993)

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<AsyncGenerator<GenerateContentResponse>> {
+    const resps = await this.streamEndpoint<CodeAssistResponse>(
+      'streamGenerateContent',
+      toCodeAssistRequest(req, this.projectId),
+    );
+    return (async function* (): AsyncGenerator<GenerateContentResponse> {
+      for await (const resp of resps) {
+        yield fromCodeAsistResponse(resp);
+      }
+    })();
+  }
+
+  async generateContent(
+    req: GenerateContentParameters,
+  ): Promise<GenerateContentResponse> {
+    const resp = await this.callEndpoint<CodeAssistResponse>(
+      'generateContent',
+      toCodeAssistRequest(req, this.projectId),
+    );
+    return fromCodeAsistResponse(resp);
+  }
+
+  async onboardUser(
+    req: OnboardUserRequest,
+  ): Promise<LongrunningOperationResponse> {
+    return await this.callEndpoint<LongrunningOperationResponse>(
+      'onboardUser',
+      req,
+    );
+  }
+
+  async loadCodeAssist(
+    req: LoadCodeAssistRequest,
+  ): Promise<LoadCodeAssistResponse> {
+    return await this.callEndpoint<LoadCodeAssistResponse>(
+      'loadCodeAssist',
+      req,
+    );
+  }
+
+  async countTokens(_req: CountTokensParameters): Promise<CountTokensResponse> {
+    return { totalTokens: 0 };
+  }
+
+  async embedContent(
+    _req: EmbedContentParameters,
+  ): Promise<EmbedContentResponse> {
+    throw Error();
+  }
+
+  async callEndpoint<T>(method: string, req: object): Promise<T> {
+    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<T>(
+    method: string,
+    req: object,
+  ): Promise<AsyncGenerator<T>> {
+    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<T> {
+      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}`);
+        }
+      }
+    })();
+  }
+}