feat: Filter kwargs from new token events (#2714)

Author: jacoblee93

js/src/client.ts

@@ -1066,6 +1066,27 @@ export class Client implements LangSmithTracingClientInterface {
     return outputs;
   }
 
+  /**
+   * Filter content from new_token events to prevent streaming LLM output
+   * from being uploaded via events.
+   */
+  private _filterNewTokenEvents(
+    events: KVMap[] | undefined
+  ): KVMap[] | undefined {
+    if (!events || events.length === 0) {
+      return events;
+    }
+    return events.map((event) => {
+      if (event.name === "new_token") {
+        // Remove the kwargs containing the token data
+        // eslint-disable-next-line @typescript-eslint/no-unused-vars
+        const { kwargs: _, ...rest } = event;
+        return rest;
+      }
+      return event;
+    });
+  }
+
   private async prepareRunCreateOrUpdateInputs(
     run: RunUpdate
   ): Promise<RunUpdate>;
@@ -1082,6 +1103,9 @@ export class Client implements LangSmithTracingClientInterface {
     if (runParams.outputs !== undefined) {
       runParams.outputs = await this.processOutputs(runParams.outputs);
     }
+    if (runParams.events !== undefined) {
+      runParams.events = this._filterNewTokenEvents(runParams.events);
+    }
     return runParams;
   }
 
@@ -2189,6 +2213,10 @@ export class Client implements LangSmithTracingClientInterface {
     if (run.outputs) {
       run.outputs = await this.processOutputs(run.outputs);
     }
+
+    if (run.events) {
+      run.events = this._filterNewTokenEvents(run.events);
+    }
     // TODO: Untangle types
     const data: UpdateRunParams = { ...run, id: runId };
     if (!this._filterForSampling([data], true).length) {

js/src/tests/client.test.ts

@@ -1029,4 +1029,98 @@ describe("Client", () => {
       expect(mergedHeaders["x-api-key"]).toBe("test-api-key");
     });
   });
+
+  describe("_filterNewTokenEvents", () => {
+    it("should strip kwargs from new_token events", () => {
+      const client = new Client({ apiKey: "test-api-key" });
+      const events = [
+        {
+          name: "new_token",
+          kwargs: { token: "sensitive streaming data" },
+          time: "2024-01-01T00:00:00Z",
+        },
+        {
+          name: "other_event",
+          kwargs: { data: "keep this" },
+          time: "2024-01-01T00:00:01Z",
+        },
+      ];
+
+      const filtered = (client as any)._filterNewTokenEvents(events);
+
+      expect(filtered[0].name).toBe("new_token");
+      expect(filtered[0].time).toBe("2024-01-01T00:00:00Z");
+      expect(filtered[0].kwargs).toBeUndefined();
+      expect(filtered[1].kwargs).toEqual({ data: "keep this" });
+    });
+
+    it("should handle empty events array", () => {
+      const client = new Client({ apiKey: "test-api-key" });
+      const filtered = (client as any)._filterNewTokenEvents([]);
+      expect(filtered).toEqual([]);
+    });
+
+    it("should handle undefined events", () => {
+      const client = new Client({ apiKey: "test-api-key" });
+      const filtered = (client as any)._filterNewTokenEvents(undefined);
+      expect(filtered).toBeUndefined();
+    });
+
+    it("should handle events without kwargs", () => {
+      const client = new Client({ apiKey: "test-api-key" });
+      const events = [
+        { name: "new_token", time: "2024-01-01T00:00:00Z" },
+        { name: "other_event", time: "2024-01-01T00:00:01Z" },
+      ];
+
+      const filtered = (client as any)._filterNewTokenEvents(events);
+
+      expect(filtered[0]).toEqual({
+        name: "new_token",
+        time: "2024-01-01T00:00:00Z",
+      });
+      expect(filtered[1]).toEqual({
+        name: "other_event",
+        time: "2024-01-01T00:00:01Z",
+      });
+    });
+
+    it("should preserve other event properties", () => {
+      const client = new Client({ apiKey: "test-api-key" });
+      const events = [
+        {
+          name: "new_token",
+          kwargs: { token: "data" },
+          time: "2024-01-01T00:00:00Z",
+          message: "token received",
+          custom_field: "custom_value",
+        },
+      ];
+
+      const filtered = (client as any)._filterNewTokenEvents(events);
+
+      expect(filtered[0].name).toBe("new_token");
+      expect(filtered[0].time).toBe("2024-01-01T00:00:00Z");
+      expect(filtered[0].message).toBe("token received");
+      expect(filtered[0].custom_field).toBe("custom_value");
+      expect(filtered[0].kwargs).toBeUndefined();
+    });
+
+    it("should filter multiple new_token events", () => {
+      const client = new Client({ apiKey: "test-api-key" });
+      const events = [
+        { name: "new_token", kwargs: { token: "chunk1" }, time: "t1" },
+        { name: "new_token", kwargs: { token: "chunk2" }, time: "t2" },
+        { name: "new_token", kwargs: { token: "chunk3" }, time: "t3" },
+      ];
+
+      const filtered = (client as any)._filterNewTokenEvents(events);
+
+      expect(filtered).toHaveLength(3);
+      filtered.forEach((event: any) => {
+        expect(event.kwargs).toBeUndefined();
+        expect(event.name).toBe("new_token");
+      });
+    });
+  });
 });

js/src/tests/wrapped_anthropic.int.test.ts

@@ -143,8 +143,7 @@ describe.skip("Requires Anthropic API key", () => {
     expect(tokenEvents.length).toBeGreaterThan(0);
     tokenEvents.forEach((event: any) => {
       expect(event.name).toBe("new_token");
-      expect(event.kwargs).toBeDefined();
-      expect(event.kwargs.token).toBeDefined();
+      expect(event.kwargs).toBeUndefined();
       expect(event.time).toBeDefined();
     });
 

js/src/tests/wrapped_openai.int.test.ts

@@ -111,8 +111,7 @@ test("chat.completions", async () => {
   expect(tokenEvents.length).toBeGreaterThan(0);
   tokenEvents.forEach((event: any) => {
     expect(event.name).toBe("new_token");
-    expect(event.kwargs).toBeDefined();
-    expect(event.kwargs.token).toBeDefined();
+    expect(event.kwargs).toBeUndefined();
     expect(event.time).toBeDefined();
   });
 
@@ -368,8 +367,7 @@ test("chat completions with tool calling", async () => {
   expect(tokenEvents.length).toBeGreaterThan(0);
   tokenEvents.forEach((event: any) => {
     expect(event.name).toBe("new_token");
-    expect(event.kwargs).toBeDefined();
-    expect(event.kwargs.token).toBeDefined();
+    expect(event.kwargs).toBeUndefined();
     expect(event.time).toBeDefined();
   });
 

python/langsmith/client.py

@@ -1946,6 +1946,8 @@ def _run_transform(
             if copy:
                 run_create["outputs"] = ls_utils.deepish_copy(run_create["outputs"])
             run_create["outputs"] = self._hide_run_outputs(run_create["outputs"])
+        if "events" in run_create and run_create["events"] is not None:
+            run_create["events"] = self._filter_new_token_events(run_create["events"])
         # Hide metadata in extra if present
         if "extra" in run_create and isinstance(run_create["extra"], dict):
             extra = run_create["extra"]
@@ -2380,6 +2382,29 @@ def _hide_run_metadata(self, metadata: dict) -> dict:
             return metadata
         return self._hide_metadata(metadata)
 
+    @staticmethod
+    def _filter_new_token_events(
+        events: Optional[Sequence[dict]],
+    ) -> Optional[list[dict]]:
+        """Filter content from new_token events.
+
+        This prevents streaming LLM output from being uploaded via events.
+
+        Args:
+            events: The events to filter.
+
+        Returns:
+            The filtered events with kwargs removed from new_token events.
+        """
+        if not events:
+            return events  # type: ignore[return-value]
+        return [
+            {k: v for k, v in event.items() if k != "kwargs"}
+            if event.get("name") == "new_token"
+            else event
+            for event in events
+        ]
+
     def _should_flush_run_ops_buffer(self) -> bool:
         """Check if the run ops buffer should be flushed based on size or time."""
         if not self._run_ops_buffer:
@@ -3290,7 +3315,7 @@ def update_run(
                 outputs = ls_utils.deepish_copy(outputs)
             data["outputs"] = self._hide_run_outputs(outputs)
         if events is not None:
-            data["events"] = events
+            data["events"] = self._filter_new_token_events(events)
         if data["extra"]:
             self._insert_runtime_env([data])
             if metadata := data["extra"].get("metadata"):

python/tests/unit_tests/test_client.py

@@ -5508,3 +5508,88 @@ def test_env_var_not_set_leaves_dir_none(self, monkeypatch):
         )
         assert client._failed_traces_dir is None
         _clear_env_cache()
+
+    def test_filter_new_token_events_strips_kwargs(self):
+        """Test that _filter_new_token_events strips kwargs from new_token events."""
+        client = Client(api_url="http://localhost:1984", api_key="test")
+        events = [
+            {
+                "name": "new_token",
+                "kwargs": {"token": "sensitive streaming data"},
+                "time": "2024-01-01T00:00:00Z",
+            },
+            {
+                "name": "other_event",
+                "kwargs": {"data": "keep this"},
+                "time": "2024-01-01T00:00:01Z",
+            },
+        ]
+
+        filtered = client._filter_new_token_events(events)
+
+        assert filtered[0]["name"] == "new_token"
+        assert filtered[0]["time"] == "2024-01-01T00:00:00Z"
+        assert "kwargs" not in filtered[0]
+        assert filtered[1]["kwargs"] == {"data": "keep this"}
+
+    def test_filter_new_token_events_empty_events(self):
+        """Test that _filter_new_token_events handles empty events list."""
+        client = Client(api_url="http://localhost:1984", api_key="test")
+        filtered = client._filter_new_token_events([])
+        assert filtered == []
+
+    def test_filter_new_token_events_none_events(self):
+        """Test that _filter_new_token_events handles None events."""
+        client = Client(api_url="http://localhost:1984", api_key="test")
+        filtered = client._filter_new_token_events(None)
+        assert filtered is None
+
+    def test_filter_new_token_events_without_kwargs(self):
+        """Test that _filter_new_token_events handles events without kwargs."""
+        client = Client(api_url="http://localhost:1984", api_key="test")
+        events = [
+            {"name": "new_token", "time": "2024-01-01T00:00:00Z"},
+            {"name": "other_event", "time": "2024-01-01T00:00:01Z"},
+        ]
+
+        filtered = client._filter_new_token_events(events)
+
+        assert filtered[0] == {"name": "new_token", "time": "2024-01-01T00:00:00Z"}
+        assert filtered[1] == {"name": "other_event", "time": "2024-01-01T00:00:01Z"}
+
+    def test_filter_new_token_events_preserves_other_properties(self):
+        """Test that _filter_new_token_events preserves other event properties."""
+        client = Client(api_url="http://localhost:1984", api_key="test")
+        events = [
+            {
+                "name": "new_token",
+                "kwargs": {"token": "data"},
+                "time": "2024-01-01T00:00:00Z",
+                "message": "token received",
+                "custom_field": "custom_value",
+            }
+        ]
+
+        filtered = client._filter_new_token_events(events)
+
+        assert filtered[0]["name"] == "new_token"
+        assert filtered[0]["time"] == "2024-01-01T00:00:00Z"
+        assert filtered[0]["message"] == "token received"
+        assert filtered[0]["custom_field"] == "custom_value"
+        assert "kwargs" not in filtered[0]
+
+    def test_filter_new_token_events_multiple_new_token_events(self):
+        """Test that _filter_new_token_events filters multiple new_token events."""
+        client = Client(api_url="http://localhost:1984", api_key="test")
+        events = [
+            {"name": "new_token", "kwargs": {"token": "chunk1"}, "time": "t1"},
+            {"name": "new_token", "kwargs": {"token": "chunk2"}, "time": "t2"},
+            {"name": "new_token", "kwargs": {"token": "chunk3"}, "time": "t3"},
+        ]
+
+        filtered = client._filter_new_token_events(events)
+
+        assert len(filtered) == 3
+        for event in filtered:
+            assert "kwargs" not in event
+            assert event["name"] == "new_token"