feat: adding gemini support (#162)

* feat: adding gemini support

* chore: update package lock

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@langtrase/typescript-sdk",
-  "version": "5.1.2",
+  "version": "5.2.0",
   "description": "A typescript SDK for Langtrace",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
@@ -21,7 +21,7 @@
   "author": "Scale3 Labs",
   "license": "Apache-2.0",
   "dependencies": {
-    "@langtrase/trace-attributes": "7.1.2",
+    "@langtrase/trace-attributes": "7.2.0",
     "@opentelemetry/api": "^1.7.0",
     "@opentelemetry/instrumentation": "^0.49.1",
     "@opentelemetry/sdk-trace-base": "^1.22.0",
@@ -36,6 +36,7 @@
   "devDependencies": {
     "@ai-sdk/anthropic": "^0.0.41",
     "@ai-sdk/openai": "^0.0.34",
+    "@google/generative-ai": "^0.1.3",
     "@anthropic-ai/sdk": "^0.24.3",
     "@changesets/cli": "^2.27.1",
     "@opentelemetry/exporter-trace-otlp-http": "^0.52.1",
@@ -61,4 +62,4 @@
     "typescript": "^4.9.5",
     "weaviate-ts-client": "^2.2.0"
   }
-}
+}

src/examples/gemini/basic.ts

@@ -0,0 +1,36 @@
+import { init } from '@langtrace-init/init'
+import { GoogleGenerativeAI } from '@google/generative-ai'
+
+import dotenv from 'dotenv'
+dotenv.config()
+init({ batch: false, write_spans_to_console: true })
+
+if (process.env.GEMINI_API_KEY === undefined) {
+  throw new Error('GEMINI_API_KEY is not defined')
+}
+const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY)
+
+export const basicGeminiChat = async (): Promise<void> => {
+  const model = genAI.getGenerativeModel({
+    model: 'gemini-1.5-flash',
+    generationConfig: {
+      candidateCount: 1,
+      maxOutputTokens: 20,
+      temperature: 1.0
+    }
+  })
+  const prompt = 'Say hi twice'
+  const result = await model.generateContent(prompt)
+  console.log(result.response.text())
+}
+
+export const basicGeminiChatStream = async (): Promise<void> => {
+  const model = genAI.getGenerativeModel({ model: 'gemini-1.5-flash' })
+  const prompt = 'say hi 3 times'
+  const result = await model.generateContentStream(prompt)
+
+  for await (const chunk of result.stream) {
+    const chunkText = chunk.text()
+    process.stdout.write(chunkText)
+  }
+}

src/init/init.ts

@@ -26,6 +26,7 @@ import { LANGTRACE_REMOTE_URL } from '@langtrace-constants/exporter/langtrace_ex
 import { anthropicInstrumentation } from '@langtrace-instrumentation/anthropic/instrumentation'
 import { chromaInstrumentation } from '@langtrace-instrumentation/chroma/instrumentation'
 import { cohereInstrumentation } from '@langtrace-instrumentation/cohere/instrumentation'
+import { geminiInstrumentation } from '@langtrace-instrumentation/gemini/instrumentation'
 import { groqInstrumentation } from '@langtrace-instrumentation/groq/instrumentation'
 import { llamaIndexInstrumentation } from '@langtrace-instrumentation/llamaindex/instrumentation'
 import { openAIInstrumentation } from '@langtrace-instrumentation/openai/instrumentation'
@@ -151,6 +152,7 @@ export const init: LangTraceInit = ({
     openai: openAIInstrumentation,
     cohere: cohereInstrumentation,
     anthropic: anthropicInstrumentation,
+    gemini: geminiInstrumentation,
     groq: groqInstrumentation,
     pinecone: pineconeInstrumentation,
     llamaindex: llamaIndexInstrumentation,

src/instrumentation/gemini/instrumentation.ts

@@ -0,0 +1,72 @@
+/*
+ * Copyright (c) 2024 Scale3 Labs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { generateContentPatch } from '@langtrace-instrumentation/gemini/patch'
+import { diag } from '@opentelemetry/api'
+import { InstrumentationBase, InstrumentationNodeModuleDefinition, isWrapped } from '@opentelemetry/instrumentation'
+// eslint-disable-next-line no-restricted-imports
+import { name, version } from '../../../package.json'
+
+class GeminiInstrumentation extends InstrumentationBase<any> {
+  constructor () {
+    super(name, version)
+  }
+
+  init (): Array<InstrumentationNodeModuleDefinition<any>> {
+    diag.debug('init GeminiInstrumentation')
+    const module = new InstrumentationNodeModuleDefinition<any>(
+      '@google/generative-ai',
+      ['>=0.1.3'],
+      (moduleExports, moduleVersion) => {
+        diag.debug(`Patching Gemini SDK version ${moduleVersion}`)
+        this._patch(moduleExports, moduleVersion as string)
+        return moduleExports
+      },
+      (moduleExports, moduleVersion) => {
+        diag.debug(`Unpatching Gemini SDK version ${moduleVersion}`)
+        if (moduleExports !== undefined) {
+          this._unpatch(moduleExports)
+        }
+      }
+    )
+    return [module]
+  }
+
+  private _patch (gemini: any, moduleVersion?: string): void {
+    if (isWrapped(gemini.GenerativeModel.prototype)) {
+      this._unpatch(gemini)
+    }
+
+    this._wrap(gemini.GenerativeModel.prototype,
+      'generateContent',
+      (originalMethod: (...args: any[]) => any) =>
+        generateContentPatch(originalMethod, this.tracer, this.instrumentationVersion, name, moduleVersion)
+    )
+
+    this._wrap(gemini.GenerativeModel.prototype,
+      'generateContentStream',
+      (originalMethod: (...args: any[]) => any) =>
+        generateContentPatch(originalMethod, this.tracer, this.instrumentationVersion, name, moduleVersion)
+    )
+  }
+
+  private _unpatch (gemini: any): void {
+    this._unwrap(gemini.GoogleGenerativeAI.GenerativeModel, 'generateContent')
+    this._unwrap(gemini.GoogleGenerativeAI.GenerativeModel, 'generateContentStream')
+  }
+}
+
+export const geminiInstrumentation = new GeminiInstrumentation()