Support providing custom LLM for fact retrieval in the history (#289)

Author: Ololoshechkin

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/extension/AIAgentLLMActions.kt

@@ -24,6 +24,15 @@ public fun AIAgentLLMWriteSession.leaveLastNMessages(n: Int) {
     prompt = prompt.withMessages { it.takeLast(n) }
 }
 
+/**
+ * Removes the last `n` messages from the current prompt in the write session.
+ *
+ * @param n The number of messages to remove from the end of the current message list.
+ */
+public fun AIAgentLLMWriteSession.dropLastNMessages(n: Int) {
+    prompt = prompt.withMessages { it.dropLast(n) }
+}
+
 /**
  * Removes all messages from the current session's prompt that have a timestamp
  * earlier than the specified timestamp.

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/dsl/extension/AIAgentNodes.kt

@@ -18,6 +18,7 @@ import ai.koog.prompt.message.Message
 import ai.koog.prompt.structure.StructuredData
 import ai.koog.prompt.structure.StructuredDataDefinition
 import ai.koog.prompt.structure.StructuredResponse
+import io.github.oshai.kotlinlogging.KotlinLogging.logger
 import kotlinx.coroutines.flow.Flow
 
 /**
@@ -262,16 +263,26 @@ public fun AIAgentSubgraphBuilderBase<*, *>.nodeLLMRequestMultiple(name: String?
  *
  * @param name Optional node name.
  * @param strategy Determines which messages to include in compression.
+ * @param retrievalModel An optional [LLModel] that will be used for retrieval of the facts from memory.
+ *                       By default, the same model will be used as the current one in the agent's strategy.
  * @param preserveMemory Specifies whether to retain message memory after compression.
  */
 @AIAgentBuilderDslMarker
 public fun <T> AIAgentSubgraphBuilderBase<*, *>.nodeLLMCompressHistory(
     name: String? = null,
     strategy: HistoryCompressionStrategy = HistoryCompressionStrategy.WholeHistory,
+    retrievalModel: LLModel? = null,
     preserveMemory: Boolean = true
 ): AIAgentNodeDelegate<T, T> = node(name) { input ->
     llm.writeSession {
+        val initialModel = model
+        if (retrievalModel != null) {
+            model = retrievalModel
+        }
+
         replaceHistoryWithTLDR(strategy, preserveMemory)
+
+        model = initialModel
     }
 
     input

agents/agents-features/agents-features-memory/src/commonMain/kotlin/ai/koog/agents/memory/feature/AgentMemory.kt

@@ -281,18 +281,29 @@ public class AgentMemory(
      * @param concept The concept to extract facts about
      * @param subject The subject categorization for the facts (e.g., User, Project)
      * @param scope The visibility scope for the facts (e.g., Agent, Feature, Product)
+     * @param retrievalModel LLM that will be used for fact retrieval from the history (by default, the same model as the current one will be used)
      */
     public suspend fun saveFactsFromHistory(
         concept: Concept,
         subject: MemorySubject,
         scope: MemoryScope,
+        retrievalModel: LLModel? = null
     ) {
         llm.writeSession {
+            val initialModel = model
+            if (retrievalModel != null) {
+                model = retrievalModel
+                logger.info { "Using model: ${retrievalModel.id}" }
+            }
             val facts = retrieveFactsFromHistory(concept)
 
             // Save facts to memory
             agentMemory.save(facts, subject, scope)
             logger.info { "Saved fact for concept '${concept.keyword}' in scope $scope: $facts" }
+            if (retrievalModel != null) {
+                model = initialModel
+                logger.info { "Switching back to model: ${initialModel.id}" }
+            }
         }
     }
 

agents/agents-features/agents-features-memory/src/commonMain/kotlin/ai/koog/agents/memory/feature/nodes/MemoryNodes.kt

@@ -3,10 +3,14 @@ package ai.koog.agents.memory.feature.nodes
 import ai.koog.agents.core.dsl.builder.AIAgentBuilderDslMarker
 import ai.koog.agents.core.dsl.builder.AIAgentNodeDelegate
 import ai.koog.agents.core.dsl.builder.AIAgentSubgraphBuilderBase
+import ai.koog.agents.core.dsl.extension.dropLastNMessages
+import ai.koog.agents.core.dsl.extension.leaveLastNMessages
 import ai.koog.agents.memory.config.MemoryScopeType
 import ai.koog.agents.memory.feature.withMemory
 import ai.koog.agents.memory.model.*
 import ai.koog.agents.memory.prompts.MemoryPrompts
+import ai.koog.prompt.llm.LLModel
+import io.github.oshai.kotlinlogging.KotlinLogging.logger
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.json.Json
 
@@ -93,20 +97,23 @@ public fun <T> AIAgentSubgraphBuilderBase<*, *>.nodeLoadAllFactsFromMemory(
  * @param subject The subject scope of the memory (USER, PROJECT, etc.)
  * @param scope The scope of the memory (Agent, Feature, etc.)
  * @param concepts List of concepts to save in memory
+ * @param retrievalModel LLM that will be used for fact retrieval from the history (by default, the same model as the current one will be used)
  */
 @AIAgentBuilderDslMarker
 public fun <T> AIAgentSubgraphBuilderBase<*, *>.nodeSaveToMemory(
     name: String? = null,
     subject: MemorySubject,
     scope: MemoryScopeType,
     concepts: List<Concept>,
+    retrievalModel: LLModel? = null
 ): AIAgentNodeDelegate<T, T> = node(name) { input ->
     withMemory {
         concepts.forEach { concept ->
             saveFactsFromHistory(
                 concept = concept,
                 subject = subject,
                 scope = scopesProfile.getScope(scope) ?: return@forEach,
+                retrievalModel = retrievalModel
             )
         }
     }
@@ -120,14 +127,16 @@ public fun <T> AIAgentSubgraphBuilderBase<*, *>.nodeSaveToMemory(
  * @param subject The subject scope of the memory (USER, PROJECT, etc.)
  * @param scope The scope of the memory (Agent, Feature, etc.)
  * @param concept The concept to save in memory
+ * @param retrievalModel LLM that will be used for fact retrieval from the history (by default, the same model as the current one will be used)
  */
 @AIAgentBuilderDslMarker
 public fun <T> AIAgentSubgraphBuilderBase<*, *>.nodeSaveToMemory(
     name: String? = null,
     concept: Concept,
     subject: MemorySubject,
     scope: MemoryScopeType,
-): AIAgentNodeDelegate<T, T> = nodeSaveToMemory(name, subject, scope, listOf(concept))
+    retrievalModel: LLModel? = null
+): AIAgentNodeDelegate<T, T> = nodeSaveToMemory(name, subject, scope, listOf(concept), retrievalModel)
 
 /**
  * Node that automatically detects and extracts facts from the chat history and saves them to memory.
@@ -137,14 +146,21 @@ public fun <T> AIAgentSubgraphBuilderBase<*, *>.nodeSaveToMemory(
  * @param scopes List of memory scopes (Agent, Feature, etc.). By default only Agent scope would be chosen
  * @param subjects List of subjects (user, project, organization, etc.) to look for.
  * By default, all subjects will be included and looked for.
+ * @param retrievalModel LLM that will be used for fact retrieval from the history (by default, the same model as the current one will be used)
  */
 @AIAgentBuilderDslMarker
 public fun <T> AIAgentSubgraphBuilderBase<*, *>.nodeSaveToMemoryAutoDetectFacts(
     name: String? = null,
     scopes: List<MemoryScopeType> = listOf(MemoryScopeType.AGENT),
-    subjects: List<MemorySubject> = MemorySubject.registeredSubjects
+    subjects: List<MemorySubject> = MemorySubject.registeredSubjects,
+    retrievalModel: LLModel? = null
 ): AIAgentNodeDelegate<T, T> = node(name) { input ->
     llm.writeSession {
+        val initialModel = model
+        val initialPrompt = prompt.copy()
+        if (retrievalModel != null) {
+            model = retrievalModel
+        }
         updatePrompt {
             val prompt = MemoryPrompts.autoDetectFacts(subjects)
             user(prompt)
@@ -160,6 +176,11 @@ public fun <T> AIAgentSubgraphBuilderBase<*, *>.nodeSaveToMemoryAutoDetectFacts(
                 }
             }
         }
+
+        rewritePrompt { initialPrompt } // Revert the prompt to the original one
+        if (retrievalModel != null) {
+            model = initialModel
+        }
     }
 
     input