(KG-212) Introduce ResponseProcessor to fix tool call messages from weak models (#871)

<!--
Thank you for opening a pull request!

Please add a brief description of the proposed change here.
Also, please tick the appropriate points in the checklist below.
-->

[KG-212](https://youtrack.jetbrains.com/issue/KG-212) Add validation and
retry logic for tool calling from local models

PR introduces a generic `ResponseProcessor` which can process the
messages received from the llm. A few out-of-the-box implementations aim
to fix tool call messages, so that the agent follows the expected flow.

## Motivation and Context
<!-- Why is this change needed? What problem does it solve? -->

Small models (e.g. most ollama models) often fail to properly generate a
tool call messages. Some problems are
- incorrectly generated jsons
- Incorrect tool name / parameter names
- Message "I will call this tool" instead of a direct tool call

`ResponseProcessor` and its implementations are designed to deal with
such issues.

## Breaking Changes
<!-- Will users need to update their code or configurations? -->

---

#### Type of the changes
- [x] New feature (non-breaking change which adds functionality)
- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Tests improvement
- [ ] Refactoring

#### Checklist
- [x] The pull request has a description of the proposed change
- [x] I read the [Contributing
Guidelines](https://github.com/JetBrains/koog/blob/main/CONTRIBUTING.md)
before opening the pull request
- [x] The pull request uses **`develop`** as the base branch
- [x] Tests for the changes have been added
- [x] All new and existing tests passed

##### Additional steps for pull requests adding a new feature
- [x] An issue describing the proposed change exists
- [x] The pull request includes a link to the issue
- [x] The change was discussed and approved in the issue
- [ ] Docs have been added / updated

---------

Co-authored-by: Andrey Bragin <[email protected]>

Author: antoniibelyshev

agents/agents-core/build.gradle.kts

@@ -17,6 +17,7 @@ kotlin {
                 api(project(":utils"))
                 api(project(":prompt:prompt-executor:prompt-executor-model"))
                 api(project(":prompt:prompt-llm"))
+                api(project(":prompt:prompt-processor"))
                 api(project(":prompt:prompt-structure"))
 
                 api(project(":prompt:prompt-executor:prompt-executor-clients:prompt-executor-openai-client"))

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/AIAgent.kt

@@ -13,6 +13,7 @@ import ai.koog.prompt.dsl.prompt
 import ai.koog.prompt.executor.model.PromptExecutor
 import ai.koog.prompt.llm.LLModel
 import ai.koog.prompt.params.LLMParams
+import ai.koog.prompt.processor.ResponseProcessor
 import ai.koog.utils.io.Closeable
 import kotlinx.datetime.Clock
 import kotlin.reflect.typeOf
@@ -259,6 +260,7 @@ public interface AIAgent<Input, Output> : Closeable {
          *
          * @param promptExecutor The executor responsible for processing language model prompts.
          * @param llmModel The specific large language model to be used for the agent.
+         * @param responseProcessor The processor responsible for processing the model's responses.
          * @param strategy The strategy that defines the agent's workflow, defaulting to the [singleRunStrategy].
          * @param toolRegistry The set of tools available for the agent, defaulting to an empty registry.
          * @param id Unique identifier for the agent. Random UUID will be generated if set to null.
@@ -273,6 +275,7 @@ public interface AIAgent<Input, Output> : Closeable {
         public operator fun invoke(
             promptExecutor: PromptExecutor,
             llmModel: LLModel,
+            responseProcessor: ResponseProcessor? = null,
             strategy: AIAgentGraphStrategy<String, String> = singleRunStrategy(),
             toolRegistry: ToolRegistry = ToolRegistry.EMPTY,
             id: String? = null,
@@ -297,6 +300,7 @@ public interface AIAgent<Input, Output> : Closeable {
                 },
                 model = llmModel,
                 maxAgentIterations = maxIterations,
+                responseProcessor = responseProcessor
             ),
             toolRegistry = toolRegistry,
             installFeatures = installFeatures
@@ -310,6 +314,7 @@ public interface AIAgent<Input, Output> : Closeable {
          * @param promptExecutor An instance of [PromptExecutor] responsible for executing prompts with the language model.
          * @param llmModel The language model [LLModel] to be used by the agent.
          * @param strategy The agent strategy [AIAgentGraphStrategy] defining how the agent processes inputs and outputs.
+         * @param responseProcessor The processor responsible for processing the model's responses.
          * @param toolRegistry An optional [ToolRegistry] specifying the tools available to the agent for execution. Defaults to `[ToolRegistry.EMPTY]`.
          * @param id Unique identifier for the agent. Random UUID will be generated if set to null.
          * @param clock A `Clock` instance used for time-related operations. Defaults to `Clock.System`.
@@ -325,6 +330,7 @@ public interface AIAgent<Input, Output> : Closeable {
             promptExecutor: PromptExecutor,
             llmModel: LLModel,
             strategy: AIAgentGraphStrategy<Input, Output>,
+            responseProcessor: ResponseProcessor? = null,
             toolRegistry: ToolRegistry = ToolRegistry.EMPTY,
             id: String? = null,
             clock: Clock = Clock.System,
@@ -352,6 +358,7 @@ public interface AIAgent<Input, Output> : Closeable {
                     },
                     model = llmModel,
                     maxAgentIterations = maxIterations,
+                    responseProcessor = responseProcessor
                 ),
                 toolRegistry = toolRegistry,
                 clock = clock,
@@ -367,6 +374,7 @@ public interface AIAgent<Input, Output> : Closeable {
          * @param Output The type of output produced by the agent.
          * @param promptExecutor The executor used to process prompts for the language model.
          * @param llmModel The language model configuration defining the underlying LLM instance and its behavior.
+         * @param responseProcessor The processor responsible for processing the model's responses.
          * @param toolRegistry Registry containing tools available to the agent for use during execution. Default is an empty registry.
          * @param strategy The strategy to be executed by the agent. Default is a single-run strategy.
          * @param id Unique identifier for the agent. Random UUID will be generated if set to null.
@@ -380,6 +388,7 @@ public interface AIAgent<Input, Output> : Closeable {
         public operator fun <Input, Output> invoke(
             promptExecutor: PromptExecutor,
             llmModel: LLModel,
+            responseProcessor: ResponseProcessor? = null,
             toolRegistry: ToolRegistry = ToolRegistry.EMPTY,
             strategy: AIAgentFunctionalStrategy<Input, Output>,
             id: String? = null,

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/AIAgentService.kt

@@ -11,6 +11,7 @@ import ai.koog.prompt.dsl.prompt
 import ai.koog.prompt.executor.model.PromptExecutor
 import ai.koog.prompt.llm.LLModel
 import ai.koog.prompt.params.LLMParams
+import ai.koog.prompt.processor.ResponseProcessor
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
 import kotlinx.datetime.Clock
@@ -203,6 +204,7 @@ public interface AIAgentService<Input, Output, TAgent : AIAgent<Input, Output>>
         public operator fun invoke(
             promptExecutor: PromptExecutor,
             llmModel: LLModel,
+            responseProcessor: ResponseProcessor? = null,
             strategy: AIAgentGraphStrategy<String, String> = singleRunStrategy(),
             toolRegistry: ToolRegistry = ToolRegistry.EMPTY,
             systemPrompt: String = "",
@@ -225,6 +227,7 @@ public interface AIAgentService<Input, Output, TAgent : AIAgent<Input, Output>>
                 },
                 model = llmModel,
                 maxAgentIterations = maxIterations,
+                responseProcessor = responseProcessor
             ),
             toolRegistry = toolRegistry,
             installFeatures = installFeatures

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/FunctionalAIAgent.kt

@@ -89,6 +89,7 @@ public class FunctionalAIAgent<Input, Output>(
             toolRegistry = toolRegistry,
             prompt = agentConfig.prompt,
             model = agentConfig.model,
+            responseProcessor = agentConfig.responseProcessor,
             promptExecutor = PromptExecutorProxy(
                 executor = promptExecutor,
                 pipeline = pipeline,

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/GraphAIAgent.kt

@@ -124,6 +124,7 @@ public open class GraphAIAgent<Input, Output>(
                 toolRegistry = toolRegistry,
                 prompt = agentConfig.prompt,
                 model = agentConfig.model,
+                responseProcessor = agentConfig.responseProcessor,
                 promptExecutor = PromptExecutorProxy(
                     executor = promptExecutor,
                     pipeline = pipeline,

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/config/AIAgentConfig.kt

@@ -4,6 +4,7 @@ import ai.koog.prompt.dsl.Prompt
 import ai.koog.prompt.dsl.prompt
 import ai.koog.prompt.executor.clients.openai.OpenAIModels
 import ai.koog.prompt.llm.LLModel
+import ai.koog.prompt.processor.ResponseProcessor
 
 /**
  * Configuration class for an AI agent that specifies the prompt, execution parameters, and behavior.
@@ -20,13 +21,15 @@ import ai.koog.prompt.llm.LLModel
  *        This property provides a mechanism to convert or format missing tool calls and result messages when they occur,
  *        typically due to differences in tool sets between steps or subgraphs while the same history is reused.
  *        This ensures the prompt remains consistent and readable for the model, even with undefined tools.
+ * @param responseProcessor Optional processor for the agent's responses. If specified, will modify the responses from the llm.
  */
 public class AIAgentConfig(
     override val prompt: Prompt,
     override val model: LLModel,
     public val maxAgentIterations: Int,
     public val missingToolsConversionStrategy: MissingToolsConversionStrategy =
-        MissingToolsConversionStrategy.Missing(ToolCallDescriber.JSON)
+        MissingToolsConversionStrategy.Missing(ToolCallDescriber.JSON),
+    public val responseProcessor: ResponseProcessor? = null,
 ) : AIAgentConfigBase {
 
     init {

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/context/AIAgentLLMContext.kt

@@ -12,6 +12,7 @@ import ai.koog.agents.core.utils.RWLock
 import ai.koog.prompt.dsl.Prompt
 import ai.koog.prompt.executor.model.PromptExecutor
 import ai.koog.prompt.llm.LLModel
+import ai.koog.prompt.processor.ResponseProcessor
 import kotlinx.datetime.Clock
 
 /**
@@ -44,6 +45,7 @@ public annotation class DetachedPromptExecutorAPI
  * @property toolRegistry A registry that contains metadata about available tools.
  * @property prompt The current LLM prompt being used or updated in write sessions.
  * @property model The current LLM model being used or updated in write sessions.
+ * @property responseProcessor The current response processor being used or updated in write sessions.
  * @property promptExecutor The [PromptExecutor] responsible for performing operations on the current prompt.
  * @property environment The environment that manages tool execution and interaction with external dependencies.
  * @property clock The clock used for timestamps of messages
@@ -53,6 +55,7 @@ public class AIAgentLLMContext(
     public val toolRegistry: ToolRegistry = ToolRegistry.EMPTY,
     prompt: Prompt,
     model: LLModel,
+    responseProcessor: ResponseProcessor?,
     @property:DetachedPromptExecutorAPI
     public val promptExecutor: PromptExecutor,
     private val environment: AIAgentEnvironment,
@@ -73,6 +76,13 @@ public class AIAgentLLMContext(
     public var model: LLModel = model
         private set
 
+    /**
+     * Response processor currently associated with this context.
+     */
+    @DetachedPromptExecutorAPI
+    public var responseProcessor: ResponseProcessor? = responseProcessor
+        private set
+
     /**
      * The current prompt used within the `AIAgentLLMContext`.
      *
@@ -106,6 +116,7 @@ public class AIAgentLLMContext(
         toolRegistry: ToolRegistry = this.toolRegistry,
         prompt: Prompt = this.prompt,
         model: LLModel = this.model,
+        responseProcessor: ResponseProcessor? = this.responseProcessor,
         promptExecutor: PromptExecutor = this.promptExecutor,
         environment: AIAgentEnvironment = this.environment,
         config: AIAgentConfig = this.config,
@@ -116,6 +127,7 @@ public class AIAgentLLMContext(
             toolRegistry = toolRegistry,
             prompt = prompt,
             model = model,
+            responseProcessor = responseProcessor,
             promptExecutor = promptExecutor,
             environment = environment,
             config = config,
@@ -132,7 +144,7 @@ public class AIAgentLLMContext(
     @OptIn(ExperimentalStdlibApi::class)
     public suspend fun <T> writeSession(block: suspend AIAgentLLMWriteSession.() -> T): T = rwLock.withWriteLock {
         val session =
-            AIAgentLLMWriteSession(environment, promptExecutor, tools, toolRegistry, prompt, model, config, clock)
+            AIAgentLLMWriteSession(environment, promptExecutor, tools, toolRegistry, prompt, model, responseProcessor, config, clock)
 
         session.use {
             val result = it.block()
@@ -141,6 +153,7 @@ public class AIAgentLLMContext(
             this.prompt = it.prompt
             this.tools = it.tools
             this.model = it.model
+            this.responseProcessor = it.responseProcessor
 
             result
         }
@@ -152,7 +165,7 @@ public class AIAgentLLMContext(
      */
     @OptIn(ExperimentalStdlibApi::class)
     public suspend fun <T> readSession(block: suspend AIAgentLLMReadSession.() -> T): T = rwLock.withReadLock {
-        val session = AIAgentLLMReadSession(tools, promptExecutor, prompt, model, config)
+        val session = AIAgentLLMReadSession(tools, promptExecutor, prompt, model, responseProcessor, config)
 
         session.use { block(it) }
     }
@@ -166,11 +179,12 @@ public class AIAgentLLMContext(
         tools: List<ToolDescriptor> = this.tools,
         prompt: Prompt = this.prompt,
         model: LLModel = this.model,
+        responseProcessor: ResponseProcessor? = this.responseProcessor,
         promptExecutor: PromptExecutor = this.promptExecutor,
         environment: AIAgentEnvironment = this.environment,
         config: AIAgentConfig = this.config,
         clock: Clock = this.clock
     ): AIAgentLLMContext {
-        return AIAgentLLMContext(tools, toolRegistry, prompt, model, promptExecutor, environment, config, clock)
+        return AIAgentLLMContext(tools, toolRegistry, prompt, model, responseProcessor, promptExecutor, environment, config, clock)
     }
 }

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/entity/AIAgentSubgraph.kt

@@ -16,6 +16,7 @@ import ai.koog.agents.core.tools.ToolDescriptor
 import ai.koog.agents.core.tools.annotations.LLMDescription
 import ai.koog.prompt.llm.LLModel
 import ai.koog.prompt.params.LLMParams
+import ai.koog.prompt.processor.ResponseProcessor
 import ai.koog.prompt.structure.StructureFixingParser
 import ai.koog.prompt.structure.StructuredRequest
 import ai.koog.prompt.structure.StructuredRequestConfig
@@ -42,6 +43,7 @@ import kotlin.uuid.Uuid
  * @param toolSelectionStrategy Strategy determining which tools should be available during this subgraph's execution.
  * @param llmModel Optional [LLModel] override for the subgraph execution.
  * @param llmParams Optional [LLMParams] override for the prompt for the subgraph execution.
+ * @param responseProcessor Optional [ResponseProcessor] override for the subgraph execution.
  */
 public open class AIAgentSubgraph<TInput, TOutput>(
     override val name: String,
@@ -50,6 +52,7 @@ public open class AIAgentSubgraph<TInput, TOutput>(
     private val toolSelectionStrategy: ToolSelectionStrategy,
     private val llmModel: LLModel? = null,
     private val llmParams: LLMParams? = null,
+    private val responseProcessor: ResponseProcessor? = null,
 ) : AIAgentNodeBase<TInput, TOutput>(), ExecutionPointNode {
     override val inputType: KType = start.inputType
     override val outputType: KType = finish.outputType
@@ -166,7 +169,8 @@ public open class AIAgentSubgraph<TInput, TOutput>(
                     llm = llm.copy(
                         tools = newTools,
                         model = llmModel ?: llm.model,
-                        prompt = llm.prompt.copy(params = llmParams ?: llm.prompt.params)
+                        prompt = llm.prompt.copy(params = llmParams ?: llm.prompt.params),
+                        responseProcessor = responseProcessor
                     )
                 )
             }

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/session/AIAgentLLMReadSession.kt

@@ -5,6 +5,7 @@ import ai.koog.agents.core.tools.ToolDescriptor
 import ai.koog.prompt.dsl.Prompt
 import ai.koog.prompt.executor.model.PromptExecutor
 import ai.koog.prompt.llm.LLModel
+import ai.koog.prompt.processor.ResponseProcessor
 
 /**
  * Represents a session for interacting with a language model (LLM) in a read-only context within an AI agent setup.
@@ -16,12 +17,14 @@ import ai.koog.prompt.llm.LLModel
  * @param executor The `PromptExecutor` responsible for handling execution of prompts within this session.
  * @param prompt The `Prompt` object specifying the input messages and parameters for the session.
  * @param model The language model instance to be used for processing prompts in this session.
+ * @param responseProcessor The response processor instance to be used for post-processing responses.
  * @param config The configuration settings for the AI agent session.
  */
 public class AIAgentLLMReadSession internal constructor(
     tools: List<ToolDescriptor>,
     executor: PromptExecutor,
     prompt: Prompt,
     model: LLModel,
+    responseProcessor: ResponseProcessor?,
     config: AIAgentConfig,
-) : AIAgentLLMSession(executor, tools, prompt, model, config)
+) : AIAgentLLMSession(executor, tools, prompt, model, responseProcessor, config)