@posthog/ai 5.2.2 → 6.0.0

package/src/langchain/callbacks.ts

@@ -1,640 +0,0 @@
-import { PostHog } from 'posthog-node'
-import { withPrivacyMode, getModelParams } from '../utils'
-import { BaseCallbackHandler } from '@langchain/core/callbacks/base'
-import type { Serialized } from '@langchain/core/load/serializable'
-import type { ChainValues } from '@langchain/core/utils/types'
-import type { LLMResult } from '@langchain/core/outputs'
-import type { AgentAction, AgentFinish } from '@langchain/core/agents'
-import type { DocumentInterface } from '@langchain/core/documents'
-import { ToolCall } from '@langchain/core/messages/tool'
-import { BaseMessage } from '@langchain/core/messages'
-
-interface SpanMetadata {
-  /** Name of the trace/span (e.g. chain name) */
-  name: string
-  /** Timestamp (in ms) when the run started */
-  startTime: number
-  /** Timestamp (in ms) when the run ended (if already finished) */
-  endTime?: number
-  /** The input state */
-  input?: any
-}
-
-interface GenerationMetadata extends SpanMetadata {
-  /** Provider used (e.g. openai, anthropic) */
-  provider?: string
-  /** Model name used in the generation */
-  model?: string
-  /** The model parameters (temperature, max_tokens, etc.) */
-  modelParams?: Record<string, any>
-  /** The base URL—for example, the API base used */
-  baseUrl?: string
-  /** The tools used in the generation */
-  tools?: Record<string, any>
-}
-
-/** A run may either be a Span or a Generation */
-type RunMetadata = SpanMetadata | GenerationMetadata
-
-/** Storage for run metadata */
-type RunMetadataStorage = { [runId: string]: RunMetadata }
-
-export class LangChainCallbackHandler extends BaseCallbackHandler {
-  public name = 'PosthogCallbackHandler'
-  private client: PostHog
-  private distinctId?: string | number
-  private traceId?: string | number
-  private properties: Record<string, any>
-  private privacyMode: boolean
-  private groups: Record<string, any>
-  private debug: boolean
-
-  private runs: RunMetadataStorage = {}
-  private parentTree: { [runId: string]: string } = {}
-
-  constructor(options: {
-    client: PostHog
-    distinctId?: string | number
-    traceId?: string | number
-    properties?: Record<string, any>
-    privacyMode?: boolean
-    groups?: Record<string, any>
-    debug?: boolean
-  }) {
-    if (!options.client) {
-      throw new Error('PostHog client is required')
-    }
-    super()
-    this.client = options.client
-    this.distinctId = options.distinctId
-    this.traceId = options.traceId
-    this.properties = options.properties || {}
-    this.privacyMode = options.privacyMode || false
-    this.groups = options.groups || {}
-    this.debug = options.debug || false
-  }
-
-  // ===== CALLBACK METHODS =====
-
-  public handleChainStart(
-    chain: Serialized,
-    inputs: ChainValues,
-    runId: string,
-    parentRunId?: string,
-    tags?: string[],
-    metadata?: Record<string, unknown>,
-    runType?: string,
-    runName?: string
-  ): void {
-    this._logDebugEvent('on_chain_start', runId, parentRunId, { inputs, tags })
-    this._setParentOfRun(runId, parentRunId)
-    this._setTraceOrSpanMetadata(chain, inputs, runId, parentRunId, metadata, tags, runName)
-  }
-
-  public handleChainEnd(
-    outputs: ChainValues,
-    runId: string,
-    parentRunId?: string,
-    tags?: string[],
-    // eslint-disable-next-line @typescript-eslint/no-unused-vars
-    kwargs?: { inputs?: Record<string, unknown> }
-  ): void {
-    this._logDebugEvent('on_chain_end', runId, parentRunId, { outputs, tags })
-    this._popRunAndCaptureTraceOrSpan(runId, parentRunId, outputs)
-  }
-
-  public handleChainError(
-    error: Error,
-    runId: string,
-    parentRunId?: string,
-    tags?: string[],
-    // eslint-disable-next-line @typescript-eslint/no-unused-vars
-    kwargs?: { inputs?: Record<string, unknown> }
-  ): void {
-    this._logDebugEvent('on_chain_error', runId, parentRunId, { error, tags })
-    this._popRunAndCaptureTraceOrSpan(runId, parentRunId, error)
-  }
-
-  public handleChatModelStart(
-    serialized: Serialized,
-    messages: BaseMessage[][],
-    runId: string,
-    parentRunId?: string,
-    extraParams?: Record<string, unknown>,
-    tags?: string[],
-    metadata?: Record<string, unknown>,
-    runName?: string
-  ): void {
-    this._logDebugEvent('on_chat_model_start', runId, parentRunId, { messages, tags })
-    this._setParentOfRun(runId, parentRunId)
-    // Flatten the two-dimensional messages and convert each message to a plain object
-    const input = messages.flat().map((m) => this._convertMessageToDict(m))
-    this._setLLMMetadata(serialized, runId, input, metadata, extraParams, runName)
-  }
-
-  public handleLLMStart(
-    serialized: Serialized,
-    prompts: string[],
-    runId: string,
-    parentRunId?: string,
-    extraParams?: Record<string, unknown>,
-    tags?: string[],
-    metadata?: Record<string, unknown>,
-    runName?: string
-  ): void {
-    this._logDebugEvent('on_llm_start', runId, parentRunId, { prompts, tags })
-    this._setParentOfRun(runId, parentRunId)
-    this._setLLMMetadata(serialized, runId, prompts, metadata, extraParams, runName)
-  }
-
-  public handleLLMEnd(
-    output: LLMResult,
-    runId: string,
-    parentRunId?: string,
-    tags?: string[],
-    // eslint-disable-next-line @typescript-eslint/no-unused-vars
-    extraParams?: Record<string, unknown>
-  ): void {
-    this._logDebugEvent('on_llm_end', runId, parentRunId, { output, tags })
-    this._popRunAndCaptureGeneration(runId, parentRunId, output)
-  }
-
-  public handleLLMError(
-    err: Error,
-    runId: string,
-    parentRunId?: string,
-    tags?: string[],
-    // eslint-disable-next-line @typescript-eslint/no-unused-vars
-    extraParams?: Record<string, unknown>
-  ): void {
-    this._logDebugEvent('on_llm_error', runId, parentRunId, { err, tags })
-    this._popRunAndCaptureGeneration(runId, parentRunId, err)
-  }
-
-  public handleToolStart(
-    tool: Serialized,
-    input: string,
-    runId: string,
-    parentRunId?: string,
-    tags?: string[],
-    metadata?: Record<string, unknown>,
-    runName?: string
-  ): void {
-    this._logDebugEvent('on_tool_start', runId, parentRunId, { input, tags })
-    this._setParentOfRun(runId, parentRunId)
-    this._setTraceOrSpanMetadata(tool, input, runId, parentRunId, metadata, tags, runName)
-  }
-
-  public handleToolEnd(output: any, runId: string, parentRunId?: string, tags?: string[]): void {
-    this._logDebugEvent('on_tool_end', runId, parentRunId, { output, tags })
-    this._popRunAndCaptureTraceOrSpan(runId, parentRunId, output)
-  }
-
-  public handleToolError(err: Error, runId: string, parentRunId?: string, tags?: string[]): void {
-    this._logDebugEvent('on_tool_error', runId, parentRunId, { err, tags })
-    this._popRunAndCaptureTraceOrSpan(runId, parentRunId, err)
-  }
-
-  public handleRetrieverStart(
-    retriever: Serialized,
-    query: string,
-    runId: string,
-    parentRunId?: string,
-    tags?: string[],
-    metadata?: Record<string, unknown>,
-    name?: string
-  ): void {
-    this._logDebugEvent('on_retriever_start', runId, parentRunId, { query, tags })
-    this._setParentOfRun(runId, parentRunId)
-    this._setTraceOrSpanMetadata(retriever, query, runId, parentRunId, metadata, tags, name)
-  }
-
-  public handleRetrieverEnd(
-    documents: DocumentInterface[],
-    runId: string,
-    parentRunId?: string,
-    tags?: string[]
-  ): void {
-    this._logDebugEvent('on_retriever_end', runId, parentRunId, { documents, tags })
-    this._popRunAndCaptureTraceOrSpan(runId, parentRunId, documents)
-  }
-
-  public handleRetrieverError(err: Error, runId: string, parentRunId?: string, tags?: string[]): void {
-    this._logDebugEvent('on_retriever_error', runId, parentRunId, { err, tags })
-    this._popRunAndCaptureTraceOrSpan(runId, parentRunId, err)
-  }
-
-  public handleAgentAction(action: AgentAction, runId: string, parentRunId?: string, tags?: string[]): void {
-    this._logDebugEvent('on_agent_action', runId, parentRunId, { action, tags })
-    this._setParentOfRun(runId, parentRunId)
-    this._setTraceOrSpanMetadata(null, action, runId, parentRunId)
-  }
-
-  public handleAgentEnd(action: AgentFinish, runId: string, parentRunId?: string, tags?: string[]): void {
-    this._logDebugEvent('on_agent_finish', runId, parentRunId, { action, tags })
-    this._popRunAndCaptureTraceOrSpan(runId, parentRunId, action)
-  }
-
-  // ===== PRIVATE HELPERS =====
-
-  private _setParentOfRun(runId: string, parentRunId?: string): void {
-    if (parentRunId) {
-      this.parentTree[runId] = parentRunId
-    }
-  }
-
-  private _popParentOfRun(runId: string): void {
-    delete this.parentTree[runId]
-  }
-
-  private _findRootRun(runId: string): string {
-    let id = runId
-    while (this.parentTree[id]) {
-      id = this.parentTree[id]
-    }
-    return id
-  }
-
-  private _setTraceOrSpanMetadata(
-    serialized: any,
-    input: any,
-    runId: string,
-    parentRunId?: string,
-    ...args: any[]
-  ): void {
-    // Use default names if not provided: if this is a top-level run, we mark it as a trace, otherwise as a span.
-    const defaultName = parentRunId ? 'span' : 'trace'
-    const runName = this._getLangchainRunName(serialized, ...args) || defaultName
-    this.runs[runId] = {
-      name: runName,
-      input,
-      startTime: Date.now(),
-    } as SpanMetadata
-  }
-
-  private _setLLMMetadata(
-    serialized: Serialized | null,
-    runId: string,
-    messages: any,
-    metadata?: any,
-    extraParams?: any,
-    runName?: string
-  ): void {
-    const runNameFound = this._getLangchainRunName(serialized, { extraParams, runName }) || 'generation'
-    const generation: GenerationMetadata = {
-      name: runNameFound,
-      input: messages,
-      startTime: Date.now(),
-    }
-    if (extraParams) {
-      generation.modelParams = getModelParams(extraParams.invocation_params)
-    }
-    if (metadata) {
-      if (metadata.ls_model_name) {
-        generation.model = metadata.ls_model_name
-      }
-      if (metadata.ls_provider) {
-        generation.provider = metadata.ls_provider
-      }
-    }
-    if (serialized && 'kwargs' in serialized && serialized.kwargs.openai_api_base) {
-      generation.baseUrl = serialized.kwargs.openai_api_base
-    }
-    this.runs[runId] = generation
-  }
-
-  private _popRunMetadata(runId: string): RunMetadata | undefined {
-    const endTime = Date.now()
-    const run = this.runs[runId]
-    if (!run) {
-      console.warn(`No run metadata found for run ${runId}`)
-      return undefined
-    }
-    run.endTime = endTime
-    delete this.runs[runId]
-    return run
-  }
-
-  private _getTraceId(runId: string): string {
-    return this.traceId ? String(this.traceId) : this._findRootRun(runId)
-  }
-
-  private _getParentRunId(traceId: string, runId: string, parentRunId?: string): string | undefined {
-    // Replace the parent-run if not found in our stored parent tree.
-    if (parentRunId && !this.parentTree[parentRunId]) {
-      return traceId
-    }
-    return parentRunId
-  }
-
-  private _popRunAndCaptureTraceOrSpan(
-    runId: string,
-    parentRunId: string | undefined,
-    outputs: ChainValues | DocumentInterface[] | AgentFinish | Error | any
-  ): void {
-    const traceId = this._getTraceId(runId)
-    this._popParentOfRun(runId)
-    const run = this._popRunMetadata(runId)
-    if (!run) {
-      return
-    }
-    if ('modelParams' in run) {
-      console.warn(`Run ${runId} is a generation, but attempted to be captured as a trace/span.`)
-      return
-    }
-    const actualParentRunId = this._getParentRunId(traceId, runId, parentRunId)
-    this._captureTraceOrSpan(traceId, runId, run as SpanMetadata, outputs, actualParentRunId)
-  }
-
-  private _captureTraceOrSpan(
-    traceId: string,
-    runId: string,
-    run: SpanMetadata,
-    outputs: ChainValues | DocumentInterface[] | AgentFinish | Error | any,
-    parentRunId?: string
-  ): void {
-    const eventName = parentRunId ? '$ai_span' : '$ai_trace'
-    const latency = run.endTime ? (run.endTime - run.startTime) / 1000 : 0
-    const eventProperties: Record<string, any> = {
-      $ai_trace_id: traceId,
-      $ai_input_state: withPrivacyMode(this.client, this.privacyMode, run.input),
-      $ai_latency: latency,
-      $ai_span_name: run.name,
-      $ai_span_id: runId,
-    }
-    if (parentRunId) {
-      eventProperties['$ai_parent_id'] = parentRunId
-    }
-
-    Object.assign(eventProperties, this.properties)
-    if (!this.distinctId) {
-      eventProperties['$process_person_profile'] = false
-    }
-    if (outputs instanceof Error) {
-      eventProperties['$ai_error'] = outputs.toString()
-      eventProperties['$ai_is_error'] = true
-    } else if (outputs !== undefined) {
-      eventProperties['$ai_output_state'] = withPrivacyMode(this.client, this.privacyMode, outputs)
-    }
-    this.client.capture({
-      distinctId: this.distinctId ? this.distinctId.toString() : runId,
-      event: eventName,
-      properties: eventProperties,
-      groups: this.groups,
-    })
-  }
-
-  private _popRunAndCaptureGeneration(
-    runId: string,
-    parentRunId: string | undefined,
-    response: LLMResult | Error
-  ): void {
-    const traceId = this._getTraceId(runId)
-    this._popParentOfRun(runId)
-    const run = this._popRunMetadata(runId)
-    if (!run || typeof run !== 'object' || !('modelParams' in run)) {
-      console.warn(`Run ${runId} is not a generation, but attempted to be captured as such.`)
-      return
-    }
-    const actualParentRunId = this._getParentRunId(traceId, runId, parentRunId)
-    this._captureGeneration(traceId, runId, run as GenerationMetadata, response, actualParentRunId)
-  }
-
-  private _captureGeneration(
-    traceId: string,
-    runId: string,
-    run: GenerationMetadata,
-    output: LLMResult | Error,
-    parentRunId?: string
-  ): void {
-    const latency = run.endTime ? (run.endTime - run.startTime) / 1000 : 0
-    const eventProperties: Record<string, any> = {
-      $ai_trace_id: traceId,
-      $ai_span_id: runId,
-      $ai_span_name: run.name,
-      $ai_parent_id: parentRunId,
-      $ai_provider: run.provider,
-      $ai_model: run.model,
-      $ai_model_parameters: run.modelParams,
-      $ai_input: withPrivacyMode(this.client, this.privacyMode, run.input),
-      $ai_http_status: 200,
-      $ai_latency: latency,
-      $ai_base_url: run.baseUrl,
-    }
-
-    if (run.tools) {
-      eventProperties['$ai_tools'] = withPrivacyMode(this.client, this.privacyMode, run.tools)
-    }
-
-    if (output instanceof Error) {
-      eventProperties['$ai_http_status'] = (output as any).status || 500
-      eventProperties['$ai_error'] = output.toString()
-      eventProperties['$ai_is_error'] = true
-    } else {
-      // Handle token usage
-      const [inputTokens, outputTokens, additionalTokenData] = this.parseUsage(output)
-      eventProperties['$ai_input_tokens'] = inputTokens
-      eventProperties['$ai_output_tokens'] = outputTokens
-
-      // Add additional token data to properties
-      if (additionalTokenData.cacheReadInputTokens) {
-        eventProperties['$ai_cache_read_tokens'] = additionalTokenData.cacheReadInputTokens
-      }
-      if (additionalTokenData.reasoningTokens) {
-        eventProperties['$ai_reasoning_tokens'] = additionalTokenData.reasoningTokens
-      }
-
-      // Handle generations/completions
-      let completions
-      if (output.generations && Array.isArray(output.generations)) {
-        const lastGeneration = output.generations[output.generations.length - 1]
-        if (Array.isArray(lastGeneration)) {
-          completions = lastGeneration.map((gen) => {
-            return { role: 'assistant', content: gen.text }
-          })
-        }
-      }
-
-      if (completions) {
-        eventProperties['$ai_output_choices'] = withPrivacyMode(this.client, this.privacyMode, completions)
-      }
-    }
-
-    Object.assign(eventProperties, this.properties)
-    if (!this.distinctId) {
-      eventProperties['$process_person_profile'] = false
-    }
-
-    this.client.capture({
-      distinctId: this.distinctId ? this.distinctId.toString() : traceId,
-      event: '$ai_generation',
-      properties: eventProperties,
-      groups: this.groups,
-    })
-  }
-
-  private _logDebugEvent(eventName: string, runId: string, parentRunId: string | undefined, extra: any): void {
-    if (this.debug) {
-      console.log(`Event: ${eventName}, runId: ${runId}, parentRunId: ${parentRunId}, extra:`, extra)
-    }
-  }
-
-  private _getLangchainRunName(serialized: any, ...args: any): string | undefined {
-    if (args && args.length > 0) {
-      for (const arg of args) {
-        if (arg && typeof arg === 'object' && 'name' in arg) {
-          return arg.name
-        } else if (arg && typeof arg === 'object' && 'runName' in arg) {
-          return arg.runName
-        }
-      }
-    }
-
-    if (serialized && serialized.name) {
-      return serialized.name
-    }
-    if (serialized && serialized.id) {
-      return Array.isArray(serialized.id) ? serialized.id[serialized.id.length - 1] : serialized.id
-    }
-    return undefined
-  }
-
-  private _convertLcToolCallsToOai(toolCalls: ToolCall[]): Record<string, any>[] {
-    return toolCalls.map((toolCall: ToolCall) => ({
-      type: 'function',
-      id: toolCall.id,
-      function: {
-        name: toolCall.name,
-        arguments: JSON.stringify(toolCall.args),
-      },
-    }))
-  }
-
-  private _convertMessageToDict(message: any): Record<string, any> {
-    let messageDict: Record<string, any> = {}
-
-    const messageType: string = message.getType()
-
-    switch (messageType) {
-      case 'human':
-        messageDict = { role: 'user', content: message.content }
-        break
-      case 'ai':
-        messageDict = { role: 'assistant', content: message.content }
-
-        if (message.tool_calls) {
-          messageDict.tool_calls = this._convertLcToolCallsToOai(message.tool_calls)
-        }
-
-        break
-      case 'system':
-        messageDict = { role: 'system', content: message.content }
-        break
-      case 'tool':
-        messageDict = { role: 'tool', content: message.content }
-        break
-      case 'function':
-        messageDict = { role: 'function', content: message.content }
-        break
-      default:
-        messageDict = { role: messageType, content: String(message.content) }
-        break
-    }
-
-    if (message.additional_kwargs) {
-      messageDict = { ...messageDict, ...message.additional_kwargs }
-    }
-
-    return messageDict
-  }
-
-  private _parseUsageModel(usage: any): [number, number, Record<string, any>] {
-    const conversionList: Array<[string, 'input' | 'output']> = [
-      ['promptTokens', 'input'],
-      ['completionTokens', 'output'],
-      ['input_tokens', 'input'],
-      ['output_tokens', 'output'],
-      ['prompt_token_count', 'input'],
-      ['candidates_token_count', 'output'],
-      ['inputTokenCount', 'input'],
-      ['outputTokenCount', 'output'],
-      ['input_token_count', 'input'],
-      ['generated_token_count', 'output'],
-    ]
-
-    const parsedUsage = conversionList.reduce(
-      (acc: { input: number; output: number }, [modelKey, typeKey]) => {
-        const value = usage[modelKey]
-        if (value != null) {
-          const finalCount = Array.isArray(value)
-            ? value.reduce((sum: number, tokenCount: number) => sum + tokenCount, 0)
-            : value
-          acc[typeKey] = finalCount
-        }
-        return acc
-      },
-      { input: 0, output: 0 }
-    )
-
-    // Extract additional token details like cached tokens and reasoning tokens
-    const additionalTokenData: Record<string, any> = {}
-
-    // Check for cached tokens in various formats
-    if (usage.prompt_tokens_details?.cached_tokens != null) {
-      additionalTokenData.cacheReadInputTokens = usage.prompt_tokens_details.cached_tokens
-    } else if (usage.input_token_details?.cache_read != null) {
-      additionalTokenData.cacheReadInputTokens = usage.input_token_details.cache_read
-    } else if (usage.cachedPromptTokens != null) {
-      additionalTokenData.cacheReadInputTokens = usage.cachedPromptTokens
-    }
-
-    // Check for reasoning tokens in various formats
-    if (usage.completion_tokens_details?.reasoning_tokens != null) {
-      additionalTokenData.reasoningTokens = usage.completion_tokens_details.reasoning_tokens
-    } else if (usage.output_token_details?.reasoning != null) {
-      additionalTokenData.reasoningTokens = usage.output_token_details.reasoning
-    } else if (usage.reasoningTokens != null) {
-      additionalTokenData.reasoningTokens = usage.reasoningTokens
-    }
-
-    return [parsedUsage.input, parsedUsage.output, additionalTokenData]
-  }
-
-  private parseUsage(response: LLMResult): [number, number, Record<string, any>] {
-    let llmUsage: [number, number, Record<string, any>] = [0, 0, {}]
-    const llmUsageKeys = ['token_usage', 'usage', 'tokenUsage']
-
-    if (response.llmOutput != null) {
-      const key = llmUsageKeys.find((k) => response.llmOutput?.[k] != null)
-      if (key) {
-        llmUsage = this._parseUsageModel(response.llmOutput[key])
-      }
-    }
-
-    // If top-level usage info was not found, try checking the generations.
-    if (llmUsage[0] === 0 && llmUsage[1] === 0 && response.generations) {
-      for (const generation of response.generations) {
-        for (const genChunk of generation) {
-          // Check other paths for usage information
-          if (genChunk.generationInfo?.usage_metadata) {
-            llmUsage = this._parseUsageModel(genChunk.generationInfo.usage_metadata)
-            return llmUsage
-          }
-
-          const messageChunk = genChunk.generationInfo ?? {}
-          const responseMetadata = messageChunk.response_metadata ?? {}
-          const chunkUsage =
-            responseMetadata['usage'] ??
-            responseMetadata['amazon-bedrock-invocationMetrics'] ??
-            messageChunk.usage_metadata
-          if (chunkUsage) {
-            llmUsage = this._parseUsageModel(chunkUsage)
-            return llmUsage
-          }
-        }
-      }
-    }
-
-    return llmUsage
-  }
-}

package/src/langchain/index.ts

@@ -1 +0,0 @@
-export * from './callbacks'