Skip to content

Commit 86ca590

Browse files
k11kirkyk11kirky
authored
feat: cached + reasoning tokens (#420)
* feat: cached tokens * fix: prettier + logs * fix: vercel types * fix preittier
1 parent 6498195 commit 86ca590

File tree

9 files changed

+194
-21
lines changed

9 files changed

+194
-21
lines changed

‎posthog-ai/CHANGELOG.md

+6
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,9 @@
1+
# 3.3.0 - 2025-03-08
2+
3+
- feat: add reasoning and cache tokens to openai and anthropic
4+
- feat: add tool support for vercel
5+
- feat: add support for other media types vercel
6+
17
# 3.2.1 - 2025-02-11
28

39
- fix: add experimental_wrapLanguageModel to vercel middleware supporting older versions of ai

‎posthog-ai/package.json

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"name": "@posthog/ai",
3-
"version": "3.2.1",
3+
"version": "3.3.0",
44
"description": "PostHog Node.js AI integrations",
55
"repository": {
66
"type": "git",

‎posthog-ai/src/anthropic/index.ts

+12-1
Original file line numberDiff line numberDiff line change
@@ -70,9 +70,16 @@ export class WrappedMessages extends AnthropicOriginal.Messages {
7070
if (anthropicParams.stream) {
7171
return parentPromise.then((value) => {
7272
let accumulatedContent = ''
73-
const usage: { inputTokens: number; outputTokens: number } = {
73+
const usage: {
74+
inputTokens: number
75+
outputTokens: number
76+
cacheCreationInputTokens?: number
77+
cacheReadInputTokens?: number
78+
} = {
7479
inputTokens: 0,
7580
outputTokens: 0,
81+
cacheCreationInputTokens: 0,
82+
cacheReadInputTokens: 0,
7683
}
7784
if ('tee' in value) {
7885
const [stream1, stream2] = value.tee()
@@ -87,6 +94,8 @@ export class WrappedMessages extends AnthropicOriginal.Messages {
8794
}
8895
if (chunk.type == 'message_start') {
8996
usage.inputTokens = chunk.message.usage.input_tokens ?? 0
97+
usage.cacheCreationInputTokens = chunk.message.usage.cache_creation_input_tokens ?? 0
98+
usage.cacheReadInputTokens = chunk.message.usage.cache_read_input_tokens ?? 0
9099
}
91100
if ('usage' in chunk) {
92101
usage.outputTokens = chunk.usage.output_tokens ?? 0
@@ -156,6 +165,8 @@ export class WrappedMessages extends AnthropicOriginal.Messages {
156165
usage: {
157166
inputTokens: result.usage.input_tokens ?? 0,
158167
outputTokens: result.usage.output_tokens ?? 0,
168+
cacheCreationInputTokens: result.usage.cache_creation_input_tokens ?? 0,
169+
cacheReadInputTokens: result.usage.cache_read_input_tokens ?? 0,
159170
},
160171
})
161172
}

‎posthog-ai/src/langchain/callbacks.ts

+6
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ interface GenerationMetadata extends SpanMetadata {
2828
modelParams?: Record<string, any>
2929
/** The base URL—for example, the API base used */
3030
baseUrl?: string
31+
/** The tools used in the generation */
32+
tools?: Record<string, any>
3133
}
3234

3335
/** A run may either be a Span or a Generation */
@@ -420,6 +422,10 @@ export class LangChainCallbackHandler extends BaseCallbackHandler {
420422
$ai_base_url: run.baseUrl,
421423
}
422424

425+
if (run.tools) {
426+
eventProperties['$ai_tools'] = withPrivacyMode(this.client, this.privacyMode, run.tools)
427+
}
428+
423429
if (output instanceof Error) {
424430
eventProperties['$ai_http_status'] = (output as any).status || 500
425431
eventProperties['$ai_error'] = output.toString()

‎posthog-ai/src/openai/azure.ts

+10-1
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,12 @@ export class WrappedCompletions extends AzureOpenAI.Chat.Completions {
8686
if (openAIParams.stream) {
8787
return parentPromise.then((value) => {
8888
let accumulatedContent = ''
89-
let usage: { inputTokens: number; outputTokens: number } = {
89+
let usage: {
90+
inputTokens: number
91+
outputTokens: number
92+
reasoningTokens?: number
93+
cacheReadInputTokens?: number
94+
} = {
9095
inputTokens: 0,
9196
outputTokens: 0,
9297
}
@@ -105,6 +110,8 @@ export class WrappedCompletions extends AzureOpenAI.Chat.Completions {
105110
usage = {
106111
inputTokens: chunk.usage.prompt_tokens ?? 0,
107112
outputTokens: chunk.usage.completion_tokens ?? 0,
113+
reasoningTokens: chunk.usage.completion_tokens_details?.reasoning_tokens ?? 0,
114+
cacheReadInputTokens: chunk.usage.prompt_tokens_details?.cached_tokens ?? 0,
108115
}
109116
}
110117
}
@@ -176,6 +183,8 @@ export class WrappedCompletions extends AzureOpenAI.Chat.Completions {
176183
usage: {
177184
inputTokens: result.usage?.prompt_tokens ?? 0,
178185
outputTokens: result.usage?.completion_tokens ?? 0,
186+
reasoningTokens: result.usage?.completion_tokens_details?.reasoning_tokens ?? 0,
187+
cacheReadInputTokens: result.usage?.prompt_tokens_details?.cached_tokens ?? 0,
179188
},
180189
})
181190
}

‎posthog-ai/src/openai/index.ts

+13-2
Original file line numberDiff line numberDiff line change
@@ -88,11 +88,18 @@ export class WrappedCompletions extends OpenAIOrignal.Chat.Completions {
8888
return parentPromise.then((value) => {
8989
if ('tee' in value) {
9090
const [stream1, stream2] = value.tee()
91-
// Use one stream for tracking
9291
;(async () => {
9392
try {
9493
let accumulatedContent = ''
95-
let usage = { inputTokens: 0, outputTokens: 0 }
94+
let usage: {
95+
inputTokens?: number
96+
outputTokens?: number
97+
reasoningTokens?: number
98+
cacheReadInputTokens?: number
99+
} = {
100+
inputTokens: 0,
101+
outputTokens: 0,
102+
}
96103

97104
for await (const chunk of stream1) {
98105
const delta = chunk?.choices?.[0]?.delta?.content ?? ''
@@ -101,6 +108,8 @@ export class WrappedCompletions extends OpenAIOrignal.Chat.Completions {
101108
usage = {
102109
inputTokens: chunk.usage.prompt_tokens ?? 0,
103110
outputTokens: chunk.usage.completion_tokens ?? 0,
111+
reasoningTokens: chunk.usage.completion_tokens_details?.reasoning_tokens ?? 0,
112+
cacheReadInputTokens: chunk.usage.prompt_tokens_details?.cached_tokens ?? 0,
104113
}
105114
}
106115
}
@@ -165,6 +174,8 @@ export class WrappedCompletions extends OpenAIOrignal.Chat.Completions {
165174
usage: {
166175
inputTokens: result.usage?.prompt_tokens ?? 0,
167176
outputTokens: result.usage?.completion_tokens ?? 0,
177+
reasoningTokens: result.usage?.completion_tokens_details?.reasoning_tokens ?? 0,
178+
cacheReadInputTokens: result.usage?.prompt_tokens_details?.cached_tokens ?? 0,
168179
},
169180
})
170181
}

‎posthog-ai/src/utils.ts

+17-1
Original file line numberDiff line numberDiff line change
@@ -118,10 +118,17 @@ export type SendEventToPosthogParams = {
118118
latency: number
119119
baseURL: string
120120
httpStatus: number
121-
usage?: { inputTokens?: number; outputTokens?: number }
121+
usage?: {
122+
inputTokens?: number
123+
outputTokens?: number
124+
reasoningTokens?: any
125+
cacheReadInputTokens?: any
126+
cacheCreationInputTokens?: any
127+
}
122128
params: (ChatCompletionCreateParamsBase | MessageCreateParams) & MonitoringParams
123129
isError?: boolean
124130
error?: string
131+
tools?: any
125132
}
126133

127134
export const sendEventToPosthog = ({
@@ -139,6 +146,7 @@ export const sendEventToPosthog = ({
139146
usage = {},
140147
isError = false,
141148
error,
149+
tools,
142150
}: SendEventToPosthogParams): void => {
143151
if (client.capture) {
144152
let errorData = {}
@@ -159,6 +167,12 @@ export const sendEventToPosthog = ({
159167
}
160168
}
161169

170+
let additionalTokenValues = {
171+
...(usage.reasoningTokens ? { $ai_reasoning_tokens: usage.reasoningTokens } : {}),
172+
...(usage.cacheReadInputTokens ? { $ai_cache_read_input_tokens: usage.cacheReadInputTokens } : {}),
173+
...(usage.cacheCreationInputTokens ? { $ai_cache_creation_input_tokens: usage.cacheCreationInputTokens } : {}),
174+
}
175+
162176
client.capture({
163177
distinctId: distinctId ?? traceId,
164178
event: '$ai_generation',
@@ -171,11 +185,13 @@ export const sendEventToPosthog = ({
171185
$ai_http_status: httpStatus,
172186
$ai_input_tokens: usage.inputTokens ?? 0,
173187
$ai_output_tokens: usage.outputTokens ?? 0,
188+
...additionalTokenValues,
174189
$ai_latency: latency,
175190
$ai_trace_id: traceId,
176191
$ai_base_url: baseURL,
177192
...params.posthogProperties,
178193
...(distinctId ? {} : { $process_person_profile: false }),
194+
...(tools ? { $ai_tools: tools } : {}),
179195
...errorData,
180196
...costOverrideData,
181197
},

‎posthog-ai/src/vercel/middleware.ts

+93-15
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,13 @@ interface CreateInstrumentationMiddlewareOptions {
2727
}
2828

2929
interface PostHogInput {
30-
content: string
3130
role: string
31+
type?: string
32+
content?:
33+
| string
34+
| {
35+
[key: string]: any
36+
}
3237
}
3338

3439
const mapVercelParams = (params: any): Record<string, any> => {
@@ -45,18 +50,60 @@ const mapVercelParams = (params: any): Record<string, any> => {
4550

4651
const mapVercelPrompt = (prompt: LanguageModelV1Prompt): PostHogInput[] => {
4752
return prompt.map((p) => {
48-
let content = ''
53+
let content = {}
4954
if (Array.isArray(p.content)) {
50-
content = p.content
51-
.map((c) => {
52-
if (c.type === 'text') {
53-
return c.text
55+
content = p.content.map((c) => {
56+
if (c.type === 'text') {
57+
return {
58+
type: 'text',
59+
content: c.text,
5460
}
55-
return ''
56-
})
57-
.join('')
61+
} else if (c.type === 'image') {
62+
return {
63+
type: 'image',
64+
content: {
65+
// if image is a url use it, or use "none supported"
66+
image: c.image instanceof URL ? c.image.toString() : 'raw images not supported',
67+
mimeType: c.mimeType,
68+
},
69+
}
70+
} else if (c.type === 'file') {
71+
return {
72+
type: 'file',
73+
content: {
74+
file: c.data instanceof URL ? c.data.toString() : 'raw files not supported',
75+
mimeType: c.mimeType,
76+
},
77+
}
78+
} else if (c.type === 'tool-call') {
79+
return {
80+
type: 'tool-call',
81+
content: {
82+
toolCallId: c.toolCallId,
83+
toolName: c.toolName,
84+
args: c.args,
85+
},
86+
}
87+
} else if (c.type === 'tool-result') {
88+
return {
89+
type: 'tool-result',
90+
content: {
91+
toolCallId: c.toolCallId,
92+
toolName: c.toolName,
93+
result: c.result,
94+
isError: c.isError,
95+
},
96+
}
97+
}
98+
return {
99+
content: '',
100+
}
101+
})
58102
} else {
59-
content = p.content
103+
content = {
104+
type: 'text',
105+
text: p.content,
106+
}
60107
}
61108
return {
62109
role: p.role,
@@ -91,10 +138,22 @@ export const createInstrumentationMiddleware = (
91138
options.posthogModelOverride ?? (result.response?.modelId ? result.response.modelId : model.modelId)
92139
const provider = options.posthogProviderOverride ?? extractProvider(model)
93140
const baseURL = '' // cannot currently get baseURL from vercel
94-
let content = result.text
95-
if (!content) {
96-
// support generate Object
97-
content = result.toolCalls?.[0].args || JSON.stringify(result)
141+
let content = result.text || JSON.stringify(result)
142+
// let tools = result.toolCalls
143+
let providerMetadata = result.providerMetadata
144+
let additionalTokenValues = {
145+
...(providerMetadata?.openai?.reasoningTokens
146+
? { reasoningTokens: providerMetadata.openai.reasoningTokens }
147+
: {}),
148+
...(providerMetadata?.openai?.cachedPromptToken
149+
? { cacheReadInputTokens: providerMetadata.openai.cachedPromptTokens }
150+
: {}),
151+
...(providerMetadata?.anthropic
152+
? {
153+
cacheReadInputTokens: providerMetadata.anthropic.cacheReadInputTokens,
154+
cacheCreationInputTokens: providerMetadata.anthropic.cacheCreationInputTokens,
155+
}
156+
: {}),
98157
}
99158
sendEventToPosthog({
100159
client: phClient,
@@ -111,6 +170,7 @@ export const createInstrumentationMiddleware = (
111170
usage: {
112171
inputTokens: result.usage.promptTokens,
113172
outputTokens: result.usage.completionTokens,
173+
...additionalTokenValues,
114174
},
115175
})
116176

@@ -143,7 +203,13 @@ export const createInstrumentationMiddleware = (
143203
wrapStream: async ({ doStream, params }) => {
144204
const startTime = Date.now()
145205
let generatedText = ''
146-
let usage: { inputTokens?: number; outputTokens?: number } = {}
206+
let usage: {
207+
inputTokens?: number
208+
outputTokens?: number
209+
reasoningTokens?: any
210+
cacheReadInputTokens?: any
211+
cacheCreationInputTokens?: any
212+
} = {}
147213
const mergedParams = {
148214
...options,
149215
...mapVercelParams(params),
@@ -164,6 +230,18 @@ export const createInstrumentationMiddleware = (
164230
inputTokens: chunk.usage?.promptTokens,
165231
outputTokens: chunk.usage?.completionTokens,
166232
}
233+
if (chunk.providerMetadata?.openai?.reasoningTokens) {
234+
usage.reasoningTokens = chunk.providerMetadata.openai.reasoningTokens
235+
}
236+
if (chunk.providerMetadata?.openai?.cachedPromptToken) {
237+
usage.cacheReadInputTokens = chunk.providerMetadata.openai.cachedPromptToken
238+
}
239+
if (chunk.providerMetadata?.anthropic?.cacheReadInputTokens) {
240+
usage.cacheReadInputTokens = chunk.providerMetadata.anthropic.cacheReadInputTokens
241+
}
242+
if (chunk.providerMetadata?.anthropic?.cacheCreationInputTokens) {
243+
usage.cacheCreationInputTokens = chunk.providerMetadata.anthropic.cacheCreationInputTokens
244+
}
167245
}
168246
controller.enqueue(chunk)
169247
},

‎posthog-ai/tests/openai.test.ts

+36
Original file line numberDiff line numberDiff line change
@@ -225,4 +225,40 @@ describe('PostHogOpenAI - Jest test suite', () => {
225225
expect(properties['$ai_stream']).toBe(false)
226226
expect(properties['foo']).toBe('bar')
227227
})
228+
229+
conditionalTest('reasoning and cache tokens', async () => {
230+
// Set up mock response with standard token usage
231+
mockOpenAiChatResponse.usage = {
232+
prompt_tokens: 20,
233+
completion_tokens: 10,
234+
total_tokens: 30,
235+
// Add the detailed token usage that OpenAI would return
236+
completion_tokens_details: {
237+
reasoning_tokens: 15,
238+
},
239+
prompt_tokens_details: {
240+
cached_tokens: 5,
241+
},
242+
}
243+
244+
// Create a completion with additional token tracking
245+
await client.chat.completions.create({
246+
model: 'gpt-4',
247+
messages: [{ role: 'user', content: 'Hello' }],
248+
posthogDistinctId: 'test-id',
249+
posthogProperties: { foo: 'bar' },
250+
})
251+
252+
expect(mockPostHogClient.capture).toHaveBeenCalledTimes(1)
253+
const [captureArgs] = (mockPostHogClient.capture as jest.Mock).mock.calls
254+
const { properties } = captureArgs[0]
255+
256+
// Check standard token properties
257+
expect(properties['$ai_input_tokens']).toBe(20)
258+
expect(properties['$ai_output_tokens']).toBe(10)
259+
260+
// Check the new token properties
261+
expect(properties['$ai_reasoning_tokens']).toBe(15)
262+
expect(properties['$ai_cache_read_input_tokens']).toBe(5)
263+
})
228264
})

0 commit comments

Comments
 (0)