fix: allow interceptor to stop tool execution

flock-community · May 1, 2024 · b5b050c · b5b050c
1 parent b89e3bc
commit b5b050c
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 5 deletions.
diff --git a/src/core/src/commonMain/kotlin/community/flock/aigentic/core/agent/AgentExecutor.kt b/src/core/src/commonMain/kotlin/community/flock/aigentic/core/agent/AgentExecutor.kt
@@ -1,17 +1,16 @@
 package community.flock.aigentic.core.agent
 
+import ToolInterceptor
 import community.flock.aigentic.core.agent.tool.FinishReason
 import community.flock.aigentic.core.agent.tool.FinishedOrStuck
 import community.flock.aigentic.core.agent.tool.finishOrStuckTool
 import community.flock.aigentic.core.message.*
 import community.flock.aigentic.core.tool.ToolName
 import community.flock.aigentic.core.model.ModelResponse
-import community.flock.aigentic.core.tool.Tool
 import kotlinx.coroutines.*
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.datetime.Clock
-
-interface ToolInterceptor { suspend fun intercept(agent: Agent, tool: Tool, toolCall: ToolCall) }
+import runInterceptors
 
 class AgentExecutor(val toolInterceptors: List<ToolInterceptor> = emptyList()) {
     val agents: MutableList<Agent> = mutableListOf()
@@ -100,7 +99,13 @@ class AgentExecutor(val toolInterceptors: List<ToolInterceptor> = emptyList()) {
     private suspend fun Agent.execute(toolCall: ToolCall): Message.ToolResult {
         val functionArgs = toolCall.argumentsAsJson()
         val tool = tools[ToolName(toolCall.name)] ?: error("Tool not registered: $toolCall")
-        toolInterceptors.forEach { it.intercept(this, tool, toolCall) }
+
+        val cancelMessage = runInterceptors(this, tool, toolCall)
+        if (cancelMessage != null) {
+            setRunningState(AgentRunningState.RUNNING)
+            return cancelMessage
+        }
+
         setRunningState(AgentRunningState.EXECUTING_TOOL)
         val result = tool.handler(functionArgs)
         setRunningState(AgentRunningState.RUNNING)

diff --git a/src/core/src/commonMain/kotlin/community/flock/aigentic/core/agent/ToolInterceptor.kt b/src/core/src/commonMain/kotlin/community/flock/aigentic/core/agent/ToolInterceptor.kt
@@ -0,0 +1,26 @@
+import community.flock.aigentic.core.agent.Agent
+import community.flock.aigentic.core.agent.AgentExecutor
+import community.flock.aigentic.core.message.Message
+import community.flock.aigentic.core.message.ToolCall
+import community.flock.aigentic.core.message.ToolResultContent
+import community.flock.aigentic.core.tool.Tool
+
+data class ToolInterceptorResult(val cancelExecution: Boolean, val reason: String?)
+
+interface ToolInterceptor {
+    suspend fun intercept(agent: Agent, tool: Tool, toolCall: ToolCall): ToolInterceptorResult
+}
+
+suspend fun AgentExecutor.runInterceptors(
+    agent: Agent,
+    tool: Tool,
+    toolCall: ToolCall
+): Message.ToolResult? = toolInterceptors
+    .map { it.intercept(agent, tool, toolCall) }
+    .firstOrNull { it.cancelExecution }?.let {
+        Message.ToolResult(
+            toolCall.id,
+            toolCall.name,
+            ToolResultContent(it.reason ?: "Tool execution blocked by interceptor")
+        )
+    }
diff --git a/src/core/src/commonMain/kotlin/community/flock/aigentic/core/dsl/AgentExecutorConfig.kt b/src/core/src/commonMain/kotlin/community/flock/aigentic/core/dsl/AgentExecutorConfig.kt
@@ -1,8 +1,8 @@
 package community.flock.aigentic.core.dsl
 
+import ToolInterceptor
 import community.flock.aigentic.core.agent.Agent
 import community.flock.aigentic.core.agent.AgentExecutor
-import community.flock.aigentic.core.agent.ToolInterceptor
 
 fun agentExecutor(agentExecutorBuilder: AgentExecutorConfig.() -> Unit): AgentExecutor {
     return AgentExecutorConfig().apply(agentExecutorBuilder).build()