@posthog/ai 2.3.1 → 3.0.0

package/src/utils.ts

@@ -1,7 +1,9 @@
 import { PostHog } from 'posthog-node'
 import OpenAIOrignal from 'openai'
+import AnthropicOriginal from '@anthropic-ai/sdk'
 
 type ChatCompletionCreateParamsBase = OpenAIOrignal.Chat.Completions.ChatCompletionCreateParams
+type MessageCreateParams = AnthropicOriginal.Messages.MessageCreateParams
 
 export interface MonitoringParams {
   posthogDistinctId?: string
@@ -11,7 +13,12 @@ export interface MonitoringParams {
   posthogGroups?: Record<string, any>
 }
 
-export const getModelParams = (params: ChatCompletionCreateParamsBase & MonitoringParams): Record<string, any> => {
+export const getModelParams = (
+  params: ((ChatCompletionCreateParamsBase | MessageCreateParams) & MonitoringParams) | null
+): Record<string, any> => {
+  if (!params) {
+    return {}
+  }
   const modelParams: Record<string, any> = {}
   const paramKeys = [
     'temperature',
@@ -76,7 +83,7 @@ export const formatResponseOpenAI = (response: any): Array<{ role: string; conte
   return output
 }
 
-export const mergeSystemPrompt = (params: ChatCompletionCreateParamsBase & MonitoringParams, provider: string): any => {
+export const mergeSystemPrompt = (params: MessageCreateParams & MonitoringParams, provider: string): any => {
   if (provider == 'anthropic') {
     const messages = params.messages || []
     if (!(params as any).system) {
@@ -104,7 +111,7 @@ export type SendEventToPosthogParams = {
   baseURL: string
   httpStatus: number
   usage?: { inputTokens?: number; outputTokens?: number }
-  params: ChatCompletionCreateParamsBase & MonitoringParams
+  params: (ChatCompletionCreateParamsBase | MessageCreateParams) & MonitoringParams
   isError?: boolean
   error?: string
 }

package/src/vercel/middleware.ts

@@ -1,10 +1,5 @@
-import { experimental_wrapLanguageModel as wrapLanguageModel } from 'ai'
-import type {
-  LanguageModelV1,
-  Experimental_LanguageModelV1Middleware as LanguageModelV1Middleware,
-  LanguageModelV1Prompt,
-  LanguageModelV1StreamPart,
-} from 'ai'
+import { wrapLanguageModel } from 'ai'
+import type { LanguageModelV1, LanguageModelV1Middleware, LanguageModelV1Prompt, LanguageModelV1StreamPart } from 'ai'
 import { v4 as uuidv4 } from 'uuid'
 import { PostHog } from 'posthog-node'
 import { sendEventToPosthog } from '../utils'
@@ -17,6 +12,12 @@ interface ClientOptions {
   posthogGroups?: Record<string, any>
   posthogModelOverride?: string
   posthogProviderOverride?: string
+  posthogCostOverride?: CostOverride
+}
+
+interface CostOverride {
+  inputTokens: number
+  outputTokens: number
 }
 
 interface CreateInstrumentationMiddlewareOptions {
@@ -27,6 +28,7 @@ interface CreateInstrumentationMiddlewareOptions {
   posthogGroups?: Record<string, any>
   posthogModelOverride?: string
   posthogProviderOverride?: string
+  posthogCostOverride?: CostOverride
 }
 
 interface PostHogInput {
@@ -83,7 +85,7 @@ export const createInstrumentationMiddleware = (
   const middleware: LanguageModelV1Middleware = {
     wrapGenerate: async ({ doGenerate, params }) => {
       const startTime = Date.now()
-      let mergedParams = {
+      const mergedParams = {
         ...options,
         ...mapVercelParams(params),
       }
@@ -143,7 +145,7 @@ export const createInstrumentationMiddleware = (
       const startTime = Date.now()
       let generatedText = ''
       let usage: { inputTokens?: number; outputTokens?: number } = {}
-      let mergedParams = {
+      const mergedParams = {
         ...options,
         ...mapVercelParams(params),
       }