feat: vertexai instrumentation support

src/typescript/constants/common.ts

@@ -15,6 +15,7 @@ import { WeaviateFunctionNames, WeaviateFunctions } from './weaviate';
 import { TiktokenModel, TiktokenEncoding } from 'js-tiktoken';
 import { OllamaFunctionNames, OllamaFunctions } from './ollama';
 import { VercelAIFunctionNames, VercelAIFunctions } from './ai';
+import { VertexAIFunctionNames, VertexAIFunctions } from './vertexai';
 
 function getTypedKeys<T extends object>(obj: T): Array<keyof T> {
   return Object.keys(obj) as Array<keyof T>;
@@ -37,7 +38,8 @@ export const Vendors = {
   WEAVIATE: 'weaviate',
   PG: 'pg',
   VERCEL: 'ai',
-  OLLAMA: 'ollama'
+  OLLAMA: 'ollama',
+  VERTEXAI: 'vertexai',
 } as const;
 
 export enum Event {
@@ -65,6 +67,7 @@ interface VendorInstrumentationFunctions {
   pg: PgFunctions[];
   ai: VercelAIFunctions[];
   ollama: OllamaFunctions[];
+  vertexai: VertexAIFunctions[];
 }
 
 export type VendorTracedFunctions = {
@@ -84,7 +87,8 @@ export const TracedFunctionsByVendor: VendorTracedFunctions = {
   qdrant: QdrantFunctionNames,
   weaviate: WeaviateFunctionNames,
   ai: VercelAIFunctionNames,
-  ollama: OllamaFunctionNames
+  ollama: OllamaFunctionNames,
+  vertexai: VertexAIFunctionNames,
 } as const;
 
 export const TIKTOKEN_MODEL_MAPPING: Record<TiktokenModel | string, TiktokenEncoding> = {

src/typescript/constants/vertexai.ts

@@ -0,0 +1,32 @@
+/*
+ * 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.
+ */
+
+export const APIS = {
+	GENERATE_CONTENT: {
+		METHOD: "vertexAI.GenerativeModel.generateContent",
+  },
+	GENERATE_CONTENT_STREAM: {
+		METHOD: "vertexAI.GenerativeModel.generateContentStream",
+	},
+	SEND_MESSAGE: {
+		METHOD: "vertexAI.ChatSession.sendMessage",
+	},
+	SEND_MESSAGE_STREAM: {
+		METHOD: "vertexAI.ChatSession.sendMessageStream",
+	},
+} as const;
+export type VertexAIFunctions = typeof APIS[keyof typeof APIS]['METHOD']
+export const VertexAIFunctionNames: VertexAIFunctions[] = Object.values(APIS).map((api) => api.METHOD)

src/typescript/index.ts

@@ -14,6 +14,7 @@ import { APIS as PineConeAPIS } from "./constants/pinecone";
 import { APIS as QdrantAPIS } from "./constants/qdrant";
 import { APIS as vercelAIAPIS } from "./constants/ai";
 import { APIS as ollamaAPIS } from "./constants/ollama";
+import { APIS as vertexAIAPIS } from "./constants/vertexai";
 import { queryTypeToFunctionToProps } from "./constants/weaviate";
 import { TIKTOKEN_MODEL_MAPPING } from "./constants/common";
 
@@ -50,7 +51,8 @@ const APIS = {
   pinecone: PineConeAPIS,
   qdrant: QdrantAPIS,
   ai: vercelAIAPIS,
-  ollama: ollamaAPIS
+  ollama: ollamaAPIS,
+  vertexai: vertexAIAPIS,
 }
 export {
   LLMSpanAttributeNames,

src/typescript/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@langtrase/trace-attributes",
-    "version": "7.2.0",
+    "version": "7.3.0",
     "description": "LangTrace - Trace Attributes",
     "main": "dist/index.js",
     "types": "dist/index.d.ts",
@@ -19,4 +19,4 @@
         "ncp": "^2.0.0",
         "typescript": "^5.5.3"
     }
-}
+}