fix(js/plugins/google-cloud): Write input and output logs for generate step (#3107)

Author: bryanatkinson

js/plugins/google-cloud/src/gcpOpenTelemetry.ts

@@ -413,27 +413,33 @@ class AdjustingTraceExporter implements SpanExporter {
       pathsTelemetry.tick(span, paths, this.logInputAndOutput, this.projectId);
       // Set root status explicitly
       span.attributes['genkit:rootState'] = span.attributes['genkit:state'];
-    }
-    if (type === 'action' && subtype === 'model') {
-      // Report generate metrics () for all model actions
-      generateTelemetry.tick(
-        span,
-        unused,
-        this.logInputAndOutput,
-        this.projectId
-      );
-    }
-    if (type === 'action' && subtype === 'tool') {
-      // TODO: Report input and output for tool actions
-    }
-    if (type === 'action' || type === 'flow' || type == 'flowStep') {
-      // Report request and latency metrics for all actions
-      actionTelemetry.tick(
-        span,
-        unused,
-        this.logInputAndOutput,
-        this.projectId
-      );
+    } else {
+      if (type === 'action' && subtype === 'model') {
+        // Report generate metrics () for all model actions
+        generateTelemetry.tick(
+          span,
+          unused,
+          this.logInputAndOutput,
+          this.projectId
+        );
+      }
+      if (type === 'action' && subtype === 'tool') {
+        // TODO: Report input and output for tool actions
+      }
+      if (
+        type === 'action' ||
+        type === 'flow' ||
+        type == 'flowStep' ||
+        type == 'util'
+      ) {
+        // Report request and latency metrics for all actions
+        actionTelemetry.tick(
+          span,
+          unused,
+          this.logInputAndOutput,
+          this.projectId
+        );
+      }
     }
     if (type === 'userEngagement') {
       // Report user acceptance and feedback metrics

js/plugins/google-cloud/src/telemetry/action.ts

@@ -32,21 +32,23 @@ class ActionTelemetry implements Telemetry {
     logInputAndOutput: boolean,
     projectId?: string
   ): void {
+    if (!logInputAndOutput) {
+      return;
+    }
     const attributes = span.attributes;
-
     const actionName = (attributes['genkit:name'] as string) || '<unknown>';
     const subtype = attributes['genkit:metadata:subtype'] as string;
-    const path = (attributes['genkit:path'] as string) || '<unknown>';
-    let featureName = extractOuterFeatureNameFromPath(path);
-    if (!featureName || featureName === '<unknown>') {
-      featureName = actionName;
-    }
 
-    if (subtype === 'tool' && logInputAndOutput) {
+    if (subtype === 'tool' || actionName === 'generate') {
+      const path = (attributes['genkit:path'] as string) || '<unknown>';
       const input = truncate(attributes['genkit:input'] as string);
       const output = truncate(attributes['genkit:output'] as string);
       const sessionId = attributes['genkit:sessionId'] as string;
       const threadName = attributes['genkit:threadName'] as string;
+      let featureName = extractOuterFeatureNameFromPath(path);
+      if (!featureName || featureName === '<unknown>') {
+        featureName = actionName;
+      }
 
       if (input) {
         this.writeLog(

js/plugins/google-cloud/tests/logs_test.ts

@@ -336,7 +336,7 @@ describe('GoogleCloudLogs', () => {
     await getExportedSpans();
 
     const logs = await getLogs(1, 100, logLines);
-    expect(logs.length).toEqual(8);
+    expect(logs.length).toEqual(10);
     const logObjectMessages = getStructuredLogMessages(logs);
     expect(logObjectMessages).toContain(
       'Config[testFlow > sub1 > sub2 > generate > testModel, testModel]'
@@ -349,6 +349,12 @@ describe('GoogleCloudLogs', () => {
     );
     expect(logObjectMessages).toContain('Input[testFlow, testFlow]');
     expect(logObjectMessages).toContain('Output[testFlow, testFlow]');
+    expect(logObjectMessages).toContain(
+      'Input[testFlow > sub1 > sub2 > generate, testFlow]'
+    );
+    expect(logObjectMessages).toContain(
+      'Output[testFlow > sub1 > sub2 > generate, testFlow]'
+    );
     // Ensure the model input/output has an associated role
     logs.map((log) => {
       const structuredLog = JSON.parse(log as string);
@@ -367,6 +373,58 @@ describe('GoogleCloudLogs', () => {
     });
   });
 
+  it('writes feature logs for generate without flow', async () => {
+    const testModel = createModel(ai, 'testModel', async () => {
+      return {
+        message: {
+          role: 'model',
+          content: [
+            {
+              text: 'response',
+            },
+          ],
+        },
+        finishReason: 'stop',
+        usage: {
+          inputTokens: 10,
+          outputTokens: 14,
+          inputCharacters: 8,
+          outputCharacters: 16,
+          inputImages: 1,
+          outputImages: 3,
+        },
+      };
+    });
+
+    await ai.generate({
+      model: testModel,
+      prompt: `a test prompt`,
+      config: {
+        temperature: 1.0,
+        topK: 3,
+        topP: 5,
+        maxOutputTokens: 7,
+      },
+    });
+
+    await getExportedSpans();
+
+    const logs = await getLogs(1, 100, logLines);
+    expect(logs.length).toEqual(7);
+    const logObjectMessages = getStructuredLogMessages(logs);
+    expect(logObjectMessages).toContain(
+      'Config[generate > testModel, testModel]'
+    );
+    expect(logObjectMessages).toContain(
+      'Input[generate > testModel, testModel] '
+    );
+    expect(logObjectMessages).toContain(
+      'Output[generate > testModel, testModel] '
+    );
+    expect(logObjectMessages).toContain('Input[generate, generate]');
+    expect(logObjectMessages).toContain('Output[generate, generate]');
+  });
+
   it('writes user feedback log', async () => {
     await appendSpan(
       'trace1',