(feat) Add support for using @google/generative-ai JS with LiteLLM Proxy (#6899)

* feat - allow using gemini js SDK with LiteLLM

* add auth for gemini_proxy_route

* basic local test for js

* test cost tagging gemini js requests

* add js sdk test for gemini with litellm

* add docs on gemini JS SDK

* run node.js tests

* fix google ai studio tests

* fix vertex js spend test

tests/pass_through_tests/test_gemini_with_spend.test.js

@@ -0,0 +1,123 @@
+const { GoogleGenerativeAI } = require("@google/generative-ai");
+const fs = require('fs');
+const path = require('path');
+
+// Import fetch if the SDK uses it
+const originalFetch = global.fetch || require('node-fetch');
+
+let lastCallId;
+
+// Monkey-patch the fetch used internally
+global.fetch = async function patchedFetch(url, options) {
+    const response = await originalFetch(url, options);
+    
+    // Store the call ID if it exists
+    lastCallId = response.headers.get('x-litellm-call-id');
+        
+    return response;
+};
+
+describe('Gemini AI Tests', () => {
+    test('should successfully generate non-streaming content with tags', async () => {
+        const genAI = new GoogleGenerativeAI("sk-1234"); // litellm proxy API key
+
+        const requestOptions = {
+            baseUrl: 'http://127.0.0.1:4000/gemini',
+            customHeaders: {
+                "tags": "gemini-js-sdk,pass-through-endpoint"
+            }
+        };
+
+        const model = genAI.getGenerativeModel({
+            model: 'gemini-pro'
+        }, requestOptions);
+
+        const prompt = 'Say "hello test" and nothing else';
+
+        const result = await model.generateContent(prompt);
+        expect(result).toBeDefined();
+        
+        // Use the captured callId
+        const callId = lastCallId;
+        console.log("Captured Call ID:", callId);
+
+        // Wait for spend to be logged
+        await new Promise(resolve => setTimeout(resolve, 15000));
+
+        // Check spend logs
+        const spendResponse = await fetch(
+            `http://127.0.0.1:4000/spend/logs?request_id=${callId}`,
+            {
+                headers: {
+                    'Authorization': 'Bearer sk-1234'
+                }
+            }
+        );
+        
+        const spendData = await spendResponse.json();
+        console.log("spendData", spendData)
+        expect(spendData).toBeDefined();
+        expect(spendData[0].request_id).toBe(callId);
+        expect(spendData[0].call_type).toBe('pass_through_endpoint');
+        expect(spendData[0].request_tags).toEqual(['gemini-js-sdk', 'pass-through-endpoint']);
+        expect(spendData[0].metadata).toHaveProperty('user_api_key');
+        expect(spendData[0].model).toContain('gemini');
+        expect(spendData[0].spend).toBeGreaterThan(0);
+    }, 25000);
+
+    test('should successfully generate streaming content with tags', async () => {
+        const genAI = new GoogleGenerativeAI("sk-1234"); // litellm proxy API key
+
+        const requestOptions = {
+            baseUrl: 'http://127.0.0.1:4000/gemini',
+            customHeaders: {
+                "tags": "gemini-js-sdk,pass-through-endpoint"
+            }
+        };
+
+        const model = genAI.getGenerativeModel({
+            model: 'gemini-pro'
+        }, requestOptions);
+
+        const prompt = 'Say "hello test" and nothing else';
+
+        const streamingResult = await model.generateContentStream(prompt);
+        expect(streamingResult).toBeDefined();
+
+        for await (const chunk of streamingResult.stream) {
+            console.log('stream chunk:', JSON.stringify(chunk));
+            expect(chunk).toBeDefined();
+        }
+
+        const aggregatedResponse = await streamingResult.response;
+        console.log('aggregated response:', JSON.stringify(aggregatedResponse));
+        expect(aggregatedResponse).toBeDefined();
+
+        // Use the captured callId
+        const callId = lastCallId;
+        console.log("Captured Call ID:", callId);
+
+        // Wait for spend to be logged
+        await new Promise(resolve => setTimeout(resolve, 15000));
+
+        // Check spend logs
+        const spendResponse = await fetch(
+            `http://127.0.0.1:4000/spend/logs?request_id=${callId}`,
+            {
+                headers: {
+                    'Authorization': 'Bearer sk-1234'
+                }
+            }
+        );
+        
+        const spendData = await spendResponse.json();
+        console.log("spendData", spendData)
+        expect(spendData).toBeDefined();
+        expect(spendData[0].request_id).toBe(callId);
+        expect(spendData[0].call_type).toBe('pass_through_endpoint');
+        expect(spendData[0].request_tags).toEqual(['gemini-js-sdk', 'pass-through-endpoint']);
+        expect(spendData[0].metadata).toHaveProperty('user_api_key');
+        expect(spendData[0].model).toContain('gemini');
+        expect(spendData[0].spend).toBeGreaterThan(0);
+    }, 25000);
+});

tests/pass_through_tests/test_local_gemini.js

@@ -0,0 +1,55 @@
+const { GoogleGenerativeAI, ModelParams, RequestOptions } = require("@google/generative-ai");
+
+const modelParams = {
+    model: 'gemini-pro',
+};
+  
+const requestOptions = {
+    baseUrl: 'http://127.0.0.1:4000/gemini',
+    customHeaders: {
+        "tags": "gemini-js-sdk,gemini-pro"
+    }
+};
+  
+const genAI = new GoogleGenerativeAI("sk-1234"); // litellm proxy API key
+const model = genAI.getGenerativeModel(modelParams, requestOptions);
+
+const testPrompt = "Explain how AI works";
+
+async function main() {
+  console.log("making request")
+  try {
+    const result = await model.generateContent(testPrompt);
+    console.log(result.response.text());
+  } catch (error) {
+    console.error('Error details:', {
+      name: error.name,
+      message: error.message,
+      cause: error.cause,
+      // Check if there's a network error
+      isNetworkError: error instanceof TypeError && error.message === 'fetch failed'
+    });
+    
+    // Check if the server is running
+    if (error instanceof TypeError && error.message === 'fetch failed') {
+      console.error('Make sure your local server is running at http://localhost:4000');
+    }
+  }
+}
+
+
+async function main_streaming() {
+    try {
+        const streamingResult = await model.generateContentStream(testPrompt);
+        for await (const item of streamingResult.stream) {
+            console.log('stream chunk: ', JSON.stringify(item));
+        }
+        const aggregatedResponse = await streamingResult.response;
+        console.log('aggregated response: ', JSON.stringify(aggregatedResponse));
+    } catch (error) {
+        console.error('Error details:', error);
+    }
+}
+
+// main();
+main_streaming();

tests/pass_through_tests/test_vertex_with_spend.test.js

@@ -60,9 +60,9 @@ function loadVertexAiCredentials() {
 }
 
 // Run credential loading before tests
-// beforeAll(() => {
-//     loadVertexAiCredentials();
-// });
+beforeAll(() => {
+    loadVertexAiCredentials();
+});