Reading current prompt within Tool code

[samDobsonDev]

Hi there, how would I go about accessing the current prompt from inside a tool? One of the tools I've got sends an email out with the conversation history thus far, but I'm not sure how to access said history within my tool so I can add it to the email content.

[Ololoshechkin]

Hi @samDobsonDev ! Thanks for your question and sorry for our delay in the answer!

Currently there's no direct way to get the prompt from inside the tool call, but you can workaround this using Persistency feature (see https://docs.koog.ai/agent-persistency/ ) like this:

1. Implement the ToolSet class so that it would hold the PersistencyStorageProvider (from the Persistency feature) that contains the latest saved full message history of your agent, and define all your tools inside.
   For example, your sendEmail tool could be similar to this:

class MyTools(private val agentState: PersistencyStorageProvider) : ToolSet {
        @Tool
        fun sendEmail(emailText: String, userId: String): String {
            val messages = runBlocking {
                agentState.getLatestCheckpoint()?.messageHistory?.map { message ->
                    "From ${message.role}: ${message.content}"
                }.orEmpty().joinToString("-----------------------\n")
            }

            // TODO: send email using the `messages` (AI Agent message history)

            return "Email sent to $userId"
        }

        @Tool
        fun someOtherTool() {...}
    }

2. Create some in-memory state storage :
   val agentStateStorage = InMemoryPersistencyStorageProvider("temporary-unique-id")

3. Create your agent with the tools from your ToolSet and with Persistency feature looking at the same state storage:

val agent = AIAgent(
            promptExecutor = ...,
            strategy = ...,
            agentConfig = ...,
            toolRegistry = ToolRegistry {
                tools(MyTools(agentStateStorage))
            }
) {
    install(Persistency) {
        storage = agentStateStorage
        enableAutomaticPersistency = true
    }
    ...
}

That's a bit a hacky (yet working) way to do that, but initially Koog wasn't designed for the intent of exposing the message history or prompt to the tools. Hope this helps!

Also it would be really interesting to learn more about your use case so that we can adapt better to it in the future :)

Cheers!

[samDobsonDev]

Yeah, so one of the tools I'm designing hands a user over to a human when requested, like if the AI agent is unable to address the user's query or if the user explicitly asks for it. This handover is just an email that we send to the human agent with user information etc. Another thing I need to add to the email body is the conversation history up to that point, so that the human agent has more context of the problem that the user is facing.

[Ololoshechkin]

Thanks for the clarification, makes a lot of sense now!

I can suggest an alternative approach to the one proposed above -- make this via an explicit node after you classified such situation.

1. Let's define the sendEmail tool:

@Tool
fun sendEmail(pastMessages: List<Message>, currentQuestion: String, userId: String): Unit {
    TODO("send email")
}

2. Also, let's define the possible classification results:

@Serializable
sealed interface ClassificationResult {
    @Serializable
    class ContinueWithAI(val message: String) : ClassificationResult

    @Serializable
    class HandOverToHuman(val message: String) : ClassificationResult

    @Serializable
    class Done(val message: String) : ClassificationResult

    @Serializable
    enum class Type {
        CONTINUE_WITH_AI,
        HAND_OVER_TO_HUMAN,
        WORK_IS_DONE;
    }
}

3. And then, in your strategy you can write something, for example, like this:

val handOverToHuman by node<HandOverToHuman, Unit> { question ->
    llm.writeSession {
        callTool(::sendEmail, prompt.messages, question.message, "current-user-id")
    }
}

val classifyRequest by node<String, ClassificationResult> { userInput ->
    llm.writeSession {
        val initialPrompt = prompt.copy()

        updatePrompt {
            user {
                markdown {
                    h1("Question (MAIN GOAL!)")
                    text("What is the intent of the user in the following message?")
                    h2("Format of the answer")
                    bulleted {
                        item("${ClassificationResult.Type.CONTINUE_WITH_AI} -- user doesn't seem to be against continuing with AI chat bot")
                        item("${ClassificationResult.Type.HAND_OVER_TO_HUMAN} -- user clearly wants to hand over this conversation to a human")
                        item("${ClassificationResult.Type.WORK_IS_DONE} -- user clearly wants to finish the conversation and stop this chat because the problem was resolved (or the user is done with the conversation)")
                    }
                    h2("The user's message to analyze:")
                    blockquote(userInput)
                }
            }
        }

        val intentType = requestLLMStructured(
            JsonStructuredData.createJsonStructure<ClassificationResult.Type>(),
            retries = 2,
            fixingModel = model // or change to something cheaper than the current model, if needed
        )
            .onSuccess { it.structure }
            .getOrElse { ClassificationResult.Type.CONTINUE_WITH_AI }

        // restore the initial prompt
        rewritePrompt { initialPrompt }

        when (intentType) {
            ClassificationResult.Type.HAND_OVER_TO_HUMAN -> HandOverToHuman(userInput)
            else -> ContinueWithAI(userInput)
        }
    }
}

val askUser by node<Any, String> { readln() } // or any other form of input
val askLLM by nodeLLMRequest()
val callTool by nodeExecuteTool()

edge(askUser forwardTo classifyRequest)
edge(classifyRequest forwardTo handOverToHuman onIsInstance HandOverToHuman::class)
edge(classifyRequest forwardTo askLLM onIsInstance ContinueWithAI::class transformed { it.message })
edge(classifyRequest forwardTo nodeFinish onIsInstance ClassificationResult.Done::class)
edge(askLLM forwardTo callTool onToolCall { true })
edge(askLLM forwardTo askUser onAssistantMessage { true })

[Ololoshechkin]

Also I forgot the edge(handOverToHuman forwardTo nodeFinish) -- otherwise the agent will hang after giving the work to a human :)