Add support to js client for LANGCHAIN_ env variable tracking in meta…

…data and revision_id in createRun (#367)
langchain-ai · Jan 19, 2024 · 969a453 · 969a453
2 parents 5b447ac + 9c38a61
commit 969a453
Show file tree

Hide file tree

Showing 4 changed files with 157 additions and 5 deletions.
diff --git a/js/src/client.ts b/js/src/client.ts
@@ -23,7 +23,11 @@ import {
   convertLangChainMessageToExample,
   isLangChainMessage,
 } from "./utils/messages.js";
-import { getEnvironmentVariable, getRuntimeEnvironment } from "./utils/env.js";
+import {
+  getEnvironmentVariable,
+  getLangChainEnvVarsMetadata,
+  getRuntimeEnvironment,
+} from "./utils/env.js";
 
 import { RunEvaluator } from "./evaluation/evaluator.js";
 import { __version__ } from "./index.js";
@@ -99,6 +103,7 @@ interface CreateRunParams {
   child_runs?: RunCreate[];
   parent_run_id?: string;
   project_name?: string;
+  revision_id?: string;
 }
 
 interface projectOptions {
@@ -349,7 +354,9 @@ export class Client {
   public async createRun(run: CreateRunParams): Promise<void> {
     const headers = { ...this.headers, "Content-Type": "application/json" };
     const extra = run.extra ?? {};
+    const metadata = extra.metadata;
     const runtimeEnv = await getRuntimeEnvironment();
+    const envVars = getLangChainEnvVarsMetadata();
     const session_name = run.project_name;
     delete run.project_name;
     const runCreate: RunCreate = {
@@ -361,6 +368,13 @@ export class Client {
           ...runtimeEnv,
           ...extra.runtime,
         },
+        metadata: {
+          ...envVars,
+          ...(envVars.revision_id || run.revision_id
+            ? { revision_id: run.revision_id ?? envVars.revision_id }
+            : {}),
+          ...metadata,
+        },
       },
     };
     runCreate.inputs = hideInputs(runCreate.inputs);

diff --git a/js/src/tests/client.int.test.ts b/js/src/tests/client.int.test.ts
@@ -507,6 +507,56 @@ test("Test create run with masked inputs/outputs", async () => {
   expect(Object.keys(run2.outputs ?? {})).toHaveLength(0);
 }, 10000);
 
+test("Test create run with revision id", async () => {
+  const langchainClient = new Client({
+    apiUrl: "http://localhost:1984",
+  });
+  // eslint-disable-next-line no-process-env
+  process.env.LANGCHAIN_REVISION_ID = "test_revision_id";
+  // eslint-disable-next-line no-process-env
+  process.env.LANGCHAIN_API_KEY = "fake_api_key";
+  // eslint-disable-next-line no-process-env
+  process.env.LANGCHAIN_OTHER_KEY = "test_other_key";
+  const projectName = "__test_create_run_with_revision_id";
+  const projects = langchainClient.listProjects();
+  for await (const project of projects) {
+    if (project.name === projectName) {
+      await langchainClient.deleteProject({ projectName });
+    }
+  }
+  const runId = "0cc29488-3b1b-4151-9476-30b5c1b24883";
+  await langchainClient.createRun({
+    id: runId,
+    project_name: projectName,
+    name: "test_run",
+    run_type: "llm",
+    inputs: { prompt: "hello world" },
+    outputs: { generation: "hi there" },
+    start_time: new Date().getTime(),
+    end_time: new Date().getTime(),
+  });
+
+  const runId2 = "82f19ed3-256f-4571-a078-2ccf11d0eba3";
+  await langchainClient.createRun({
+    id: runId2,
+    project_name: projectName,
+    name: "test_run_2",
+    run_type: "llm",
+    inputs: { messages: "hello world 2" },
+    start_time: new Date().getTime(),
+    revision_id: "different_revision_id",
+  });
+
+  const run1 = await langchainClient.readRun(runId);
+  expect(run1.extra?.metadata?.revision_id).toEqual("test_revision_id");
+  expect(run1.extra?.metadata.LANGCHAIN_OTHER_KEY).toEqual("test_other_key");
+  expect(run1.extra?.metadata).not.toHaveProperty("LANGCHAIN_API_KEY");
+  const run2 = await langchainClient.readRun(runId2);
+  expect(run2.extra?.metadata?.revision_id).toEqual("different_revision_id");
+  expect(run2.extra?.metadata.LANGCHAIN_OTHER_KEY).toEqual("test_other_key");
+  expect(run2.extra?.metadata).not.toHaveProperty("LANGCHAIN_API_KEY");
+}, 10000);
+
 describe("createChatExample", () => {
   it("should convert LangChainBaseMessage objects to examples", async () => {
     const langchainClient = new Client({

diff --git a/js/src/tests/client.test.ts b/js/src/tests/client.test.ts
@@ -1,5 +1,10 @@
 import { jest } from "@jest/globals";
 import { Client } from "../client.js";
+import {
+  getEnvironmentVariables,
+  getLangChainEnvVars,
+  getLangChainEnvVarsMetadata,
+} from "../utils/env.js";
 
 describe("Client", () => {
   describe("createLLMExample", () => {
@@ -118,4 +123,46 @@ describe("Client", () => {
       expect(result).toBe("https://smith.langchain.com");
     });
   });
+
+  describe("env functions", () => {
+    it("should return the env variables correctly", async () => {
+      // eslint-disable-next-line no-process-env
+      process.env.LANGCHAIN_REVISION_ID = "test_revision_id";
+      // eslint-disable-next-line no-process-env
+      process.env.LANGCHAIN_API_KEY = "fake_api_key";
+      // eslint-disable-next-line no-process-env
+      process.env.LANGCHAIN_OTHER_KEY = "test_other_key";
+      // eslint-disable-next-line no-process-env
+      process.env.LANGCHAIN_OTHER_NON_SENSITIVE_METADATA = "test_some_metadata";
+      // eslint-disable-next-line no-process-env
+      process.env.LANGCHAIN_ENDPOINT = "https://example.com";
+      // eslint-disable-next-line no-process-env
+      process.env.SOME_RANDOM_THING = "random";
+
+      const envVars = getEnvironmentVariables();
+      const langchainEnvVars = getLangChainEnvVars();
+      const langchainMetadataEnvVars = getLangChainEnvVarsMetadata();
+
+      expect(envVars).toMatchObject({
+        LANGCHAIN_REVISION_ID: "test_revision_id",
+        LANGCHAIN_API_KEY: "fake_api_key",
+        LANGCHAIN_OTHER_KEY: "test_other_key",
+        LANGCHAIN_ENDPOINT: "https://example.com",
+        SOME_RANDOM_THING: "random",
+        LANGCHAIN_OTHER_NON_SENSITIVE_METADATA: "test_some_metadata",
+      });
+      expect(langchainEnvVars).toMatchObject({
+        LANGCHAIN_REVISION_ID: "test_revision_id",
+        LANGCHAIN_API_KEY: "fa********ey",
+        LANGCHAIN_OTHER_KEY: "te**********ey",
+        LANGCHAIN_ENDPOINT: "https://example.com",
+        LANGCHAIN_OTHER_NON_SENSITIVE_METADATA: "test_some_metadata",
+      });
+      expect(langchainEnvVars).not.toHaveProperty("SOME_RANDOM_THING");
+      expect(langchainMetadataEnvVars).toEqual({
+        revision_id: "test_revision_id",
+        LANGCHAIN_OTHER_NON_SENSITIVE_METADATA: "test_some_metadata",
+      });
+    });
+  });
 });
diff --git a/js/src/utils/env.ts b/js/src/utils/env.ts
@@ -86,7 +86,7 @@ export async function getRuntimeEnvironment(): Promise<RuntimeEnvironment> {
 
 /**
  * Retrieves the LangChain-specific environment variables from the current runtime environment.
- * Sensitive keys (containing the word "key") have their values redacted for security.
+ * Sensitive keys (containing the word "key", "token", or "secret") have their values redacted for security.
  *
  * @returns {Record<string, string>}
  *  - A record of LangChain-specific environment variables.
@@ -102,7 +102,12 @@ export function getLangChainEnvVars(): Record<string, string> {
   }
 
   for (const key in envVars) {
-    if (key.toLowerCase().includes("key") && typeof envVars[key] === "string") {
+    if (
+      (key.toLowerCase().includes("key") ||
+        key.toLowerCase().includes("secret") ||
+        key.toLowerCase().includes("token")) &&
+      typeof envVars[key] === "string"
+    ) {
       const value = envVars[key];
       envVars[key] =
         value.slice(0, 2) + "*".repeat(value.length - 4) + value.slice(-2);
@@ -112,6 +117,43 @@ export function getLangChainEnvVars(): Record<string, string> {
   return envVars;
 }
 
+/**
+ * Retrieves the LangChain-specific metadata from the current runtime environment.
+ *
+ * @returns {Record<string, string>}
+ *  - A record of LangChain-specific metadata environment variables.
+ */
+export function getLangChainEnvVarsMetadata(): Record<string, string> {
+  const allEnvVars = getEnvironmentVariables() || {};
+  const envVars: Record<string, string> = {};
+  const excluded = [
+    "LANGCHAIN_API_KEY",
+    "LANGCHAIN_ENDPOINT",
+    "LANGCHAIN_TRACING_V2",
+    "LANGCHAIN_PROJECT",
+    "LANGCHAIN_SESSION",
+  ];
+
+  for (const [key, value] of Object.entries(allEnvVars)) {
+    if (
+      key.startsWith("LANGCHAIN_") &&
+      typeof value === "string" &&
+      !excluded.includes(key) &&
+      !key.toLowerCase().includes("key") &&
+      !key.toLowerCase().includes("secret") &&
+      !key.toLowerCase().includes("token")
+    ) {
+      if (key === "LANGCHAIN_REVISION_ID") {
+        envVars["revision_id"] = value;
+      } else {
+        envVars[key] = value;
+      }
+    }
+  }
+
+  return envVars;
+}
+
 /**
  * Retrieves the environment variables from the current runtime environment.
  *
@@ -128,15 +170,14 @@ export function getEnvironmentVariables(): Record<string, string> | undefined {
     // eslint-disable-next-line no-process-env
     if (typeof process !== "undefined" && process.env) {
       // eslint-disable-next-line no-process-env
-      Object.entries(process.env).reduce(
+      return Object.entries(process.env).reduce(
         (acc: { [key: string]: string }, [key, value]) => {
           acc[key] = String(value);
           return acc;
         },
         {}
       );
     }
-
     // For browsers and other environments, we may not have direct access to env variables
     // Return undefined or any other fallback as required.
     return undefined;