From b5a295b6ebd964e13c7ede09b5b847ea80f321db Mon Sep 17 00:00:00 2001
From: Georgiy Tarasov <gtarasov.work@gmail.com>
Date: Thu, 20 Feb 2025 13:07:38 +0100
Subject: [PATCH] feat(llm-observability): metric and feedback methods (#1709)

* feat: metric and feedback methods

* fix: events
---
 src/__tests__/ai.test.ts | 78 ++++++++++++++++++++++++++++++++++++++++
 src/posthog-core.ts      | 26 ++++++++++++++
 2 files changed, 104 insertions(+)
 create mode 100644 src/__tests__/ai.test.ts

diff --git a/src/__tests__/ai.test.ts b/src/__tests__/ai.test.ts
new file mode 100644
index 000000000..109115e84
--- /dev/null
+++ b/src/__tests__/ai.test.ts
@@ -0,0 +1,78 @@
+import { defaultPostHog } from './helpers/posthog-instance'
+import type { PostHogConfig } from '../types'
+import { uuidv7 } from '../uuidv7'
+
+describe('ai', () => {
+    beforeEach(() => {
+        console.error = jest.fn()
+    })
+
+    const setup = (config: Partial<PostHogConfig> = {}, token: string = uuidv7()) => {
+        const beforeSendMock = jest.fn().mockImplementation((e) => e)
+        const posthog = defaultPostHog().init(token, { ...config, before_send: beforeSendMock }, token)!
+        posthog.debug()
+        return { posthog, beforeSendMock }
+    }
+
+    describe('captureTraceMetric()', () => {
+        it('should capture metric', () => {
+            const { posthog, beforeSendMock } = setup()
+
+            posthog.captureTraceMetric('123', 'test', 'test')
+
+            const { event, properties } = beforeSendMock.mock.calls[0][0]
+            expect(event).toBe('$ai_metric')
+            expect(properties['$ai_trace_id']).toBe('123')
+            expect(properties['$ai_metric_name']).toBe('test')
+            expect(properties['$ai_metric_value']).toBe('test')
+        })
+
+        it('should convert numeric values', () => {
+            const { posthog, beforeSendMock } = setup()
+
+            posthog.captureTraceMetric(123, 'test', 1)
+
+            const { event, properties } = beforeSendMock.mock.calls[0][0]
+            expect(event).toBe('$ai_metric')
+            expect(properties['$ai_trace_id']).toBe('123')
+            expect(properties['$ai_metric_name']).toBe('test')
+            expect(properties['$ai_metric_value']).toBe('1')
+        })
+
+        it('should convert boolean metric_value', () => {
+            const { posthog, beforeSendMock } = setup()
+
+            posthog.captureTraceMetric('test', 'test', false)
+
+            const { event, properties } = beforeSendMock.mock.calls[0][0]
+            expect(event).toBe('$ai_metric')
+            expect(properties['$ai_trace_id']).toBe('test')
+            expect(properties['$ai_metric_name']).toBe('test')
+            expect(properties['$ai_metric_value']).toBe('false')
+        })
+    })
+
+    describe('captureTraceFeedback()', () => {
+        it('should capture feedback', () => {
+            const { posthog, beforeSendMock } = setup()
+
+            posthog.captureTraceFeedback('123', 'feedback')
+
+            const { event, properties } = beforeSendMock.mock.calls[0][0]
+            expect(event).toBe('$ai_feedback')
+            expect(properties['$ai_trace_id']).toBe('123')
+            expect(properties['$ai_feedback_text']).toBe('feedback')
+        })
+
+        it('should convert numeric values', () => {
+            const { posthog, beforeSendMock } = setup()
+
+            posthog.captureTraceFeedback(123, 'feedback')
+
+            const { event, properties } = beforeSendMock.mock.calls[0][0]
+            expect(event).toBe('$ai_feedback')
+            expect(properties['$ai_trace_id']).toBe('123')
+            expect(properties['$ai_feedback_text']).toBe('feedback')
+        })
+    })
+})
diff --git a/src/posthog-core.ts b/src/posthog-core.ts
index e4be5b877..5e5e96859 100644
--- a/src/posthog-core.ts
+++ b/src/posthog-core.ts
@@ -2166,6 +2166,32 @@ export class PostHog {
     public getPageViewId(): string | undefined {
         return this.pageViewManager._currentPageview?.pageViewId
     }
+
+    /**
+     * Capture written user feedback for a LLM trace. Numeric values are converted to strings.
+     * @param traceId The trace ID to capture feedback for.
+     * @param userFeedback The feedback to capture.
+     */
+    captureTraceFeedback(traceId: string | number, userFeedback: string) {
+        this.capture('$ai_feedback', {
+            $ai_trace_id: String(traceId),
+            $ai_feedback_text: userFeedback,
+        })
+    }
+
+    /**
+     * Capture a metric for a LLM trace. Numeric values are converted to strings.
+     * @param traceId The trace ID to capture the metric for.
+     * @param metricName The name of the metric to capture.
+     * @param metricValue The value of the metric to capture.
+     */
+    captureTraceMetric(traceId: string | number, metricName: string, metricValue: string | number | boolean) {
+        this.capture('$ai_metric', {
+            $ai_trace_id: String(traceId),
+            $ai_metric_name: metricName,
+            $ai_metric_value: String(metricValue),
+        })
+    }
 }
 
 safewrapClass(PostHog, ['identify'])