fix(vertex-ai): missing system prompt (#473)

traceloop · Nov 13, 2024 · 663e438 · 663e438
1 parent 23fdb28
commit 663e438
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 14 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/packages/instrumentation-vertexai/src/vertexai-instrumentation.ts b/packages/instrumentation-vertexai/src/vertexai-instrumentation.ts
@@ -157,10 +157,22 @@ export class VertexAIInstrumentation extends InstrumentationBase {
       }
 
       if (this._shouldSendPrompts() && "contents" in params) {
-        attributes[`${SpanAttributes.LLM_PROMPTS}.0.role`] =
-          params.contents[0].role ?? "user";
-        attributes[`${SpanAttributes.LLM_PROMPTS}.0.content`] =
-          this._formatPartsData(params.contents[0].parts);
+        let i = 0;
+
+        if (instance["systemInstruction"]) {
+          attributes[`${SpanAttributes.LLM_PROMPTS}.${i}.role`] = "system";
+          attributes[`${SpanAttributes.LLM_PROMPTS}.${i}.content`] =
+            this._formatPartsData(instance["systemInstruction"].parts);
+
+          i++;
+        }
+
+        params.contents.forEach((content, j) => {
+          attributes[`${SpanAttributes.LLM_PROMPTS}.${i + j}.role`] =
+            content.role ?? "user";
+          attributes[`${SpanAttributes.LLM_PROMPTS}.${i + j}.content`] =
+            this._formatPartsData(content.parts);
+        });
       }
     } catch (e) {
       this._diag.debug(e);

diff --git a/packages/sample-app/package.json b/packages/sample-app/package.json
@@ -36,8 +36,8 @@
     "@anthropic-ai/sdk": "^0.27.1",
     "@aws-sdk/client-bedrock-runtime": "^3.499.0",
     "@azure/openai": "^1.0.0-beta.11",
-    "@google-cloud/aiplatform": "^3.10.0",
-    "@google-cloud/vertexai": "^1.2.0",
+    "@google-cloud/aiplatform": "^3.32.0",
+    "@google-cloud/vertexai": "^1.9.0",
     "@langchain/community": "^0.2.31",
     "@pinecone-database/pinecone": "^2.0.1",
     "@traceloop/node-server-sdk": "*",

diff --git a/packages/sample-app/src/vertexai/gemini.ts b/packages/sample-app/src/vertexai/gemini.ts
@@ -20,6 +20,10 @@ async function createNonStreamingContent() {
       // Instantiate the model
       const generativeModel = vertexAI.getGenerativeModel({
         model: "gemini-1.5-flash",
+        systemInstruction: {
+          role: "system",
+          parts: [{ text: "You are a helpful assistant" }],
+        },
       });
 
       const request = {
@@ -53,6 +57,10 @@ async function createStreamingContent() {
       // Instantiate the model
       const generativeModel = vertexAI.getGenerativeModel({
         model: "gemini-1.5-flash",
+        systemInstruction: {
+          role: "system",
+          parts: [{ text: "You are a helpful assistant" }],
+        },
       });
 
       const request = {