Code assist service metrics. (#15024)

packages/core/src/code_assist/server.ts

@@ -16,6 +16,10 @@ import type {
   ClientMetadata,
   RetrieveUserQuotaRequest,
   RetrieveUserQuotaResponse,
+  ConversationOffered,
+  ConversationInteraction,
+  StreamingLatency,
+  RecordCodeAssistMetricsRequest,
 } from './types.js';
 import type {
   ListExperimentsRequest,
@@ -42,6 +46,11 @@ import {
   toCountTokenRequest,
   toGenerateContentRequest,
 } from './converter.js';
+import {
+  createConversationOffered,
+  formatProtoJsonDuration,
+} from './telemetry.js';
+import { getClientMetadata } from './experiments/client_metadata.js';
 
 /** HTTP options to be used in each of the requests. */
 export interface HttpOptions {
@@ -65,28 +74,60 @@ export class CodeAssistServer implements ContentGenerator {
     req: GenerateContentParameters,
     userPromptId: string,
   ): Promise<AsyncGenerator<GenerateContentResponse>> {
-    const resps = await this.requestStreamingPost<CaGenerateContentResponse>(
-      'streamGenerateContent',
-      toGenerateContentRequest(
-        req,
-        userPromptId,
-        this.projectId,
-        this.sessionId,
-      ),
-      req.config?.abortSignal,
-    );
-    return (async function* (): AsyncGenerator<GenerateContentResponse> {
-      for await (const resp of resps) {
-        yield fromGenerateContentResponse(resp);
+    const responses =
+      await this.requestStreamingPost<CaGenerateContentResponse>(
+        'streamGenerateContent',
+        toGenerateContentRequest(
+          req,
+          userPromptId,
+          this.projectId,
+          this.sessionId,
+        ),
+        req.config?.abortSignal,
+      );
+
+    const streamingLatency: StreamingLatency = {};
+    const start = Date.now();
+    let isFirst = true;
+
+    return (async function* (
+      server: CodeAssistServer,
+    ): AsyncGenerator<GenerateContentResponse> {
+      for await (const response of responses) {
+        if (isFirst) {
+          streamingLatency.firstMessageLatency = formatProtoJsonDuration(
+            Date.now() - start,
+          );
+          isFirst = false;
+        }
+
+        streamingLatency.totalLatency = formatProtoJsonDuration(
+          Date.now() - start,
+        );
+
+        const translatedResponse = fromGenerateContentResponse(response);
+
+        if (response.traceId) {
+          const offered = createConversationOffered(
+            translatedResponse,
+            response.traceId,
+            req.config?.abortSignal,
+            streamingLatency,
+          );
+          await server.recordConversationOffered(offered);
+        }
+
+        yield translatedResponse;
       }
-    })();
+    })(this);
   }
 
   async generateContent(
     req: GenerateContentParameters,
     userPromptId: string,
   ): Promise<GenerateContentResponse> {
-    const resp = await this.requestPost<CaGenerateContentResponse>(
+    const start = Date.now();
+    const response = await this.requestPost<CaGenerateContentResponse>(
       'generateContent',
       toGenerateContentRequest(
         req,
@@ -96,7 +137,25 @@ export class CodeAssistServer implements ContentGenerator {
       ),
       req.config?.abortSignal,
     );
-    return fromGenerateContentResponse(resp);
+    const duration = formatProtoJsonDuration(Date.now() - start);
+    const streamingLatency: StreamingLatency = {
+      totalLatency: duration,
+      firstMessageLatency: duration,
+    };
+
+    const translatedResponse = fromGenerateContentResponse(response);
+
+    if (response.traceId) {
+      const offered = createConversationOffered(
+        translatedResponse,
+        response.traceId,
+        req.config?.abortSignal,
+        streamingLatency,
+      );
+      await this.recordConversationOffered(offered);
+    }
+
+    return translatedResponse;
   }
 
   async onboardUser(
@@ -176,6 +235,40 @@ export class CodeAssistServer implements ContentGenerator {
     );
   }
 
+  async recordConversationOffered(
+    conversationOffered: ConversationOffered,
+  ): Promise<void> {
+    if (!this.projectId) {
+      return;
+    }
+
+    await this.recordCodeAssistMetrics({
+      project: this.projectId,
+      metadata: await getClientMetadata(),
+      metrics: [{ conversationOffered }],
+    });
+  }
+
+  async recordConversationInteraction(
+    interaction: ConversationInteraction,
+  ): Promise<void> {
+    if (!this.projectId) {
+      return;
+    }
+
+    await this.recordCodeAssistMetrics({
+      project: this.projectId,
+      metadata: await getClientMetadata(),
+      metrics: [{ conversationInteraction: interaction }],
+    });
+  }
+
+  async recordCodeAssistMetrics(
+    request: RecordCodeAssistMetricsRequest,
+  ): Promise<void> {
+    return this.requestPost<void>('recordCodeAssistMetrics', request);
+  }
+
   async requestPost<T>(
     method: string,
     req: object,