feat: Merge enqueued posts and patches to optimize batching (#135)

jacoblee93 · web-flow · commit d65e3019910b · 2026-04-28T10:59:20.000-07:00
* Merge enqueued posts and patches to optimize batching

* Exclude patched inputs by default

* Logs

* Fix import order

* Feedback
diff --git a/AGENTS.md b/AGENTS.md
@@ -44,6 +44,10 @@ fun format(variables: Map<String, Any>): PromptMessages {
 
 ## Kotlin idioms
 
+### Prefer immutable collection transformations
+
+Avoid mutable accumulators when `map`, `filter`, `partition`, `associate`, `buildMap`, or `buildList` express the same logic clearly. Use mutation only when it materially improves readability, performance, or is required by an API.
+
 ### Use `buildMap` / `buildList` instead of mutable + convert
 
 ```kotlin
diff --git a/langsmith-java-core/src/main/kotlin/com/langchain/smith/client/AutoBatchQueue.kt b/langsmith-java-core/src/main/kotlin/com/langchain/smith/client/AutoBatchQueue.kt
@@ -1,9 +1,11 @@
 package com.langchain.smith.client
 
+import com.fasterxml.jackson.databind.node.ObjectNode
 import com.langchain.smith.core.RequestOptions
 import com.langchain.smith.core.Timeout
 import com.langchain.smith.core.http.Headers
 import com.langchain.smith.core.http.QueryParams
+import com.langchain.smith.core.jsonMapper
 import com.langchain.smith.models.runs.Run
 import com.langchain.smith.models.runs.RunIngestBatchParams
 import java.util.concurrent.CompletionException
@@ -19,6 +21,7 @@ import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicBoolean
 import java.util.concurrent.atomic.AtomicInteger
 import java.util.concurrent.locks.ReentrantLock
+import kotlin.jvm.optionals.getOrNull
 import org.slf4j.LoggerFactory
 
 /**
@@ -238,9 +241,6 @@ class AutoBatchQueue(
      * Drains up to [maxItems] queued operations and returns batch params grouped by request
      * options.
      *
-     * TODO: Merge create + update for the same run ID before sending (like the JS/Python SDKs).
-     *   This would reduce the number of operations in each batch when a run is created and
-     *   immediately updated (common for short-lived runs).
      * TODO: Also flush/split batches based on serialized payload size, not just operation count.
      * TODO: Support multipart ingest endpoint for large payloads with attachments.
      * TODO: Support gzip compression for batch requests.
@@ -324,15 +324,58 @@ class AutoBatchQueue(
         val queryParams: QueryParams.Builder = QueryParams.builder(),
     ) {
         fun toBatch(): Batch {
+            val mergeResult = mergePostsAndPatches()
             val builder = RunIngestBatchParams.builder()
-            if (posts.isNotEmpty()) builder.post(posts)
-            if (patches.isNotEmpty()) builder.patch(patches)
+            if (mergeResult.posts.isNotEmpty()) builder.post(mergeResult.posts)
+            if (mergeResult.patches.isNotEmpty()) builder.patch(mergeResult.patches)
             builder.additionalHeaders(headers.build())
             builder.additionalQueryParams(queryParams.build())
             return Batch(params = builder.build(), requestOptions = requestOptions)
         }
+
+        private fun mergePostsAndPatches(): MergeResult {
+            if (posts.isEmpty() || patches.isEmpty()) {
+                return MergeResult(posts = posts, patches = patches)
+            }
+
+            val postsById =
+                posts.mapNotNull { post -> post.id().getOrNull()?.let { it to post } }.toMap()
+            val postsWithoutId = posts.filter { it.id().getOrNull() == null }
+            val patchesByPostId =
+                patches.mapNotNull { patch ->
+                    patch.id().getOrNull()?.takeIf(postsById::containsKey)?.let { it to patch }
+                }
+            val patchesByPostIdMap = patchesByPostId.toMap()
+            val standalonePatches =
+                patches.filter { patch ->
+                    patch.id().getOrNull()?.let(postsById::containsKey) != true
+                }
+
+            return MergeResult(
+                posts =
+                    postsWithoutId +
+                        postsById.map { (id, post) ->
+                            patchesByPostIdMap[id]?.let { mergePostAndPatch(post, it) } ?: post
+                        },
+                patches = standalonePatches,
+                mergedRunIds = patchesByPostId.map { it.first },
+            )
+        }
+
+        private fun mergePostAndPatch(post: Run, patch: Run): Run {
+            val merged = objectMapper.valueToTree<ObjectNode>(post)
+            val patchFields = objectMapper.valueToTree<ObjectNode>(patch)
+            patchFields.fields().forEach { (field, value) -> merged.set<ObjectNode>(field, value) }
+            return objectMapper.treeToValue(merged, Run::class.java)
+        }
     }
 
+    private data class MergeResult(
+        val posts: List<Run>,
+        val patches: List<Run>,
+        val mergedRunIds: List<String> = emptyList(),
+    )
+
     private data class BatchItem(
         val op: BatchOp,
         val run: Run,
@@ -348,6 +391,7 @@ class AutoBatchQueue(
 
     companion object {
         private val logger = LoggerFactory.getLogger(AutoBatchQueue::class.java)
+        private val objectMapper = jsonMapper()
 
         const val DEFAULT_BATCH_SIZE_LIMIT = 100
         const val DEFAULT_AGGREGATION_DELAY_MS = 250L
diff --git a/langsmith-java-core/src/main/kotlin/com/langchain/smith/tracing/RunTree.kt b/langsmith-java-core/src/main/kotlin/com/langchain/smith/tracing/RunTree.kt
@@ -132,11 +132,19 @@ class RunTree(
         submitSafely(e, name) { c.runs().create(runData) }
     }
 
+    /**
+     * Patches this run in LangSmith in the background, batched with other pending operations.
+     *
+     * Inputs are omitted by default so the patch does not duplicate or overwrite the create payload
+     * if the create and patch are merged during batching.
+     */
+    fun patchRun() = patchRun(excludeInputs = true)
+
     /** Patches this run in LangSmith in the background, batched with other pending operations. */
-    fun patchRun() {
+    fun patchRun(excludeInputs: Boolean) {
         val c = client ?: return
         val e = executor ?: DEFAULT_EXECUTOR
-        val runData = buildRunData()
+        val runData = buildRunData(includeInputs = !excludeInputs)
         submitSafely(e, name) {
             c.runs()
                 .update(
@@ -149,7 +157,9 @@ class RunTree(
     }
 
     /** Builds the [Run] data object for posting/patching to the LangSmith API. */
-    fun buildRunData(): Run {
+    fun buildRunData(): Run = buildRunData(includeInputs = true)
+
+    private fun buildRunData(includeInputs: Boolean): Run {
         val builder =
             Run.builder()
                 .id(id)
@@ -160,9 +170,6 @@ class RunTree(
                 .startTime(startTime)
                 .apply { this@RunTree.projectName?.let { sessionName(it) } }
                 .tags(tags)
-                .inputs(
-                    Run.Inputs.builder().putAllAdditionalProperties(toJsonValueMap(inputs)).build()
-                )
                 .extra(
                     Run.Extra.builder()
                         .apply {
@@ -177,6 +184,11 @@ class RunTree(
                         .build()
                 )
 
+        if (includeInputs) {
+            builder.inputs(
+                Run.Inputs.builder().putAllAdditionalProperties(toJsonValueMap(inputs)).build()
+            )
+        }
         endTime?.let { builder.endTime(it) }
         error?.let { builder.error(it) }
         parentRunId?.let { builder.parentRunId(it) }
diff --git a/langsmith-java-core/src/test/kotlin/com/langchain/smith/client/AutoBatchQueueTest.kt b/langsmith-java-core/src/test/kotlin/com/langchain/smith/client/AutoBatchQueueTest.kt
@@ -1,5 +1,6 @@
 package com.langchain.smith.client
 
+import com.langchain.smith.core.JsonValue
 import com.langchain.smith.core.RequestOptions
 import com.langchain.smith.core.Timeout
 import com.langchain.smith.core.http.Headers
@@ -84,6 +85,45 @@ internal class AutoBatchQueueTest {
         assertThat(capture.patchedRuns).hasSize(1)
     }
 
+    @Test
+    fun `merges post and patch for same run id`() {
+        val capture = CapturingLangsmithClient()
+        val queue = testQueue(capture)
+        val post =
+            testRun("r1")
+                .toBuilder()
+                .inputs(
+                    Run.Inputs.builder()
+                        .putAdditionalProperty("question", JsonValue.from("hello"))
+                        .build()
+                )
+                .build()
+        val patch =
+            Run.builder()
+                .id("r1")
+                .outputs(
+                    Run.Outputs.builder()
+                        .putAdditionalProperty("answer", JsonValue.from("world"))
+                        .build()
+                )
+                .endTime("2024-01-01T00:00:01Z")
+                .build()
+
+        queue.post(post)
+        queue.patch(patch)
+        queue.flush()
+
+        assertThat(capture.postedRuns).hasSize(1)
+        assertThat(capture.patchedRuns).isEmpty()
+        val mergedRun = capture.postedRuns.single()
+        assertThat(mergedRun.id().getOrNull()).isEqualTo("r1")
+        assertThat(mergedRun.inputs().get()._additionalProperties()["question"])
+            .isEqualTo(JsonValue.from("hello"))
+        assertThat(mergedRun.outputs().get()._additionalProperties()["answer"])
+            .isEqualTo(JsonValue.from("world"))
+        assertThat(mergedRun.endTime().getOrNull()).isEqualTo("2024-01-01T00:00:01Z")
+    }
+
     @Test
     fun `flush is no-op when queue is empty`() {
         val capture = CapturingLangsmithClient()
diff --git a/langsmith-java-core/src/test/kotlin/com/langchain/smith/tracing/RunTreeTest.kt b/langsmith-java-core/src/test/kotlin/com/langchain/smith/tracing/RunTreeTest.kt
@@ -0,0 +1,61 @@
+package com.langchain.smith.tracing
+
+import com.langchain.smith.testutils.CapturingLangsmithClient
+import java.util.concurrent.Executors
+import java.util.concurrent.TimeUnit
+import kotlin.jvm.optionals.getOrNull
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.jupiter.api.Test
+
+internal class RunTreeTest {
+
+    @Test
+    fun `patchRun omits inputs by default`() {
+        val capture = CapturingLangsmithClient()
+        val executor = Executors.newSingleThreadExecutor()
+        val run =
+            RunTree.builder()
+                .name("test")
+                .inputs(mapOf("question" to "hello"))
+                .outputs(mapOf("answer" to "world"))
+                .client(capture.client)
+                .executor(executor)
+                .build()
+
+        run.patchRun()
+        executor.shutdown()
+        assertThat(executor.awaitTermination(5, TimeUnit.SECONDS)).isTrue()
+
+        assertThat(capture.patchedRuns).hasSize(1)
+        assertThat(capture.patchedRuns.single().inputs().isPresent).isFalse()
+    }
+
+    @Test
+    fun `patchRun can include inputs when requested`() {
+        val capture = CapturingLangsmithClient()
+        val executor = Executors.newSingleThreadExecutor()
+        val run =
+            RunTree.builder()
+                .name("test")
+                .inputs(mapOf("question" to "hello"))
+                .client(capture.client)
+                .executor(executor)
+                .build()
+
+        run.patchRun(excludeInputs = false)
+        executor.shutdown()
+        assertThat(executor.awaitTermination(5, TimeUnit.SECONDS)).isTrue()
+
+        assertThat(capture.patchedRuns).hasSize(1)
+        assertThat(
+                capture.patchedRuns
+                    .single()
+                    .inputs()
+                    .get()
+                    ._additionalProperties()["question"]
+                    ?.asString()
+                    ?.getOrNull()
+            )
+            .isEqualTo("hello")
+    }
+}
diff --git a/langsmith-java-example/build.gradle.kts b/langsmith-java-example/build.gradle.kts
@@ -27,6 +27,9 @@ dependencies {
     implementation("org.springframework.boot:spring-boot-starter-web")
     implementation("org.springframework.boot:spring-boot-starter")
 
+    // Simple logging for examples. This lets smoke tests show SDK debug/trace logs.
+    runtimeOnly("org.slf4j:slf4j-simple:2.0.17")
+
     // Security: constrain vulnerable transitive dependencies from Spring Boot 2.7.18.
     // Spring Boot 2.7.x is EOL; these constraints override the managed versions in-place.
     // None of these affect published artifacts (this is a non-published example module).
@@ -94,6 +97,10 @@ application {
 // Export stdin to examples for readln(); require -Pexample= when running (configuration-cache safe: no project access in doFirst)
 tasks.named<JavaExec>("run") {
     standardInput = System.`in`
+    systemProperty("org.slf4j.simpleLogger.defaultLogLevel", "info")
+    systemProperty("org.slf4j.simpleLogger.log.com.langchain.smith.client.AutoBatchQueue", "trace")
+    systemProperty("org.slf4j.simpleLogger.showDateTime", "false")
+    systemProperty("org.slf4j.simpleLogger.showShortLogName", "true")
     doFirst {
         if (mainClass.get() == "Main") {
             throw GradleException(
