feat: Adds processInputs and processOutputs to traceable (#113)

* Adds processInputs and processOutputs to traceable

* Move generics into TraceProcessIO to avoid having them top level

* Fix docstring, nit

Author: jacoblee93

langsmith-java-core/src/main/kotlin/com/langchain/smith/tracing/RunTree.kt

@@ -15,6 +15,12 @@ private val logger = LoggerFactory.getLogger(RunTree::class.java)
 /**
  * Represents a run in the trace tree.
  *
+ * **Thread safety:** `RunTree` is mutable — [inputs], [outputs], [error], [endTime], [metadata],
+ * [tags], and [extra] can be modified after construction. A single `RunTree` instance should only
+ * be accessed from one thread at a time. Within a [traceable] wrapper this is guaranteed by the
+ * run-context scoping, but if you hold a reference to a `RunTree` and access it from multiple
+ * threads, external synchronization is required.
+ *
  * Can be created manually, via [traceable], or via [createChild]:
  * ```kotlin
  * // Via traceable (most common)
@@ -36,8 +42,7 @@ private val logger = LoggerFactory.getLogger(RunTree::class.java)
  *     .build();
  * ```
  */
-class RunTree
-constructor(
+class RunTree(
     /** The display name of this run. */
     val name: String = "<lambda>",
     /** The type of run (chain, llm, tool, retriever). */

langsmith-java-core/src/main/kotlin/com/langchain/smith/tracing/TraceProcessIO.kt

@@ -0,0 +1,97 @@
+@file:JvmName("TraceProcessIOHelper")
+
+package com.langchain.smith.tracing
+
+import java.util.function.Function
+
+/**
+ * Typed processors for transforming inputs and outputs before they are recorded on a traced run.
+ *
+ * The type parameters describe what the processors receive:
+ * - [PI] — the type passed to [inputs]. For 1-arg traced functions this is the raw input type; for
+ *   0/2/3-arg functions this is `Map<String, Any?>` (the packed representation).
+ * - [PO] — the type passed to [outputs]. This is always the raw output type of the traced function.
+ *
+ * When [inputs] is set, it replaces the default input serialization entirely. Same for [outputs].
+ *
+ * ## Example (Kotlin)
+ *
+ * ```kotlin
+ * val process = TraceProcessIO<String, MyResponse>(
+ *     processInputs = Function { input -> mapOf("query" to input) },
+ *     processOutputs = Function { resp -> mapOf("answer" to resp.text) },
+ * )
+ * val config = TraceConfig(name = "my-run", client = client, processTracedIO = process)
+ * ```
+ *
+ * ## Example (Java)
+ *
+ * ```java
+ * TraceProcessIO process = TraceProcessIO.<String, MyResponse>builder()
+ *     .processInputs(input -> Map.of("query", input))
+ *     .processOutputs(resp -> Map.of("answer", resp.getText()))
+ *     .build();
+ * TraceConfig config = TraceConfig.builder()
+ *     .name("my-run")
+ *     .client(client)
+ *     .processTracedIO(process)
+ *     .build();
+ * ```
+ *
+ * @see TraceConfig.processTracedIO
+ */
+class TraceProcessIO<PI, PO>(
+    processInputs: Function<PI, Map<String, Any?>>? = null,
+    processOutputs: Function<PO, Map<String, Any?>>? = null,
+) {
+    // Store erased wrappers so traceable overloads can call these without casts.
+    // erase() converts Function<T, R> to Function<Any?, R> — safe because T erases to Object.
+    internal val inputsFn: Function<Any?, Map<String, Any?>>? = processInputs?.erase()
+    internal val outputsFn: Function<Any?, Map<String, Any?>>? = processOutputs?.erase()
+
+    companion object {
+        /** Creates a new [Builder]. */
+        @JvmStatic fun <PI, PO> builder() = Builder<PI, PO>()
+
+        /**
+         * Converts a typed [Function] to an erased `Function<Any?, R>`.
+         *
+         * Safe because generic type parameters erase to `Object` at runtime — the JVM never
+         * actually checks the input type on `Function.apply`.
+         */
+        @JvmSynthetic
+        @JvmStatic
+        @Suppress("UNCHECKED_CAST")
+        private fun <T, R> Function<T, R>.erase(): Function<Any?, R> = this as Function<Any?, R>
+    }
+
+    /**
+     * A builder for [TraceProcessIO].
+     *
+     * ```java
+     * TraceProcessIO process = TraceProcessIO.<String, String>builder()
+     *     .processInputs(input -> Map.of("query", input))
+     *     .processOutputs(output -> Map.of("answer", output))
+     *     .build();
+     * ```
+     */
+    class Builder<PI, PO> internal constructor() {
+        private var processInputs: Function<PI, Map<String, Any?>>? = null
+        private var processOutputs: Function<PO, Map<String, Any?>>? = null
+
+        /** Callback to transform the input before it is recorded on the run. */
+        fun processInputs(processInputs: Function<PI, Map<String, Any?>>) = apply {
+            this.processInputs = processInputs
+        }
+
+        /** Callback to transform the output before it is recorded on the run. */
+        fun processOutputs(processOutputs: Function<PO, Map<String, Any?>>) = apply {
+            this.processOutputs = processOutputs
+        }
+
+        /** Builds the [TraceProcessIO]. */
+        fun build() = TraceProcessIO(processInputs = processInputs, processOutputs = processOutputs)
+    }
+
+    override fun toString(): String = "TraceProcessIO{inputs=$inputsFn, outputs=$outputsFn}"
+}

langsmith-java-core/src/main/kotlin/com/langchain/smith/tracing/Traceable.kt

@@ -233,6 +233,27 @@ class TraceConfig(
     val executor: ExecutorService? = null,
     /** Whether tracing is active. Inherited by children. */
     val tracingEnabled: Boolean? = null,
+    /**
+     * Optional typed processors for transforming inputs/outputs before they are recorded on a run.
+     *
+     * When set, the processors in [TraceProcessIO] replace the default input/output serialization.
+     * The generic type parameters on [TraceProcessIO] control what the processor callbacks receive
+     * — see [TraceProcessIO] for details.
+     *
+     * ```java
+     * TraceProcessIO process = TraceProcessIO.<String, MyResponse>builder()
+     *     .processInputs(input -> Map.of("query", input))
+     *     .processOutputs(resp -> Map.of("answer", resp.getText()))
+     *     .build();
+     * TraceConfig config = TraceConfig.builder()
+     *     .name("my-run")
+     *     .processTracedIO(process)
+     *     .build();
+     * ```
+     *
+     * @see TraceProcessIO
+     */
+    val processTracedIO: TraceProcessIO<*, *>? = null,
 ) {
     companion object {
         /** Creates a new [Builder] for constructing a [TraceConfig]. */
@@ -250,9 +271,6 @@ class TraceConfig(
      *     .name("my-run")
      *     .client(LangsmithOkHttpClient.fromEnv())
      *     .runType(RunType.LLM)
-     *     .metadata(Map.of("version", "1.0"))
-     *     .tags(List.of("prod"))
-     *     .projectName("my-project")
      *     .build();
      * ```
      */
@@ -265,6 +283,7 @@ class TraceConfig(
         private var projectName: String? = null
         private var executor: ExecutorService? = null
         private var tracingEnabled: Boolean? = null
+        private var processTracedIO: TraceProcessIO<*, *>? = null
 
         @JvmSynthetic
         internal fun from(config: TraceConfig) = apply {
@@ -276,6 +295,7 @@ class TraceConfig(
             projectName = config.projectName
             executor = config.executor
             tracingEnabled = config.tracingEnabled
+            processTracedIO = config.processTracedIO
         }
 
         /** The name of the run, displayed in LangSmith. */
@@ -302,6 +322,16 @@ class TraceConfig(
         /** Whether tracing is active. Inherited by children if not set. */
         fun tracingEnabled(tracingEnabled: Boolean) = apply { this.tracingEnabled = tracingEnabled }
 
+        /**
+         * Typed processors for transforming inputs/outputs before they are recorded.
+         *
+         * @see TraceConfig.processTracedIO
+         * @see TraceProcessIO
+         */
+        fun processTracedIO(processTracedIO: TraceProcessIO<*, *>) = apply {
+            this.processTracedIO = processTracedIO
+        }
+
         /** Builds the [TraceConfig]. */
         fun build() =
             TraceConfig(
@@ -313,12 +343,21 @@ class TraceConfig(
                 projectName = projectName,
                 executor = executor,
                 tracingEnabled = tracingEnabled,
+                processTracedIO = processTracedIO,
             )
     }
 }
 
 private const val DEFAULT_RUN_NAME = "<lambda>"
 
+/** Applies processInputs from the config, or returns null if not set. */
+private fun applyProcessInputs(config: TraceConfig, value: Any?): Map<String, Any?>? =
+    config.processTracedIO?.inputsFn?.apply(value)
+
+/** Applies processOutputs from the config, or returns null if not set. */
+private fun applyProcessOutputs(config: TraceConfig, value: Any?): Map<String, Any?>? =
+    config.processTracedIO?.outputsFn?.apply(value)
+
 /** Resolves the run name from the config and the function being wrapped. */
 private fun resolveName(config: TraceConfig, block: Any?): String {
     config.name?.let {
@@ -328,6 +367,12 @@ private fun resolveName(config: TraceConfig, block: Any?): String {
     return DEFAULT_RUN_NAME
 }
 
+/** Creates a copy of [config] with the resolved name set. */
+private fun resolveConfig(config: TraceConfig, block: Any?): TraceConfig {
+    val name = resolveName(config, block)
+    return if (name == config.name) config else config.toBuilder().name(name).build()
+}
+
 /**
  * Wraps a no-arg function with LangSmith tracing (Kotlin).
  *
@@ -339,14 +384,22 @@ private fun resolveName(config: TraceConfig, block: Any?): String {
  */
 @JvmSynthetic
 fun <O> traceable(block: () -> O, config: TraceConfig): () -> O {
-    val resolvedConfig = config.toBuilder().name(resolveName(config, block)).build()
-    return { executeTraced(resolvedConfig, emptyMap()) { block() } }
+    val resolvedConfig = resolveConfig(config, block)
+    return {
+        val packed = emptyMap<String, Any?>()
+        val serializedInputs = applyProcessInputs(resolvedConfig, packed) ?: packed
+        executeTraced(resolvedConfig, serializedInputs) { block() }
+    }
 }
 
 /** Wraps a no-arg function with LangSmith tracing (Java [Supplier]). */
 fun <O> traceable(block: Supplier<O>, config: TraceConfig): Supplier<O> {
-    val resolvedConfig = config.toBuilder().name(resolveName(config, block)).build()
-    return Supplier { executeTraced(resolvedConfig, emptyMap()) { block.get() } }
+    val resolvedConfig = resolveConfig(config, block)
+    return Supplier {
+        val packed = emptyMap<String, Any?>()
+        val serializedInputs = applyProcessInputs(resolvedConfig, packed) ?: packed
+        executeTraced(resolvedConfig, serializedInputs) { block.get() }
+    }
 }
 
 /**
@@ -393,10 +446,10 @@ fun <O> traceable(block: Supplier<O>, config: TraceConfig): Supplier<O> {
  */
 @JvmSynthetic
 fun <I, O> traceable(block: (I) -> O, config: TraceConfig): (I) -> O {
-    val resolvedConfig = config.toBuilder().name(resolveName(config, block)).build()
+    val resolvedConfig = resolveConfig(config, block)
     return { input ->
-        val inputs = toInputMap(input)
-        executeTraced(resolvedConfig, inputs) { block(input) }
+        val serializedInputs = applyProcessInputs(resolvedConfig, input) ?: toInputMap(input)
+        executeTraced(resolvedConfig, serializedInputs) { block(input) }
     }
 }
 
@@ -409,10 +462,11 @@ fun <I, O> traceable(block: Function<I, O>, config: TraceConfig): Function<I, O>
 /** Wraps a 2-arg function with LangSmith tracing (Kotlin). */
 @JvmSynthetic
 fun <I1, I2, O> traceable(block: (I1, I2) -> O, config: TraceConfig): (I1, I2) -> O {
-    val resolvedConfig = config.toBuilder().name(resolveName(config, block)).build()
+    val resolvedConfig = resolveConfig(config, block)
     return { i1, i2 ->
-        val inputs = mapOf("args" to listOf(i1, i2))
-        executeTraced(resolvedConfig, inputs) { block(i1, i2) }
+        val packed = mapOf<String, Any?>("args" to listOf(i1, i2))
+        val serializedInputs = applyProcessInputs(resolvedConfig, packed) ?: packed
+        executeTraced(resolvedConfig, serializedInputs) { block(i1, i2) }
     }
 }
 
@@ -428,10 +482,11 @@ fun <I1, I2, O> traceable(
 /** Wraps a 3-arg function with LangSmith tracing (Kotlin). */
 @JvmSynthetic
 fun <I1, I2, I3, O> traceable(block: (I1, I2, I3) -> O, config: TraceConfig): (I1, I2, I3) -> O {
-    val resolvedConfig = config.toBuilder().name(resolveName(config, block)).build()
+    val resolvedConfig = resolveConfig(config, block)
     return { i1, i2, i3 ->
-        val inputs = mapOf("args" to listOf(i1, i2, i3))
-        executeTraced(resolvedConfig, inputs) { block(i1, i2, i3) }
+        val packed = mapOf<String, Any?>("args" to listOf(i1, i2, i3))
+        val serializedInputs = applyProcessInputs(resolvedConfig, packed) ?: packed
+        executeTraced(resolvedConfig, serializedInputs) { block(i1, i2, i3) }
     }
 }
 
@@ -540,7 +595,7 @@ fun <I1, I2, I3, O> traceTriFunction(
  *
  * For best tracing results, pass [Map] inputs/outputs to [traceable]. Typed SDK objects (e.g.
  * `ChatCompletionCreateParams`) should be converted to maps by the caller — use
- * `processInputs`/`processOutputs` callbacks (when available) or manual conversion.
+ * [TraceConfig.processTracedIO] callbacks or manual conversion.
  */
 private fun serializeValue(value: Any?): Any? =
     when (value) {
@@ -630,7 +685,7 @@ private fun <T> executeTraced(config: TraceConfig, inputs: Map<String, Any?>?, b
         try {
             val result = block()
             run.endTime = nowIso()
-            run.outputs = toOutputMap(result)
+            run.outputs = applyProcessOutputs(config, result) ?: toOutputMap(result)
             run.patchRun()
             result
         } catch (e: Throwable) {

langsmith-java-core/src/test/java/com/langchain/smith/tracing/TraceableJavaTest.java

@@ -3,6 +3,7 @@
 import static org.assertj.core.api.Assertions.assertThat;
 
 import com.langchain.smith.client.LangsmithClient;
+import java.util.Collections;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -158,4 +159,50 @@ void tracingDisabled_getCurrentRunReturnsNull() {
         Function<String, RunTree> traced = Tracing.traceFunction(input -> Tracing.getCurrentRunTree(), disabled);
         assertThat(traced.apply("hello")).isNull();
     }
+
+    // ---- processInputs / processOutputs ----
+
+    @Test
+    void processInputs_viaProcessTracedIO() {
+        TraceConfig cfg = TraceConfig.builder()
+                .name("process-inputs-test")
+                .client(client)
+                .tracingEnabled(true)
+                .processTracedIO(TraceProcessIO.<String, RunTree>builder()
+                        .processInputs(input -> Collections.singletonMap("query", input))
+                        .build())
+                .build();
+        Function<String, RunTree> traced = Tracing.traceFunction(input -> Tracing.getCurrentRunTree(), cfg);
+        RunTree run = traced.apply("hello");
+        assertThat(run.getInputs()).containsEntry("query", "hello");
+    }
+
+    @Test
+    void processOutputs_viaProcessTracedIO() {
+        TraceConfig cfg = TraceConfig.builder()
+                .name("process-outputs-test")
+                .client(client)
+                .tracingEnabled(true)
+                .processTracedIO(TraceProcessIO.<String, String>builder()
+                        .processOutputs(output -> Collections.singletonMap("answer", output))
+                        .build())
+                .build();
+        Function<String, String> traced = Tracing.traceFunction(input -> "result", cfg);
+        traced.apply("hello");
+    }
+
+    @Test
+    void processInputsAndOutputs_together() {
+        TraceConfig cfg = TraceConfig.builder()
+                .name("both-processors")
+                .client(client)
+                .tracingEnabled(true)
+                .processTracedIO(TraceProcessIO.<String, String>builder()
+                        .processInputs(input -> Collections.singletonMap("q", input))
+                        .processOutputs(output -> Collections.singletonMap("a", output))
+                        .build())
+                .build();
+        Function<String, String> traced = Tracing.traceFunction(input -> "answer", cfg);
+        assertThat(traced.apply("question")).isEqualTo("answer");
+    }
 }