tes vertex JS sdk

tests/pass_through_tests/test_vertex.js

@@ -1,36 +1,30 @@
+const { VertexAI, RequestOptions } = require('@google-cloud/vertexai');
 
-const {
+const vertexAI = new VertexAI({
-    FunctionDeclarationSchemaType,
+    project: 'adroit-crow-413218',
-    HarmBlockThreshold,
+    location: 'us-central1',
-    HarmCategory,
+    apiEndpoint: "localhost:4000/vertex-ai"
-    VertexAI,
+});
-    RequestOptions
-  } = require('@google-cloud/vertexai');
 
-  const project = 'adroit-crow-413218';
+// Create customHeaders using Headers
-  const location = 'us-central1';
+const customHeaders = new Headers({
-  const textModel =  'gemini-1.0-pro';
+    "X-Litellm-Api-Key": "sk-1234"
-  const visionModel = 'gemini-1.0-pro-vision';
+});
 
+// Use customHeaders in RequestOptions
+const requestOptions = {
+    customHeaders: customHeaders
+};
 
-  const vertexAI = new VertexAI({project: project, location: location, apiEndpoint: "localhost:4000/vertex-ai"});
+const generativeModel = vertexAI.getGenerativeModel(
+    { model: 'gemini-1.0-pro' },
+    requestOptions
+);
 
-  // Instantiate Gemini models
+async function testModel() {
-  const generativeModel = vertexAI.getGenerativeModel({
+    try {
-      model: textModel,
-      // The following parameters are optional
-      // They can also be passed to individual content generation requests
-      safetySettings: [{category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE}],
-      generationConfig: {maxOutputTokens: 256},
-      systemInstruction: {
-        role: 'system',
-        parts: [{"text": `For example, you are a helpful customer service agent. tell me your name. in 5 pages`}]
-      },
-  })
-
-async function streamGenerateContent() {
         const request = {
-      contents: [{role: 'user', parts: [{text: 'How are you doing today?'}]}],
+            contents: [{role: 'user', parts: [{text: 'How are you doing today tell me your name?'}]}],
           };
         const streamingResult = await generativeModel.generateContentStream(request);
         for await (const item of streamingResult.stream) {
@@ -38,6 +32,9 @@ async function streamGenerateContent() {
         }
         const aggregatedResponse = await streamingResult.response;
         console.log('aggregated response: ', JSON.stringify(aggregatedResponse));
-  };
+    } catch (error) {
+        console.error('Error:', error);
+    }
+}
 
-  streamGenerateContent();
+testModel();