patch(python, js): Fix sampling logic for patch runs (#1881)

### Description
We were not correctly filtering patch runs when the post would get
filtered, meaning large amounts of patch requests were being sent to
ingest API and counting toward usage limits, but not actually getting
ingested.

### Changes
Check patch trace id against trace id list, only remove trace ids from
list for root runs.

### Alternative Solutions
Use consistent hashing based on run_id to sample both post and patch run
ids equally, prevents storage of uuids (minimal memory overhead)

## Tests
Added unit tests confirming filtering logic correctly applies to both
post and patch

Author: angus-langchain

js/src/client.ts

@@ -935,10 +935,10 @@ export class Client implements LangSmithTracingClientInterface {
     if (patch) {
       const sampled = [];
       for (const run of runs) {
-        if (!this.filteredPostUuids.has(run.id)) {
+        if (!this.filteredPostUuids.has(run.trace_id)) {
           sampled.push(run);
-        } else {
-          this.filteredPostUuids.delete(run.id);
+        } else if (run.id === run.trace_id) {
+          this.filteredPostUuids.delete(run.trace_id);
         }
       }
       return sampled;

js/src/tests/client.test.ts

@@ -266,4 +266,218 @@ describe("Client", () => {
       });
     });
   });
+
+  describe("_filterForSampling patch logic", () => {
+    it("should filter patch runs based on trace_id instead of run.id", () => {
+      const client = new Client({
+        apiKey: "test-api-key",
+        tracingSamplingRate: 0.5,
+      });
+
+      // Mock the _shouldSample method to control sampling decisions
+      let counter = 0;
+      jest.spyOn(client as any, "_shouldSample").mockImplementation(() => {
+        counter += 1;
+        return counter % 2 === 0; // Accept even-numbered calls (2nd, 4th, etc.)
+      });
+
+      // Create two traces
+      const traceId1 = "trace-1";
+      const traceId2 = "trace-2";
+      const childRunId1 = "child-1";
+      const childRunId2 = "child-2";
+
+      // Create root runs (these will be sampled)
+      const rootRuns = [
+        {
+          id: traceId1,
+          trace_id: traceId1,
+          name: "root_run_1",
+          run_type: "llm" as const,
+          inputs: { text: "hello" },
+        },
+        {
+          id: traceId2,
+          trace_id: traceId2,
+          name: "root_run_2",
+          run_type: "llm" as const,
+          inputs: { text: "world" },
+        },
+      ];
+
+      // Test POST filtering (initial sampling)
+      const postFiltered = (client as any)._filterForSampling(rootRuns, false);
+
+      // Based on our mock, first call returns false, second returns true
+      // So only root_run_2 should be sampled
+      expect(postFiltered).toHaveLength(1);
+      expect(postFiltered[0].id).toBe(traceId2);
+
+      // Verify that traceId1 is in filtered set, traceId2 is not
+      expect((client as any).filteredPostUuids.has(traceId1)).toBe(true);
+      expect((client as any).filteredPostUuids.has(traceId2)).toBe(false);
+
+      // Test PATCH filtering - child runs should follow their trace's sampling decision
+      const patchRuns = [
+        {
+          id: childRunId1,
+          trace_id: traceId1,
+          name: "child_run_1",
+          run_type: "tool" as const,
+          inputs: { text: "child hello" },
+          outputs: { result: "child result 1" },
+        },
+        {
+          id: childRunId2,
+          trace_id: traceId2,
+          name: "child_run_2",
+          run_type: "tool" as const,
+          inputs: { text: "child world" },
+          outputs: { result: "child result 2" },
+        },
+      ];
+
+      const patchFiltered = (client as any)._filterForSampling(patchRuns, true);
+
+      // Only child_run_2 should be included (its trace was sampled)
+      // child_run_1 should be filtered out (its trace was not sampled)
+      expect(patchFiltered).toHaveLength(1);
+      expect(patchFiltered[0].id).toBe(childRunId2);
+      expect(patchFiltered[0].trace_id).toBe(traceId2);
+    });
+
+    it("should remove trace_id from filtered set when processing root run patches", () => {
+      const client = new Client({
+        apiKey: "test-api-key",
+        tracingSamplingRate: 0.5,
+      });
+
+      // Mock the _shouldSample method to reject first trace, accept second
+      let counter = 0;
+      jest.spyOn(client as any, "_shouldSample").mockImplementation(() => {
+        counter += 1;
+        return counter % 2 === 0;
+      });
+
+      const traceId1 = "trace-1";
+      const traceId2 = "trace-2";
+
+      // Create root runs and sample them
+      const rootRuns = [
+        {
+          id: traceId1,
+          trace_id: traceId1,
+          name: "root_run_1",
+          run_type: "llm" as const,
+          inputs: { text: "hello" },
+        },
+        {
+          id: traceId2,
+          trace_id: traceId2,
+          name: "root_run_2",
+          run_type: "llm" as const,
+          inputs: { text: "world" },
+        },
+      ];
+
+      (client as any)._filterForSampling(rootRuns, false);
+
+      // Verify initial state
+      expect((client as any).filteredPostUuids.has(traceId1)).toBe(true);
+      expect((client as any).filteredPostUuids.has(traceId2)).toBe(false);
+
+      // Test PATCH filtering for root runs (updates to the root runs themselves)
+      const rootPatchRuns = [
+        {
+          id: traceId1,
+          trace_id: traceId1,
+          name: "root_run_1",
+          run_type: "llm" as const,
+          inputs: { text: "hello" },
+          outputs: { result: "root result 1" },
+        },
+        {
+          id: traceId2,
+          trace_id: traceId2,
+          name: "root_run_2",
+          run_type: "llm" as const,
+          inputs: { text: "world" },
+          outputs: { result: "root result 2" },
+        },
+      ];
+
+      const rootPatchFiltered = (client as any)._filterForSampling(
+        rootPatchRuns,
+        true
+      );
+
+      // Only root_run_2 should be included, and traceId1 should be removed from filtered set
+      // since we're updating the root run that was originally filtered
+      expect(rootPatchFiltered).toHaveLength(1);
+      expect(rootPatchFiltered[0].id).toBe(traceId2);
+
+      // traceId1 should be removed from filtered set since we processed its root run
+      expect((client as any).filteredPostUuids.has(traceId1)).toBe(false);
+      expect((client as any).filteredPostUuids.has(traceId2)).toBe(false);
+    });
+
+    it("should handle mixed traces with patch sampling", () => {
+      const client = new Client({
+        apiKey: "test-api-key",
+        tracingSamplingRate: 0.5,
+      });
+
+      // Mock sampling to accept every other trace
+      let counter = 0;
+      jest.spyOn(client as any, "_shouldSample").mockImplementation(() => {
+        counter += 1;
+        return counter % 2 === 1; // Accept odd-numbered calls (1st, 3rd, etc.)
+      });
+
+      // Create multiple traces
+      const traceIds = ["trace-0", "trace-1", "trace-2", "trace-3"];
+      const childRunIds = ["child-0", "child-1", "child-2", "child-3"];
+
+      // Create root runs
+      const rootRuns = traceIds.map((traceId, i) => ({
+        id: traceId,
+        trace_id: traceId,
+        name: `root_run_${i}`,
+        run_type: "llm" as const,
+        inputs: { text: `hello ${i}` },
+      }));
+
+      // Sample the root runs
+      const postFiltered = (client as any)._filterForSampling(rootRuns, false);
+
+      // Based on our mock: 1st and 3rd calls return true (indices 0, 2)
+      expect(postFiltered).toHaveLength(2);
+      const sampledTraceIds = new Set(postFiltered.map((run: any) => run.id));
+      expect(sampledTraceIds.has(traceIds[0])).toBe(true);
+      expect(sampledTraceIds.has(traceIds[2])).toBe(true);
+
+      // Create child runs for all traces
+      const childRuns = traceIds.map((traceId, i) => ({
+        id: childRunIds[i],
+        trace_id: traceId,
+        name: `child_run_${i}`,
+        run_type: "tool" as const,
+        inputs: { text: `child ${i}` },
+        outputs: { result: `child result ${i}` },
+      }));
+
+      // Test patch filtering for child runs
+      const patchFiltered = (client as any)._filterForSampling(childRuns, true);
+
+      // Only children of sampled traces should be included
+      expect(patchFiltered).toHaveLength(2);
+      const patchTraceIds = new Set(
+        patchFiltered.map((run: any) => run.trace_id)
+      );
+      expect(patchTraceIds.has(traceIds[0])).toBe(true);
+      expect(patchTraceIds.has(traceIds[2])).toBe(true);
+      expect(patchTraceIds.has(traceIds[1])).toBe(false);
+      expect(patchTraceIds.has(traceIds[3])).toBe(false);
+    });
+  });
 });

python/langsmith/__init__.py

@@ -21,7 +21,7 @@
 
 # Avoid calling into importlib on every call to __version__
 
-__version__ = "0.4.8"
+__version__ = "0.4.9"
 version = __version__  # for backwards compatibility
 
 

python/langsmith/client.py

@@ -1284,11 +1284,11 @@ def _filter_for_sampling(
         if patch:
             sampled = []
             for run in runs:
-                run_id = _as_uuid(run["id"])
-                if run_id not in self._filtered_post_uuids:
+                trace_id = _as_uuid(run["trace_id"])
+                if trace_id not in self._filtered_post_uuids:
                     sampled.append(run)
-                else:
-                    self._filtered_post_uuids.remove(run_id)
+                elif run["id"] == trace_id:
+                    self._filtered_post_uuids.remove(trace_id)
             return sampled
         else:
             sampled = []

python/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langsmith"
-version = "0.4.8"
+version = "0.4.9"
 description = "Client library to connect to the LangSmith LLM Tracing and Evaluation Platform."
 authors = ["LangChain <[email protected]>"]
 license = "MIT"