From e3b5b2447db3fd33f7f7900063aac60dfcd33b9c Mon Sep 17 00:00:00 2001
From: Greg Richardson <greg.nmr@gmail.com>
Date: Thu, 14 Nov 2024 11:13:42 -0700
Subject: [PATCH] feat(telemetry): tool result name and success

---
 apps/postgres-new/app/api/chat/route.ts | 37 +++++++++++++++++++++++--
 apps/postgres-new/utils/telemetry.ts    |  6 ++++
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/apps/postgres-new/app/api/chat/route.ts b/apps/postgres-new/app/api/chat/route.ts
index ac5784d0..4a7a0b62 100644
--- a/apps/postgres-new/app/api/chat/route.ts
+++ b/apps/postgres-new/app/api/chat/route.ts
@@ -1,11 +1,11 @@
 import { createOpenAI } from '@ai-sdk/openai'
 import { Ratelimit } from '@upstash/ratelimit'
 import { kv } from '@vercel/kv'
-import { ToolInvocation, convertToCoreMessages, streamText } from 'ai'
+import { convertToCoreMessages, streamText, ToolInvocation, ToolResultPart } from 'ai'
 import { codeBlock } from 'common-tags'
 import { convertToCoreTools, maxMessageContext, maxRowLimit, tools } from '~/lib/tools'
 import { createClient } from '~/utils/supabase/server'
-import { logEvent } from '~/utils/telemetry'
+import { ChatInferenceEventToolResult, logEvent } from '~/utils/telemetry'
 
 // Allow streaming responses up to 30 seconds
 export const maxDuration = 30
@@ -134,6 +134,12 @@ export async function POST(req: Request) {
 
       // `tool` role indicates a tool result, `user` role indicates a user message
       const inputType = inputMessage.role === 'tool' ? 'tool-result' : 'user-message'
+      const toolResults =
+        inputMessage.role === 'tool'
+          ? inputMessage.content
+              .map((toolResult) => getEventToolResult(toolResult))
+              .filter((eventToolResult) => eventToolResult !== undefined)
+          : undefined
 
       // +1 for the assistant message just received
       const messageCount = coreMessages.length + 1
@@ -143,6 +149,7 @@ export async function POST(req: Request) {
         userId,
         messageCount,
         inputType,
+        toolResults,
         inputTokens: usage.promptTokens,
         outputTokens: usage.completionTokens,
         finishReason,
@@ -153,3 +160,29 @@ export async function POST(req: Request) {
 
   return result.toAIStreamResponse()
 }
+
+function getEventToolResult(toolResult: ToolResultPart): ChatInferenceEventToolResult | undefined {
+  try {
+    if (
+      !('result' in toolResult) ||
+      !toolResult.result ||
+      typeof toolResult.result !== 'object' ||
+      !('success' in toolResult.result) ||
+      typeof toolResult.result.success !== 'boolean'
+    ) {
+      return undefined
+    }
+
+    const {
+      toolName,
+      result: { success },
+    } = toolResult
+
+    return {
+      toolName,
+      success,
+    }
+  } catch (error) {
+    return undefined
+  }
+}
diff --git a/apps/postgres-new/utils/telemetry.ts b/apps/postgres-new/utils/telemetry.ts
index b093cf31..17315e5f 100644
--- a/apps/postgres-new/utils/telemetry.ts
+++ b/apps/postgres-new/utils/telemetry.ts
@@ -13,6 +13,11 @@ export type ChatRateLimitEvent = {
   }
 }
 
+export type ChatInferenceEventToolResult = {
+  toolName: string
+  success: boolean
+}
+
 /**
  * Event for an AI chat inference request-response.
  * Includes both input and output metadata.
@@ -24,6 +29,7 @@ export type ChatInferenceEvent = {
     userId: string
     messageCount: number
     inputType: 'user-message' | 'tool-result'
+    toolResults?: ChatInferenceEventToolResult[]
     inputTokens: number
     outputTokens: number
     finishReason: