From 42b2f9c3258fbdb197eee3c4edc8f539052467a0 Mon Sep 17 00:00:00 2001 From: Kamil Tomaszek Date: Wed, 4 Nov 2026 09:52:09 -0800 Subject: [PATCH] feat(tokt): stream a Flowable FunctionTool in a live run PiperOrigin-RevId: 1013932657 --- tokt/README.md | 3 + .../kotlin/com/google/adk/tokt/JavaAdkToKt.kt | 6 +- .../google/adk/tokt/adapters/JavaToolToKt.kt | 29 ++- .../google/adk/tokt/KtRunnerInteropTest.kt | 220 +++++++++++++++++- 4 files changed, 253 insertions(+), 5 deletions(-) diff --git a/tokt/README.md b/tokt/README.md index 0c20beacd..78defdc8c 100644 --- a/tokt/README.md +++ b/tokt/README.md @@ -94,6 +94,9 @@ liveRunner `image/*` as a video frame. Any other blob, or one with no MIME type, is dropped with a warning, because the Kotlin engine has no untyped media input. +- An ADK Java `FunctionTool` that returns a `Flowable` streams its items to + the model in a live run; its `inputStream` parameter, if any, receives + `null`. - `RunConfig` settings the Kotlin engine cannot honor fail the stream rather than being dropped. - `runLive(session, ...)` runs the copy of that session the Kotlin runner's diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/JavaAdkToKt.kt b/tokt/src/main/kotlin/com/google/adk/tokt/JavaAdkToKt.kt index 7628b46cd..3304f5a0b 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/JavaAdkToKt.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/JavaAdkToKt.kt @@ -70,7 +70,11 @@ import kotlinx.coroutines.CoroutineDispatcher */ object JavaAdkToKt { - /** Adapts an ADK Java tool, hopping to `dispatcher` for its (possibly blocking) calls. */ + /** + * Adapts an ADK Java tool, hopping to `dispatcher` for its (possibly blocking) calls. A + * `FunctionTool` that returns a `Flowable` streams its items to the model in a live run; its + * `inputStream` parameter, if any, receives `null`. + */ @JvmStatic @JvmOverloads fun asKtTool( diff --git a/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaToolToKt.kt b/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaToolToKt.kt index e7bdc6edf..ccca5d64a 100644 --- a/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaToolToKt.kt +++ b/tokt/src/main/kotlin/com/google/adk/tokt/adapters/JavaToolToKt.kt @@ -21,19 +21,27 @@ import com.google.adk.kt.tools.BaseTool import com.google.adk.kt.tools.ToolContext import com.google.adk.kt.types.FunctionDeclaration import com.google.adk.tokt.codecs.FunctionDeclarationCodec +import com.google.adk.tokt.codecs.WorkflowCodec import com.google.adk.tokt.context.ktToolContextToJava import com.google.adk.tools.BaseTool as JavaBaseTool +import com.google.adk.tools.FunctionTool as JavaFunctionTool +import io.reactivex.rxjava3.core.Flowable +import java.lang.reflect.InvocationTargetException import kotlin.jvm.optionals.getOrNull import kotlinx.coroutines.CoroutineDispatcher +import kotlinx.coroutines.flow.flowOn +import kotlinx.coroutines.flow.map +import kotlinx.coroutines.reactive.asFlow import kotlinx.coroutines.reactive.awaitFirstOrNull import kotlinx.coroutines.reactive.awaitSingle import kotlinx.coroutines.withContext /** * Exposes a Java [JavaBaseTool] as a Kotlin [BaseTool] so the Kotlin engine can invoke it. [run] - * awaits the tool's RxJava result off the engine dispatcher (`dispatcher`), with actions delegating - * live to the Kotlin context. [processLlmRequest] re-applies a tool's request edits while - * preserving Kotlin-only fields (model, toolsDict). + * awaits the tool's RxJava result off the engine dispatcher (`dispatcher`), or returns a streaming + * `FunctionTool`'s items as a Flow, with actions delegating live to the Kotlin context. + * [processLlmRequest] re-applies a tool's request edits while preserving Kotlin-only fields (model, + * toolsDict). */ internal class JavaToolToKt( internal val javaTool: JavaBaseTool, @@ -51,6 +59,21 @@ internal class JavaToolToKt( override suspend fun run(context: ToolContext, args: Map): Any { val javaContext = ktToolContextToJava(context, dispatcher) + if (javaTool is JavaFunctionTool && javaTool.isStreaming) { + // Returned in any mode: the engine owns the Flow, and a live run streams it to the model. + val items: Flowable<*> = + withContext(dispatcher) { + try { + javaTool.callLive(args, javaContext, javaContext.invocationContext()) + } catch (e: InvocationTargetException) { + // Surface the tool's own exception, as calling its method directly would. + throw e.cause ?: e + } + } + // What the body set before returning still reaches the call's answer. + reconcileActionsToKt(javaContext.actions(), context.actions) + return items.asFlow().map { WorkflowCodec.outputFromJava(it) }.flowOn(dispatcher) + } // Run off the engine dispatcher so synchronous RxJava or blocking tool I/O cannot stall the // agent loop; the live actions view is backed by concurrent maps, so another thread is safe. val result = diff --git a/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt b/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt index 10f5b0eee..1225aef15 100644 --- a/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt +++ b/tokt/src/test/kotlin/com/google/adk/tokt/KtRunnerInteropTest.kt @@ -33,6 +33,7 @@ import com.google.adk.kt.agents.BaseAgent as KtBaseAgent import com.google.adk.kt.agents.CallbackContext as KtCallbackContext import com.google.adk.kt.agents.Context as KtContext import com.google.adk.kt.agents.InvocationContext as KtInvocationContext +import com.google.adk.kt.agents.LiveRequestQueue as KtLiveRequestQueue import com.google.adk.kt.agents.LiveRunConfig as KtLiveRunConfig import com.google.adk.kt.agents.LlmAgent as KtLlmAgent import com.google.adk.kt.agents.ResumabilityConfig as KtResumabilityConfig @@ -72,6 +73,7 @@ import com.google.adk.kt.types.CitationMetadata as KtCitationMetadata import com.google.adk.kt.types.Content as KtContent import com.google.adk.kt.types.FileData as KtFileData import com.google.adk.kt.types.FinishReason as KtFinishReason +import com.google.adk.kt.types.FunctionCall as KtFunctionCall import com.google.adk.kt.types.FunctionCallingConfig as KtFunctionCallingConfig import com.google.adk.kt.types.FunctionCallingConfigMode as KtFunctionCallingConfigMode import com.google.adk.kt.types.FunctionResponse as KtFunctionResponse @@ -164,8 +166,10 @@ import com.google.adk.tokt.codecs.toGenai import com.google.adk.tokt.codecs.toKt import com.google.adk.tokt.codecs.typedDataFromJava import com.google.adk.tokt.codecs.unsupportedFields +import com.google.adk.tools.Annotations.Schema import com.google.adk.tools.BaseTool as JavaBaseTool import com.google.adk.tools.BaseToolset as JavaBaseToolset +import com.google.adk.tools.FunctionTool as JavaFunctionTool import com.google.adk.tools.ToolContext as JavaToolContext import com.google.adk.workflow.NodeInfo as JavaNodeInfo import com.google.adk.workflow.Route as JavaRoute @@ -248,6 +252,7 @@ import kotlin.test.assertContentEquals import kotlin.test.assertEquals import kotlin.test.assertFailsWith import kotlin.test.assertFalse +import kotlin.test.assertIs import kotlin.test.assertNotNull import kotlin.test.assertNull import kotlin.test.assertSame @@ -2872,6 +2877,144 @@ class KtRunnerInteropTest { assertThat(genaiTranscriptionConfig().toKt()).isEqualTo(ktTranscriptionConfig()) } + @Test + fun asKtTool_streamingJavaTool_returnsItsItemsAsAFlow(): Unit = runBlocking { + // Returned in any mode: the engine owns the Flow, and a live run streams it to the model. + val tool = + JavaAdkToKt.asKtTool(JavaFunctionTool.create(StreamingJavaTools::class.java, "countUp")) + + val result = tool.run(KtToolContext(ktInvocationContext()), emptyMap()) + + assertEquals(listOf("1", "2"), assertIs>(result).toList()) + } + + @Test + fun asKtTool_nonStreamingJavaTool_runsInALiveRun(): Unit = runBlocking { + // Only a Flowable return streams; a queue parameter alone doesn't. + val tool = + JavaAdkToKt.asKtTool(JavaFunctionTool.create(StreamingJavaTools::class.java, "listen")) + val live = ktInvocationContext(KtRunConfig(streamingMode = KtStreamingMode.BIDI)) + + val result = tool.run(KtToolContext(live), emptyMap()) + + assertEquals(mapOf("result" to "heard"), result) + } + + @Test + fun asKtTool_streamingJavaTool_carriesTheActionsItsBodySets(): Unit = runBlocking { + val tool = + JavaAdkToKt.asKtTool( + JavaFunctionTool.create(StreamingJavaTools::class.java, "countUpQuietly") + ) + val context = KtToolContext(ktInvocationContext()) + + val result = tool.run(context, emptyMap()) + + assertEquals(listOf("1"), assertIs>(result).toList()) + assertTrue(context.actions.skipSummarization) + } + + @Test + fun asKtTool_streamingJavaToolWithAQueueParameter_getsNoQueue(): Unit = runBlocking { + val tool = + JavaAdkToKt.asKtTool(JavaFunctionTool.create(StreamingJavaTools::class.java, "countHeard")) + + val result = tool.run(KtToolContext(ktInvocationContext()), emptyMap()) + + assertEquals(listOf("no queue"), assertIs>(result).toList()) + } + + @Test + fun asKtTool_streamingJavaToolThatThrows_failsWithItsOwnException(): Unit = runBlocking { + val tool = + JavaAdkToKt.asKtTool(JavaFunctionTool.create(StreamingJavaTools::class.java, "failToStart")) + + assertFailsWith { + tool.run(KtToolContext(ktInvocationContext()), emptyMap()) + } + } + + @Test + fun asKtTool_streamingJavaTool_runsAndEmitsOnTheDispatcher(): Unit = runBlocking { + val executor = Executors.newSingleThreadExecutor { Thread(it, "tool-dispatcher") } + try { + val tool = + JavaAdkToKt.asKtTool( + JavaFunctionTool.create(StreamingJavaTools::class.java, "threadNames"), + executor.asCoroutineDispatcher(), + ) + + val result = tool.run(KtToolContext(ktInvocationContext()), emptyMap()) + + // Coroutine debug mode adds " @coroutine#n" to thread names. + val threads = assertIs>(result).toList().map { "$it".substringBefore(" @") } + assertEquals(listOf("tool-dispatcher", "tool-dispatcher"), threads) + } finally { + executor.shutdown() + } + } + + @Test + fun asJavaRunner_runLive_streamingJavaTool_streamsItsItemsToTheModel() { + val model = ToolCallLiveModel("countUp") + val countUp = + JavaAdkToKt.asKtTool(JavaFunctionTool.create(StreamingJavaTools::class.java, "countUp")) + val queue = LiveRequestQueue() + val events = + KotlinAdkToJava.asJavaRunner( + KtInMemoryRunner( + agent = KtLlmAgent(name = "a", model = model, tools = listOf(countUp)), + appName = "app", + ) + ) + .runLive("u", "s", queue, autoCreateSessionConfig()) + .test() + + try { + awaitTrue("both items should reach the model") { + model.results().containsAll(listOf("1", "2")) + } + } finally { + queue.close() + } + + events.awaitDone(10, TimeUnit.SECONDS).assertComplete().assertNoErrors() + } + + @Test + fun ktRunner_runLive_streamingJavaToolInStreamingModeNone_streamsItsItems() { + // A native live run keeps the caller's RunConfig, NONE by default; the engine still streams. + val model = ToolCallLiveModel("countUp") + val countUp = + JavaAdkToKt.asKtTool(JavaFunctionTool.create(StreamingJavaTools::class.java, "countUp")) + val runner = + KtInMemoryRunner( + agent = KtLlmAgent(name = "a", model = model, tools = listOf(countUp)), + appName = "app", + ) + val queue = KtLiveRequestQueue() + val failure = AtomicReference() + val run = thread { + try { + runBlocking { runner.runLive("u", "s", queue, KtRunConfig()).toList() } + } catch (e: Throwable) { + failure.set(e) + } + } + + try { + awaitTrue("both items should reach the model") { + model.results().containsAll(listOf("1", "2")) + } + } finally { + queue.close() + } + + run.join(TimeUnit.SECONDS.toMillis(10)) + assertFalse(run.isAlive, "the run should end once its queue closes") + assertNull(failure.get()) + } + @Test fun asJavaRunner_pluginManager_listsUnwrappedJavaPlugins() { // A Java plugin adapted onto the Kotlin runner is unwrapped back to the original instance, so a @@ -6888,6 +7031,80 @@ class KtRunnerInteropTest { } } + /** Static methods FunctionTool can wrap: those returning a Flowable stream; [listen] doesn't. */ + object StreamingJavaTools { + @JvmStatic fun countUp(): Flowable = Flowable.just("1", "2") + + @JvmStatic + fun countUpQuietly( + @Schema(name = "toolContext") toolContext: JavaToolContext + ): Flowable { + toolContext.actions().setSkipSummarization(true) + return Flowable.just("1") + } + + @JvmStatic + fun countHeard(@Schema(name = "inputStream") inputStream: LiveRequestQueue?): Flowable = + Flowable.just(if (inputStream == null) "no queue" else "queue") + + @JvmStatic fun failToStart(): Flowable = throw IllegalStateException("cannot start") + + /** Emits the thread that called it, then the thread that emits. */ + @JvmStatic + fun threadNames(): Flowable { + val caller = Thread.currentThread().name + return Flowable.fromCallable { Thread.currentThread().name }.startWithItem(caller) + } + + @JvmStatic + fun listen(@Schema(name = "inputStream") inputStream: LiveRequestQueue?): String = "heard" + } + + /** A live model that calls [toolName] once, then records every function response it is sent. */ + private class ToolCallLiveModel(private val toolName: String) : KtModel { + override val name = "tool-call-live-model" + private val responses = CopyOnWriteArrayList() + + /** The `result` value of each function response so far. */ + fun results(): List = responses.mapNotNull { it.response?.get("result") } + + override fun generateContent(request: KtLlmRequest, stream: Boolean): Flow = + throw UnsupportedOperationException("ToolCallLiveModel answers live connections only") + + override suspend fun connect(request: KtLlmRequest): KtLiveConnection = + object : KtLiveConnection { + private val turns = AtomicInteger() + private val untilClosed = Channel() + + override suspend fun sendHistory(history: List) {} + + override suspend fun sendContent(content: KtContent, partial: Boolean) { + responses += content.parts.mapNotNull { it.functionResponse } + } + + override suspend fun sendRealtime(input: KtRealtimeInput) {} + + override fun receive(): Flow = flow { + // The first turn asks for the tool; a later one waits for the session to close. + if (turns.getAndIncrement() == 0) { + val call = KtFunctionCall(name = toolName, id = "call-1") + emit( + KtLlmResponse( + content = KtContent(role = "model", parts = listOf(KtPart(functionCall = call))) + ) + ) + emit(KtLlmResponse(turnComplete = true)) + } else { + for (reply in untilClosed) emit(reply) + } + } + + override suspend fun closeSession() { + untilClosed.close() + } + } + } + private companion object { fun modelFunctionCall(name: String, args: Map): GenaiContent = GenaiContent.builder() @@ -6932,11 +7149,12 @@ class KtRunnerInteropTest { .build() /** A Kotlin invocation context on a fresh session, for driving a plugin adapter directly. */ - fun ktInvocationContext() = + fun ktInvocationContext(runConfig: KtRunConfig? = null) = KtInvocationContext( session = KtSession(KtSessionKey("app", "u", "s")), agent = KtLlmAgent(name = "a", model = JavaAdkToKt.asKtModel(SequentialJavaModel(emptyList()))), + runConfig = runConfig, ) /** A Kotlin callback context over [ktInvocationContext]. */