@posthog/ai 4.4.0 → 5.0.0

package/lib/index.d.ts

@@ -6,6 +6,7 @@ import { LanguageModelV1 } from 'ai';
 import AnthropicOriginal from '@anthropic-ai/sdk';
 import { RequestOptions as RequestOptions$1, APIPromise as APIPromise$1 } from '@anthropic-ai/sdk/core';
 import { Stream as Stream$1 } from '@anthropic-ai/sdk/streaming';
+import { GoogleGenAI } from '@google/genai';
 import { BaseCallbackHandler } from '@langchain/core/callbacks/base';
 import { Serialized } from '@langchain/core/load/serializable';
 import { ChainValues } from '@langchain/core/utils/types';
@@ -123,6 +124,46 @@ declare class WrappedMessages extends AnthropicOriginal.Messages {
     create(body: MessageCreateParamsBase & MonitoringParams, options?: RequestOptions$1): APIPromise$1<Stream$1<RawMessageStreamEvent> | Message>;
 }
 
+type GenerateContentRequest = {
+    model: string;
+    contents: any;
+    config?: any;
+    [key: string]: any;
+};
+type GenerateContentResponse = {
+    text?: string;
+    candidates?: any[];
+    usageMetadata?: {
+        promptTokenCount?: number;
+        candidatesTokenCount?: number;
+        totalTokenCount?: number;
+    };
+    [key: string]: any;
+};
+interface MonitoringGeminiConfig {
+    apiKey?: string;
+    vertexai?: boolean;
+    project?: string;
+    location?: string;
+    apiVersion?: string;
+    posthog: PostHog;
+}
+declare class PostHogGoogleGenAI {
+    private readonly phClient;
+    private readonly client;
+    models: WrappedModels;
+    constructor(config: MonitoringGeminiConfig);
+}
+declare class WrappedModels {
+    private readonly phClient;
+    private readonly client;
+    constructor(client: GoogleGenAI, phClient: PostHog);
+    generateContent(params: GenerateContentRequest & MonitoringParams): Promise<GenerateContentResponse>;
+    generateContentStream(params: GenerateContentRequest & MonitoringParams): AsyncGenerator<any, void, unknown>;
+    private formatInput;
+    private formatOutput;
+}
+
 declare class LangChainCallbackHandler extends BaseCallbackHandler {
     name: string;
     private client;
@@ -181,4 +222,4 @@ declare class LangChainCallbackHandler extends BaseCallbackHandler {
     private parseUsage;
 }
 
-export { PostHogAnthropic as Anthropic, PostHogAzureOpenAI as AzureOpenAI, LangChainCallbackHandler, PostHogOpenAI as OpenAI, wrapVercelLanguageModel as withTracing };
+export { PostHogAnthropic as Anthropic, PostHogAzureOpenAI as AzureOpenAI, PostHogGoogleGenAI as GoogleGenAI, LangChainCallbackHandler, PostHogOpenAI as OpenAI, wrapVercelLanguageModel as withTracing };

package/lib/index.mjs

@@ -4,6 +4,7 @@ import { v4 } from 'uuid';
 import { Buffer } from 'buffer';
 import { experimental_wrapLanguageModel } from 'ai';
 import AnthropicOriginal from '@anthropic-ai/sdk';
+import { GoogleGenAI } from '@google/genai';
 
 // limit large outputs by truncating to 200kb (approx 200k bytes)
 const MAX_OUTPUT_SIZE = 200000;
@@ -1037,6 +1038,234 @@ class WrappedMessages extends AnthropicOriginal.Messages {
   }
 }
 
+class PostHogGoogleGenAI {
+  constructor(config) {
+    const {
+      posthog,
+      ...geminiConfig
+    } = config;
+    this.phClient = posthog;
+    this.client = new GoogleGenAI(geminiConfig);
+    this.models = new WrappedModels(this.client, this.phClient);
+  }
+}
+class WrappedModels {
+  constructor(client, phClient) {
+    this.client = client;
+    this.phClient = phClient;
+  }
+  async generateContent(params) {
+    const {
+      posthogDistinctId,
+      posthogTraceId,
+      posthogProperties,
+      posthogGroups,
+      posthogCaptureImmediate,
+      ...geminiParams
+    } = params;
+    const traceId = posthogTraceId ?? v4();
+    const startTime = Date.now();
+    try {
+      const response = await this.client.models.generateContent(geminiParams);
+      const latency = (Date.now() - startTime) / 1000;
+      await sendEventToPosthog({
+        client: this.phClient,
+        distinctId: posthogDistinctId ?? traceId,
+        traceId,
+        model: geminiParams.model,
+        provider: 'gemini',
+        input: this.formatInput(geminiParams.contents),
+        output: this.formatOutput(response),
+        latency,
+        baseURL: 'https://generativelanguage.googleapis.com',
+        params: params,
+        httpStatus: 200,
+        usage: {
+          inputTokens: response.usageMetadata?.promptTokenCount ?? 0,
+          outputTokens: response.usageMetadata?.candidatesTokenCount ?? 0
+        },
+        captureImmediate: posthogCaptureImmediate
+      });
+      return response;
+    } catch (error) {
+      const latency = (Date.now() - startTime) / 1000;
+      await sendEventToPosthog({
+        client: this.phClient,
+        distinctId: posthogDistinctId ?? traceId,
+        traceId,
+        model: geminiParams.model,
+        provider: 'gemini',
+        input: this.formatInput(geminiParams.contents),
+        output: [],
+        latency,
+        baseURL: 'https://generativelanguage.googleapis.com',
+        params: params,
+        httpStatus: error?.status ?? 500,
+        usage: {
+          inputTokens: 0,
+          outputTokens: 0
+        },
+        isError: true,
+        error: JSON.stringify(error),
+        captureImmediate: posthogCaptureImmediate
+      });
+      throw error;
+    }
+  }
+  async *generateContentStream(params) {
+    const {
+      posthogDistinctId,
+      posthogTraceId,
+      posthogProperties,
+      posthogGroups,
+      posthogCaptureImmediate,
+      ...geminiParams
+    } = params;
+    const traceId = posthogTraceId ?? v4();
+    const startTime = Date.now();
+    let accumulatedContent = '';
+    let usage = {
+      inputTokens: 0,
+      outputTokens: 0
+    };
+    try {
+      const stream = await this.client.models.generateContentStream(geminiParams);
+      for await (const chunk of stream) {
+        if (chunk.text) {
+          accumulatedContent += chunk.text;
+        }
+        if (chunk.usageMetadata) {
+          usage = {
+            inputTokens: chunk.usageMetadata.promptTokenCount ?? 0,
+            outputTokens: chunk.usageMetadata.candidatesTokenCount ?? 0
+          };
+        }
+        yield chunk;
+      }
+      const latency = (Date.now() - startTime) / 1000;
+      await sendEventToPosthog({
+        client: this.phClient,
+        distinctId: posthogDistinctId ?? traceId,
+        traceId,
+        model: geminiParams.model,
+        provider: 'gemini',
+        input: this.formatInput(geminiParams.contents),
+        output: [{
+          content: accumulatedContent,
+          role: 'assistant'
+        }],
+        latency,
+        baseURL: 'https://generativelanguage.googleapis.com',
+        params: params,
+        httpStatus: 200,
+        usage,
+        captureImmediate: posthogCaptureImmediate
+      });
+    } catch (error) {
+      const latency = (Date.now() - startTime) / 1000;
+      await sendEventToPosthog({
+        client: this.phClient,
+        distinctId: posthogDistinctId ?? traceId,
+        traceId,
+        model: geminiParams.model,
+        provider: 'gemini',
+        input: this.formatInput(geminiParams.contents),
+        output: [],
+        latency,
+        baseURL: 'https://generativelanguage.googleapis.com',
+        params: params,
+        httpStatus: error?.status ?? 500,
+        usage: {
+          inputTokens: 0,
+          outputTokens: 0
+        },
+        isError: true,
+        error: JSON.stringify(error),
+        captureImmediate: posthogCaptureImmediate
+      });
+      throw error;
+    }
+  }
+  formatInput(contents) {
+    if (typeof contents === 'string') {
+      return [{
+        role: 'user',
+        content: contents
+      }];
+    }
+    if (Array.isArray(contents)) {
+      return contents.map(item => {
+        if (typeof item === 'string') {
+          return {
+            role: 'user',
+            content: item
+          };
+        }
+        if (item && typeof item === 'object') {
+          if (item.text) {
+            return {
+              role: item.role || 'user',
+              content: item.text
+            };
+          }
+          if (item.content) {
+            return {
+              role: item.role || 'user',
+              content: item.content
+            };
+          }
+        }
+        return {
+          role: 'user',
+          content: String(item)
+        };
+      });
+    }
+    if (contents && typeof contents === 'object') {
+      if (contents.text) {
+        return [{
+          role: 'user',
+          content: contents.text
+        }];
+      }
+      if (contents.content) {
+        return [{
+          role: 'user',
+          content: contents.content
+        }];
+      }
+    }
+    return [{
+      role: 'user',
+      content: String(contents)
+    }];
+  }
+  formatOutput(response) {
+    if (response.text) {
+      return [{
+        role: 'assistant',
+        content: response.text
+      }];
+    }
+    if (response.candidates && Array.isArray(response.candidates)) {
+      return response.candidates.map(candidate => {
+        if (candidate.content && candidate.content.parts) {
+          const text = candidate.content.parts.filter(part => part.text).map(part => part.text).join('');
+          return {
+            role: 'assistant',
+            content: text
+          };
+        }
+        return {
+          role: 'assistant',
+          content: String(candidate)
+        };
+      });
+    }
+    return [];
+  }
+}
+
 var decamelize = function (str, sep) {
     if (typeof str !== 'string') {
         throw new TypeError('Expected a string');
@@ -1960,5 +2189,5 @@ class LangChainCallbackHandler extends BaseCallbackHandler {
   }
 }
 
-export { PostHogAnthropic as Anthropic, PostHogAzureOpenAI as AzureOpenAI, LangChainCallbackHandler, PostHogOpenAI as OpenAI, wrapVercelLanguageModel as withTracing };
+export { PostHogAnthropic as Anthropic, PostHogAzureOpenAI as AzureOpenAI, PostHogGoogleGenAI as GoogleGenAI, LangChainCallbackHandler, PostHogOpenAI as OpenAI, wrapVercelLanguageModel as withTracing };
 //# sourceMappingURL=index.mjs.map