Spaces:
Configuration error
Configuration error
File size: 4,535 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 |
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-1.5-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].custom_llm_provider).toBe('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-1.5-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);
expect(spendData[0].custom_llm_provider).toBe('gemini');
}, 25000);
});
|